From 9d321b12b82b85713c8b353c681087fcfaac6a73 Mon Sep 17 00:00:00 2001
From: Evstigneev Denis <danteevstigneev@yandex.ru>
Date: Wed, 6 Mar 2024 18:44:45 +0300
Subject: [PATCH] add errGroup

---
 tcp.go | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/tcp.go b/tcp.go
index 2efa88a..eb558da 100644
--- a/tcp.go
+++ b/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
-	<-ctx.Done()
+	g.Go(func() error {
+		select {
+		case <-tm.C:
+			return nil
+		case <-gctx.Done():
+			return gctx.Err()
+		}
+	})
+
+	g.Wait()
+	return
 }
 
 func (h *tcpServer) String() string {