Fixed the issue where quorum retry could not catch non-timeout errors. (#259)

* Fixed the issue where quorum retry could not catch non-timeout errors

* remove log

* update

* fix

* fix

* add ut

* avoid too small sleep_time

* fix ut

Author: zhengchenyu

src/manager.rs

@@ -252,34 +252,44 @@ impl Manager {
 
             let result = tokio::time::timeout(timeout, client.quorum(request)).await;
 
+            let mut sleep_time = 100;
             match result {
-                Ok(response) => {
-                    return response;
+                Ok(Ok(response)) => {
+                    return Ok(response);
                 }
-                Err(e) => {
+                Ok(Err(e)) => {
                     info_with_replica!(
                         self.replica_id,
                         "lighthouse quorum failed. error: {}",
                         e.to_string()
                     );
+                    sleep_time = timeout.as_millis() as u64
+                        / (std::cmp::max(self.quorum_retries + 1, 1 as i64)) as u64;
+                    sleep_time = std::cmp::max(100, sleep_time);
+                }
+                Err(e) => {
+                    info_with_replica!(
+                        self.replica_id,
+                        "lighthouse quorum timeout. error: {}",
+                        e.to_string()
+                    );
+                }
+            }
 
-                    if retry_count == self.quorum_retries {
-                        return Err(Status::internal(format!(
-                            "lighthouse quorum failed after {} retries. error: {}",
-                            retry_count,
-                            e.to_string(),
-                        )));
-                    }
+            if retry_count == self.quorum_retries {
+                return Err(Status::internal(format!(
+                    "lighthouse quorum failed after {} retries.",
+                    retry_count,
+                )));
+            }
 
-                    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+            tokio::time::sleep(Duration::from_millis(sleep_time)).await;
 
-                    // Reset the client since the lighthouse server might have failed
-                    // If this also fails, consider increasing `connect_timeout`.
-                    let _ = self.create_lighthouse_client().await;
+            // Reset the client since the lighthouse server might have failed
+            // If this also fails, consider increasing `connect_timeout`.
+            let _ = self.create_lighthouse_client().await;
 
-                    retry_count += 1;
-                }
-            }
+            retry_count += 1;
         }
     }
 
@@ -607,6 +617,10 @@ fn compute_quorum_results(
 mod tests {
     use super::*;
     use crate::lighthouse::{Lighthouse, LighthouseOpt};
+    use crate::torchftpb::lighthouse_service_server::{LighthouseService, LighthouseServiceServer};
+    use crate::torchftpb::LighthouseHeartbeatResponse;
+    use tokio::net::TcpListener;
+    use tonic::codegen::tokio_stream::wrappers::TcpListenerStream;
 
     async fn should_commit(group_rank: i64, should_commit: bool) -> Result<ShouldCommitResponse> {
         let mut client = manager_client_new(
@@ -1078,4 +1092,114 @@ mod tests {
 
         Ok(())
     }
+
+    #[derive(Default)]
+    pub struct MockLighthouse {
+        request_count: Arc<Mutex<usize>>,
+        fail_count: usize,
+    }
+
+    impl MockLighthouse {
+        pub fn new(fail_count: usize) -> Self {
+            Self {
+                request_count: Arc::new(Mutex::new(0)),
+                fail_count,
+            }
+        }
+    }
+
+    #[tonic::async_trait]
+    impl LighthouseService for MockLighthouse {
+        async fn quorum(
+            &self,
+            request: Request<LighthouseQuorumRequest>,
+        ) -> Result<Response<LighthouseQuorumResponse>, Status> {
+            let mut count = self.request_count.lock().await;
+            *count += 1;
+            let current_count = *count;
+
+            if current_count <= self.fail_count {
+                return Err(Status::internal(format!(
+                    "simulated failure (request {}/{})",
+                    current_count, self.fail_count
+                )));
+            }
+
+            let requester = request.into_inner().requester.unwrap_or_default();
+
+            Ok(Response::new(LighthouseQuorumResponse {
+                quorum: Some(Quorum {
+                    participants: vec![requester],
+                    quorum_id: 1,
+                    created: None,
+                }),
+            }))
+        }
+
+        async fn heartbeat(
+            &self,
+            _request: Request<LighthouseHeartbeatRequest>,
+        ) -> Result<Response<LighthouseHeartbeatResponse>, Status> {
+            Ok(Response::new(LighthouseHeartbeatResponse {}))
+        }
+    }
+
+    pub async fn start_mock_lighthouse(
+        fail_count: usize,
+    ) -> Result<String, Box<dyn std::error::Error>> {
+        let listener = TcpListener::bind("[::]:0").await?;
+        let addr = listener.local_addr()?;
+        let mock_service = LighthouseServiceServer::new(MockLighthouse::new(fail_count));
+        let incoming = TcpListenerStream::new(listener);
+        tokio::spawn(async move {
+            Server::builder()
+                .add_service(mock_service)
+                .serve_with_incoming(incoming)
+                .await
+                .expect("Mock Lighthouse server failed to start");
+        });
+        Ok(format!(
+            "http://{}:{}",
+            gethostname::gethostname().into_string().unwrap(),
+            addr.port()
+        ))
+    }
+
+    #[tokio::test]
+    async fn test_get_quorum_when_lighthouse_down() -> Result<()> {
+        let addr = start_mock_lighthouse(1).await.unwrap();
+
+        let manager = Manager::new(
+            "rep_id".to_string(),
+            addr,
+            "localhost".to_string(),
+            "[::]:0".to_string(),
+            "store_addr".to_string(),
+            1,
+            Duration::from_millis(100),
+            Duration::from_secs(10),
+            1,
+        )
+        .await?;
+        let manager_fut = tokio::spawn(manager.clone().run());
+
+        let mut client = manager_client_new(manager.address(), Duration::from_secs(3)).await?;
+
+        let mut request = tonic::Request::new(ManagerQuorumRequest {
+            group_rank: 0,
+            step: 123,
+            checkpoint_metadata: "addr".to_string(),
+            shrink_only: false,
+            init_sync: true,
+            commit_failures: 3,
+        });
+        request.set_timeout(Duration::from_secs(3));
+        let resp = client.quorum(request).await?.into_inner();
+
+        manager_fut.abort();
+
+        assert_eq!(resp.quorum_id, 1);
+
+        Ok(())
+    }
 }