From a972fafb5f144df3211789601106377fa79cfbf5 Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Fri, 17 Sep 2021 17:49:58 +0800
Subject: [PATCH] Fix proxy ut may hang with dn graceful stop (#8124)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 internal/distributed/datanode/service.go | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/internal/distributed/datanode/service.go b/internal/distributed/datanode/service.go
index b962863cfa..e548b684ad 100644
--- a/internal/distributed/datanode/service.go
+++ b/internal/distributed/datanode/service.go
@@ -141,7 +141,21 @@ func (s *Server) Stop() error {
 	}
 	s.cancel()
 	if s.grpcServer != nil {
-		s.grpcServer.GracefulStop()
+		// make graceful stop has a timeout
+		stopped := make(chan struct{})
+		go func() {
+			s.grpcServer.GracefulStop()
+			close(stopped)
+		}()
+
+		t := time.NewTimer(10 * time.Second)
+		select {
+		case <-t.C:
+			// hard stop since grace timeout
+			s.grpcServer.Stop()
+		case <-stopped:
+			t.Stop()
+		}
 	}
 
 	err := s.datanode.Stop()