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

* add timeout for gs.GracefulStop

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

@@ -879,8 +879,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))