Add timeout to metastore GRPC client (#5716)

* Add timeout to metastore GRPC client

* Explicit the test content

Author: rdettai

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

@@ -15,6 +15,7 @@
 use std::fmt;
 
 use quickwit_actors::AskError;
+use quickwit_common::tower::TimeoutExceeded;
 use quickwit_proto::error::GrpcServiceError;
 pub use quickwit_proto::error::{grpc_error_to_grpc_status, grpc_status_to_service_error};
 use quickwit_proto::{ServiceError, ServiceErrorCode};
@@ -72,3 +73,9 @@ where E: fmt::Debug
         HelloError::Internal(format!("{error:?}"))
     }
 }
+
+impl From<TimeoutExceeded> for HelloError {
+    fn from(_: TimeoutExceeded) -> Self {
+        HelloError::Timeout("client".to_string())
+    }
+}

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

@@ -162,7 +162,7 @@ mod tests {
 
     use bytesize::ByteSize;
     use quickwit_actors::{Actor, ActorContext, ActorExitStatus, Handler, Universe};
-    use quickwit_common::tower::{BalanceChannel, Change};
+    use quickwit_common::tower::{BalanceChannel, Change, TimeoutLayer};
     use tokio::sync::mpsc::error::TrySendError;
     use tokio_stream::StreamExt;
     use tonic::transport::{Endpoint, Server};
@@ -767,4 +767,61 @@ mod tests {
             .unwrap_err();
         assert!(matches!(error, HelloError::Timeout(_)));
     }
+
+    #[tokio::test]
+    async fn test_balanced_channel_timeout_with_server_crash() {
+        let addr_str = "127.0.0.1:11112";
+        let addr: SocketAddr = addr_str.parse().unwrap();
+        // We want to abruptly stop a server without even sending the connection
+        // RST packet. Simply dropping the tonic Server is not enough, so we
+        // spawn a thread and freeze it with thread::park().
+        std::thread::spawn(move || {
+            let server_fut = async {
+                let hello = HelloImpl {
+                    // delay the response so that the server freezes in the middle of the request
+                    delay: Duration::from_millis(1000),
+                };
+                let grpc_server_adapter = HelloGrpcServerAdapter::new(hello);
+                let grpc_server = HelloGrpcServer::new(grpc_server_adapter);
+                tokio::select! {
+                    // wait just enough to let the client perform its request
+                    _ = tokio::time::sleep(Duration::from_millis(100)) => {}
+                    _ = Server::builder().add_service(grpc_server).serve(addr) => {}
+                };
+                std::thread::park();
+                println!("Thread unparked, unexpected");
+            };
+            tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .unwrap()
+                .block_on(server_fut);
+        });
+
+        // create a client that will try to connect to the server
+        let (balance_channel, balance_channel_tx) = BalanceChannel::new();
+        let channel = Endpoint::from_str(&format!("http://{addr_str}"))
+            .unwrap()
+            .connect_lazy();
+        balance_channel_tx
+            .send(Change::Insert(addr, channel))
+            .unwrap();
+
+        let grpc_client = HelloClient::tower()
+            // 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));
+
+        let response_fut = async move {
+            grpc_client
+                .hello(HelloRequest {
+                    name: "World".to_string(),
+                })
+                .await
+        };
+        response_fut
+            .await
+            .expect_err("should have timed out at the client level");
+    }
 }

quickwit/quickwit-indexing/src/actors/indexing_service.rs

