tcp.go

@@ -16,6 +16,7 @@ import (
 	"go.unistack.org/micro/v3/register"
 	"go.unistack.org/micro/v3/server"
 	"golang.org/x/net/netutil"
+	"golang.org/x/sync/errgroup"
 )
 
 type tcpServer struct {
@@ -427,15 +428,27 @@ func (h *tcpServer) Stop() error {
 }
 
 func (h *tcpServer) gracefulStop() {
-	ctx, cancel := context.WithTimeout(context.Background(), h.opts.GracefulTimeout)
+	tm := time.NewTimer(h.opts.GracefulTimeout)
+	defer tm.Stop()
 
-	go func() {
+	g, gctx := errgroup.WithContext(context.Background())
+
+	g.Go(func() error {
 		h.opts.Wait.Wait()
-		cancel()
+		return nil
-	}()
+	})
 
-	// wait timeout or finish all connects
+	g.Go(func() error {
-	<-ctx.Done()
+		select {
+		case <-tm.C:
+			return nil
+		case <-gctx.Done():
+			return gctx.Err()
+		}
+	})
+
+	g.Wait()
+	return
 }
 
 func (h *tcpServer) String() string {