Add metastore read replica routing

Author: shuheiktgw

config/quickwit.yaml

@@ -103,6 +103,11 @@ version: 0.8
 # metastore_uri: s3://your-bucket/indexes
 # metastore_uri: postgres://username:password@host:port/db
 #
+# Optional PostgreSQL read replica URI used by metastore nodes for gRPC requests
+# carrying the `qw-use-read-replica: true` metadata. Defaults to unset.
+#
+# metastore_read_replica_uri: postgres://username:password@read-replica-host:port/db
+#
 # When using a file-backed metastore, the state of the metastore will be cached forever.
 # If you are indexing and searching from different processes, it is possible to periodically
 # refresh the state of the metastore on the searcher using the `polling_interval` hashtag.

docs/configuration/node-config.md

@@ -30,6 +30,7 @@ A commented example is available here: [quickwit.yaml](https://github.com/quickw
 | `peer_seeds` | List of IP addresses or hostnames used to bootstrap the cluster and discover the complete set of nodes. This list may contain the current node address and does not need to be exhaustive. If the list of peer seeds contains a host name, Quickwit will resolve it by querying the DNS every minute. On kubernetes for instance, it is a good practise to set it to a [headless service](https://kubernetes.io/docs/concepts/services-networking/service/#headless-services). | `QW_PEER_SEEDS` | |
 | `data_dir` | Path to directory where data (tmp data, splits kept for caching purpose) is persisted. This is mostly used in indexing. | `QW_DATA_DIR` | `./qwdata` |
 | `metastore_uri` | Metastore URI. Can be a local directory or `s3://my-bucket/indexes` or `postgres://username:password@localhost:5432/metastore`. [Learn more about the metastore configuration](metastore-config.md). | `QW_METASTORE_URI` | `{data_dir}/indexes` |
+| `metastore_read_replica_uri` | Optional PostgreSQL read replica URI used by metastore nodes when a gRPC request carries the `qw-use-read-replica: true` metadata. If unset, those requests use `metastore_uri`. | `QW_METASTORE_READ_REPLICA_URI` | |
 | `default_index_root_uri` | Default index root URI that defines the location where index data (splits) is stored. The index URI is built following the scheme: `{default_index_root_uri}/{index-id}` | `QW_DEFAULT_INDEX_ROOT_URI` | `{data_dir}/indexes` |
 | environment variable only | Log level of Quickwit. Can be a direct log level, or a comma separated list of `module_name=level` | `RUST_LOG` | `info` |
 

quickwit/quickwit-cluster/src/grpc_service.rs

@@ -44,7 +44,7 @@ pub(crate) async fn cluster_grpc_client(
 
     ClusterServiceClient::tower()
         .stack_layer(CLUSTER_GRPC_CLIENT_METRICS_LAYER.clone())
-        .build_from_channel(socket_addr, channel, MAX_MESSAGE_SIZE, None)
+        .build_from_channel(socket_addr, channel, MAX_MESSAGE_SIZE, None, [])
 }
 
 pub fn cluster_grpc_server(

quickwit/quickwit-codegen/example/src/codegen/hello.rs

@@ -282,7 +282,8 @@ mod tests {
             "127.0.0.1:6666".parse().unwrap(),
             Endpoint::from_static("http://127.0.0.1:6666").connect_lazy(),
         );
-        let grpc_client = HelloClient::from_balance_channel(channel, MAX_GRPC_MESSAGE_SIZE, None);
+        let grpc_client =
+            HelloClient::from_balance_channel(channel, MAX_GRPC_MESSAGE_SIZE, None, []);
 
         assert_eq!(
             grpc_client
@@ -342,7 +343,7 @@ mod tests {
         // The connectivity check fails if there is no client behind the channel.
         let (balance_channel, _): (BalanceChannel<SocketAddr>, _) = BalanceChannel::new();
         let grpc_client =
-            HelloClient::from_balance_channel(balance_channel, MAX_GRPC_MESSAGE_SIZE, None);
+            HelloClient::from_balance_channel(balance_channel, MAX_GRPC_MESSAGE_SIZE, None, []);
         assert_eq!(
             grpc_client
                 .check_connectivity()
@@ -441,6 +442,7 @@ mod tests {
             channel,
             MAX_GRPC_MESSAGE_SIZE,
             Some(CompressionEncoding::Zstd),
+            [],
         );
 
         assert_eq!(
@@ -788,7 +790,7 @@ mod tests {
             "127.0.0.1:7777".parse().unwrap(),
             Endpoint::from_static("http://127.0.0.1:7777").connect_lazy(),
         );
-        HelloClient::from_balance_channel(balance_channed, MAX_GRPC_MESSAGE_SIZE, None);
+        HelloClient::from_balance_channel(balance_channed, MAX_GRPC_MESSAGE_SIZE, None, []);
     }
 
     #[tokio::test]
@@ -876,7 +878,7 @@ mod tests {
             .timeout(Duration::from_millis(100))
             .connect_lazy();
         let max_message_size = ByteSize::mib(1);
-        let grpc_client = HelloClient::from_channel(addr, channel, max_message_size, None);
+        let grpc_client = HelloClient::from_channel(addr, channel, max_message_size, None, []);
 
         let error = grpc_client
             .hello(HelloRequest {
@@ -955,7 +957,7 @@ mod tests {
             // this test hangs forever if we comment out the TimeoutLayer, which
             // shows that a request without explicit timeout might hang forever
             .stack_layer(TimeoutLayer::new(Duration::from_secs(3)))
-            .build_from_balance_channel(balance_channel, ByteSize::mib(1), None);
+            .build_from_balance_channel(balance_channel, ByteSize::mib(1), None, []);
 
         let response_fut = async move {
             grpc_client

quickwit/quickwit-codegen/example/src/lib.rs

@@ -316,6 +316,10 @@ impl CodegenContext {
     }
 }
 
+fn is_metastore_service(context: &CodegenContext) -> bool {
+    context.package_name == "quickwit.metastore" && context.service_name == "MetastoreService"
+}
+
 fn generate_all(
     service: &Service,
     result_type_path: &str,
@@ -637,6 +641,21 @@ fn generate_client(context: &CodegenContext) -> TokenStream {
     } else {
         TokenStream::new()
     };
+    let metastore_read_replica_client_methods = if is_metastore_service(context) {
+        quote! {
+            pub fn as_grpc_service_with_read_replica(
+                &self,
+                read_replica: Option<Self>,
+                max_message_size: bytesize::ByteSize,
+            ) -> crate::grpc_read_replica::ReadReplicaGrpcService<#grpc_server_package_name::#grpc_server_name<#grpc_server_adapter_name>> {
+                let primary = self.as_grpc_service(max_message_size);
+                let read_replica = read_replica.map(|client| client.as_grpc_service(max_message_size));
+                crate::grpc_read_replica::ReadReplicaGrpcService::new(primary, read_replica)
+            }
+        }
+    } else {
+        TokenStream::new()
+    };
 
     quote! {
         #[derive(Debug, Clone)]
@@ -671,15 +690,21 @@ fn generate_client(context: &CodegenContext) -> TokenStream {
                     .max_encoding_message_size(max_message_size.0 as usize)
             }
 
+            #metastore_read_replica_client_methods
+
             pub fn from_channel(
                 addr: std::net::SocketAddr,
                 channel: tonic::transport::Channel,
                 max_message_size: bytesize::ByteSize,
                 compression_encoding_opt: Option<tonic::codec::CompressionEncoding>,
+                interceptors: impl IntoIterator<Item = quickwit_common::tower::GrpcInterceptor>,
             ) -> Self
             {
                 let (_, connection_keys_watcher) = tokio::sync::watch::channel(std::collections::HashSet::from_iter([addr]));
-                let mut client = #grpc_client_package_name::#grpc_client_name::new(channel)
+                let mut client = #grpc_client_package_name::#grpc_client_name::with_interceptor(
+                    channel,
+                    quickwit_common::tower::GrpcInterceptors::new(interceptors),
+                )
                     .max_decoding_message_size(max_message_size.0 as usize)
                     .max_encoding_message_size(max_message_size.0 as usize);
                 if let Some(compression_encoding) = compression_encoding_opt {
@@ -695,10 +720,14 @@ fn generate_client(context: &CodegenContext) -> TokenStream {
                 balance_channel: quickwit_common::tower::BalanceChannel<std::net::SocketAddr>,
                 max_message_size: bytesize::ByteSize,
                 compression_encoding_opt: Option<tonic::codec::CompressionEncoding>,
+                interceptors: impl IntoIterator<Item = quickwit_common::tower::GrpcInterceptor>,
             ) -> #client_name
             {
                 let connection_keys_watcher = balance_channel.connection_keys_watcher();
-                let mut client = #grpc_client_package_name::#grpc_client_name::new(balance_channel)
+                let mut client = #grpc_client_package_name::#grpc_client_name::with_interceptor(
+                    balance_channel,
+                    quickwit_common::tower::GrpcInterceptors::new(interceptors),
+                )
                     .max_decoding_message_size(max_message_size.0 as usize)
                     .max_encoding_message_size(max_message_size.0 as usize);
                 if let Some(compression_encoding) = compression_encoding_opt {
@@ -1024,7 +1053,6 @@ fn generate_layer_stack_impl(context: &CodegenContext) -> TokenStream {
 
         svc_attribute_idents.push(svc_attribute_name);
     }
-
     quote! {
         impl #tower_layer_stack_name {
             pub fn stack_layer<L>(mut self, layer: L) -> Self
@@ -1051,9 +1079,10 @@ fn generate_layer_stack_impl(context: &CodegenContext) -> TokenStream {
                 channel: tonic::transport::Channel,
                 max_message_size: bytesize::ByteSize,
                 compression_encoding_opt: Option<tonic::codec::CompressionEncoding>,
+                interceptors: impl IntoIterator<Item = quickwit_common::tower::GrpcInterceptor>,
             ) -> #client_name
             {
-                let client =  #client_name::from_channel(addr, channel, max_message_size, compression_encoding_opt);
+                let client =  #client_name::from_channel(addr, channel, max_message_size, compression_encoding_opt, interceptors);
                 let inner_client = client.inner;
                 self.build_from_inner_client(inner_client)
             }
@@ -1063,9 +1092,10 @@ fn generate_layer_stack_impl(context: &CodegenContext) -> TokenStream {
                 balance_channel: quickwit_common::tower::BalanceChannel<std::net::SocketAddr>,
                 max_message_size: bytesize::ByteSize,
                 compression_encoding_opt: Option<tonic::codec::CompressionEncoding>,
+                interceptors: impl IntoIterator<Item = quickwit_common::tower::GrpcInterceptor>,
             ) -> #client_name
             {
-                let client =  #client_name::from_balance_channel(balance_channel, max_message_size, compression_encoding_opt);
+                let client =  #client_name::from_balance_channel(balance_channel, max_message_size, compression_encoding_opt, interceptors);
                 let inner_client = client.inner;
                 self.build_from_inner_client(inner_client)
             }
@@ -1279,7 +1309,7 @@ fn generate_grpc_client_adapter(context: &CodegenContext) -> TokenStream {
             pub fn new(instance: T, connection_addrs_rx: tokio::sync::watch::Receiver<std::collections::HashSet<std::net::SocketAddr>>) -> Self {
                 Self {
                     inner: instance,
-                    connection_addrs_rx
+                    connection_addrs_rx,
                 }
             }
         }
@@ -1435,8 +1465,7 @@ fn generate_grpc_server_adapter_methods(context: &CodegenContext) -> TokenStream
                 let span = #span_macro;
                 let _ = <tracing::Span as tracing_opentelemetry::OpenTelemetrySpanExt>::set_parent(&span, parent_context);
                 let fut = async move {
-                    self.inner
-                        .0
+                    self.inner.0
                         .#method_name(request)
                         .await
                         .map(#into_response_type)

quickwit/quickwit-codegen/src/codegen.rs

@@ -0,0 +1,116 @@
+// Copyright 2021-Present Datadog, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use tonic::metadata::KeyAndValueRef;
+use tonic::service::Interceptor;
+
+pub type GrpcInterceptor =
+    Arc<dyn Fn(tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> + Send + Sync>;
+
+#[derive(Clone, Default)]
+pub struct GrpcInterceptors {
+    interceptors: Vec<GrpcInterceptor>,
+}
+
+impl GrpcInterceptors {
+    pub fn new(interceptors: impl IntoIterator<Item = GrpcInterceptor>) -> Self {
+        Self {
+            interceptors: interceptors.into_iter().collect(),
+        }
+    }
+}
+
+impl Interceptor for GrpcInterceptors {
+    fn call(&mut self, request: tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> {
+        self.interceptors
+            .iter()
+            .try_fold(request, |request, interceptor| interceptor(request))
+    }
+}
+
+pub fn fixed_headers_interceptor(headers: tonic::metadata::MetadataMap) -> GrpcInterceptor {
+    Arc::new(move |mut request: tonic::Request<()>| {
+        for key_and_value in headers.iter() {
+            match key_and_value {
+                KeyAndValueRef::Ascii(key, value) => {
+                    request.metadata_mut().insert(key, value.to_owned());
+                }
+                KeyAndValueRef::Binary(key, value) => {
+                    request.metadata_mut().insert_bin(key, value.to_owned());
+                }
+            }
+        }
+        Ok(request)
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use tonic::metadata::{BinaryMetadataValue, MetadataMap, MetadataValue};
+
+    use super::*;
+
+    #[test]
+    fn fixed_headers_interceptor_copies_ascii_and_binary_metadata() {
+        let mut headers = MetadataMap::new();
+        headers.insert("x-ascii", MetadataValue::from_static("ascii-value"));
+        headers.insert_bin("x-bin", BinaryMetadataValue::from_bytes(b"binary-value"));
+
+        let mut interceptors = GrpcInterceptors::new([fixed_headers_interceptor(headers)]);
+        let request = interceptors.call(tonic::Request::new(())).unwrap();
+
+        assert_eq!(
+            request.metadata().get("x-ascii").unwrap(),
+            MetadataValue::from_static("ascii-value")
+        );
+        assert_eq!(
+            request
+                .metadata()
+                .get_bin("x-bin")
+                .unwrap()
+                .to_bytes()
+                .unwrap()
+                .as_ref(),
+            b"binary-value"
+        );
+    }
+
+    #[test]
+    fn grpc_interceptors_apply_multiple_interceptors_in_order() {
+        let first: GrpcInterceptor = Arc::new(|mut request| {
+            request
+                .metadata_mut()
+                .insert("x-order", MetadataValue::from_static("first"));
+            Ok(request)
+        });
+        let second: GrpcInterceptor = Arc::new(|mut request| {
+            let value = request.metadata().get("x-order").unwrap().to_str().unwrap();
+            assert_eq!(value, "first");
+            request
+                .metadata_mut()
+                .insert("x-order", MetadataValue::from_static("second"));
+            Ok(request)
+        });
+
+        let mut interceptors = GrpcInterceptors::new([first, second]);
+        let request = interceptors.call(tonic::Request::new(())).unwrap();
+
+        assert_eq!(
+            request.metadata().get("x-order").unwrap(),
+            MetadataValue::from_static("second")
+        );
+    }
+}

quickwit/quickwit-common/src/tower/interceptor.rs

@@ -20,6 +20,7 @@ mod circuit_breaker;
 mod delay;
 mod estimate_rate;
 mod event_listener;
+mod interceptor;
 mod load_shed;
 mod metrics;
 mod one_task_per_call_layer;
@@ -43,6 +44,7 @@ pub use delay::{Delay, DelayLayer};
 pub use estimate_rate::{EstimateRate, EstimateRateLayer};
 pub use event_listener::{EventListener, EventListenerLayer};
 use futures::Future;
+pub use interceptor::{GrpcInterceptor, GrpcInterceptors, fixed_headers_interceptor};
 pub use load_shed::{LoadShed, LoadShedLayer, MakeLoadShedError};
 pub use metrics::{GrpcMetrics, GrpcMetricsLayer, RpcName};
 pub use one_task_per_call_layer::{OneTaskPerCallLayer, TaskCancelled};