tcp.go

@@ -2,7 +2,6 @@
 package tcp // import "go.unistack.org/micro-server-tcp/v3"
 
 import (
-	"context"
 	"crypto/tls"
 	"fmt"
 	"net"
@@ -16,7 +15,6 @@ 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 {
@@ -28,8 +26,6 @@ type tcpServer struct {
 	sync.RWMutex
 	registered bool
 	init       bool
-
-	wg sync.WaitGroup
 }
 
 func (h *tcpServer) newCodec(ct string) (codec.Codec, error) {
@@ -431,24 +427,17 @@ func (h *tcpServer) gracefulStop() {
 	tm := time.NewTimer(h.opts.GracefulTimeout)
 	defer tm.Stop()
 
-	g, gctx := errgroup.WithContext(context.Background())
+	done := make(chan struct{})
-
+	go func() {
-	g.Go(func() error {
 		h.opts.Wait.Wait()
-		return nil
+		done <- struct{}{}
-	})
+	}()
 
-	g.Go(func() error {
+	select {
-		select {
+	case <-tm.C:
-		case <-tm.C:
+	case <-done:
-			return nil
+		close(done)
-		case <-gctx.Done():
+	}
-			return gctx.Err()
-		}
-	})
-
-	g.Wait()
-	return
 }
 
 func (h *tcpServer) String() string {
@@ -501,8 +490,8 @@ func (h *tcpServer) serve(ln net.Listener, hd Handler) {
 
 		h.opts.Wait.Add(1)
 		go func() {
-			defer h.opts.Wait.Done()
 			hd.Serve(c)
+			h.opts.Wait.Done()
 		}()
 	}
 }