add timeout for gs.GracefulStop (#813) (#817)

pkg/util/util.go

@@ -280,3 +280,22 @@ func AdjustDuration(v *time.Duration, defValue time.Duration) {
 		*v = defValue
 	}
 }
+
+// WaitUntilTimeout creates a goroutine to run fn, and then gives up waiting for the goroutine to exit When it timeouts
+func WaitUntilTimeout(name string, fn func(), timeout time.Duration) {
+	fName := zap.String("name", name)
+	exited := make(chan struct{})
+	go func() {
+		defer func() {
+			log.Info("goroutine exit by itself (with GoAndAbortGoroutine help)", fName)
+			close(exited)
+		}()
+		fn()
+	}()
+
+	select {
+	case <-time.After(timeout):
+		log.Info("abort goroutine (with GoAndAbortGoroutine help)", fName)
+	case <-exited:
+	}
+}

pkg/util/util_test.go

@@ -326,3 +326,17 @@ func (s *retryCtxSuite) TestSuccessAfterRetry(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(callCount, Equals, 2)
 }
+
+type waitUntilTimeoutSuit struct{}
+
+var _ = Suite(&waitUntilTimeoutSuit{})
+
+func (s *waitUntilTimeoutSuit) TestGoAndAbortGoroutine(c *C) {
+	var called bool
+	WaitUntilTimeout("test", func() {
+		c := make(chan struct{})
+		called = true
+		<-c
+	}, time.Second)
+	c.Assert(called, IsTrue)
+}

pump/server.go

@@ -875,8 +875,10 @@ func (s *Server) Close() {
 
 	close(s.pullClose)
 	// stop the gRPC server
-	s.gs.GracefulStop()
-	log.Info("grpc is stopped")
+	util.WaitUntilTimeout("grpc_server.GracefulStop", func() {
+		s.gs.GracefulStop()
+		log.Info("grpc is stopped")
+	}, 10*time.Second)
 
 	if err := s.storage.Close(); err != nil {
 		log.Error("close storage failed", zap.Error(err))