@@ -749,6 +749,10 @@ impl IndexingService {
 
     /// Shuts down the pipelines with supplied ids and performs necessary cleanup.
     async fn shutdown_pipelines(&mut self, pipelines_to_shutdown: &[PipelineUid]) {
+        info!(
+            pipeline_uids=?pipelines_to_shutdown,
+            "shutdown indexing pipelines"
+        );
         let should_gc_ingest_api_queues = pipelines_to_shutdown
             .iter()
             .flat_map(|pipeline_uid| self.indexing_pipelines.get(pipeline_uid))

quickwit/quickwit-proto/src/metastore/mod.rs

@@ -16,7 +16,7 @@ use std::fmt;
 
 use quickwit_common::rate_limited_error;
 use quickwit_common::retry::Retryable;
-use quickwit_common::tower::MakeLoadShedError;
+use quickwit_common::tower::{MakeLoadShedError, TimeoutExceeded};
 use serde::{Deserialize, Serialize};
 
 use crate::types::{IndexId, IndexUid, QueueId, SourceId, SplitId};
@@ -187,6 +187,12 @@ impl From<sqlx::Error> for MetastoreError {
     }
 }
 
+impl From<TimeoutExceeded> for MetastoreError {
+    fn from(_: TimeoutExceeded) -> Self {
+        MetastoreError::Timeout("client".to_string())
+    }
+}
+
 impl ServiceError for MetastoreError {
     fn error_code(&self) -> ServiceErrorCode {
         match self {

quickwit/quickwit-serve/src/lib.rs

@@ -169,8 +169,9 @@ static METASTORE_GRPC_CLIENT_METRICS_LAYER: Lazy<GrpcMetricsLayer> =
 static METASTORE_GRPC_SERVER_METRICS_LAYER: Lazy<GrpcMetricsLayer> =
     Lazy::new(|| GrpcMetricsLayer::new("metastore", "server"));
 
-static GRPC_TIMEOUT_LAYER: Lazy<TimeoutLayer> =
-    Lazy::new(|| TimeoutLayer::new(Duration::from_secs(30)));
+static GRPC_INGESTER_SERVICE_TIMEOUT: Duration = Duration::from_secs(30);
+static GRPC_INDEXING_SERVICE_TIMEOUT: Duration = Duration::from_secs(30);
+static GRPC_METASTORE_SERVICE_TIMEOUT: Duration = Duration::from_secs(10);
 
 struct QuickwitServices {
     pub node_config: Arc<NodeConfig>,
@@ -459,16 +460,13 @@ pub async fn serve_quickwit(
             {
                 bail!("could not find any metastore node in the cluster");
             }
-            // These layers applies to all the RPCs of the metastore.
-            let shared_layers = ServiceBuilder::new()
-                .layer(RetryLayer::new(RetryPolicy::from(RetryParams::standard())))
-                .layer(METASTORE_GRPC_CLIENT_METRICS_LAYER.clone())
-                .layer(tower::limit::GlobalConcurrencyLimitLayer::new(
+            MetastoreServiceClient::tower()
+                .stack_layer(TimeoutLayer::new(GRPC_METASTORE_SERVICE_TIMEOUT))
+                .stack_layer(RetryLayer::new(RetryPolicy::from(RetryParams::standard())))
+                .stack_layer(METASTORE_GRPC_CLIENT_METRICS_LAYER.clone())
+                .stack_layer(tower::limit::GlobalConcurrencyLimitLayer::new(
                     get_metastore_client_max_concurrency(),
                 ))
-                .into_inner();
-            MetastoreServiceClient::tower()
-                .stack_layer(shared_layers)
                 .build_from_balance_channel(balance_channel, grpc_config.max_message_size)
         };
     // Instantiate a control plane server if the `control-plane` role is enabled on the node.
@@ -514,11 +512,11 @@ pub async fn serve_quickwit(
         None
     };
 
-    // Setup indexer pool.
+    // Setup the indexer pool to track cluster changes.
     setup_indexer_pool(
         &node_config,
         cluster.change_stream(),
-        indexer_pool.clone(),
+        indexer_pool,
         indexing_service_opt.clone(),
     );
 
@@ -943,7 +941,7 @@ async fn setup_ingest_v2(
                     } else {
                         let ingester_service = IngesterServiceClient::tower()
                             .stack_layer(INGEST_GRPC_CLIENT_METRICS_LAYER.clone())
-                            .stack_layer(GRPC_TIMEOUT_LAYER.clone())
+                            .stack_layer(TimeoutLayer::new(GRPC_INGESTER_SERVICE_TIMEOUT))
                             .build_from_channel(
                                 node.grpc_advertise_addr(),
                                 node.channel(),
@@ -1145,7 +1143,7 @@ fn setup_indexer_pool(
                     } else {
                         let client = IndexingServiceClient::tower()
                             .stack_layer(INDEXING_GRPC_CLIENT_METRICS_LAYER.clone())
-                            .stack_layer(GRPC_TIMEOUT_LAYER.clone())
+                            .stack_layer(TimeoutLayer::new(GRPC_INDEXING_SERVICE_TIMEOUT))
                             .build_from_channel(
                                 node.grpc_advertise_addr(),
                                 node.channel(),