go.mod

@@ -3,7 +3,7 @@ module go.unistack.org/micro-server-tcp/v3
 go 1.20
 
 require (
-	go.unistack.org/micro/v3 v3.10.42
+	go.unistack.org/micro/v3 v3.10.51
 	golang.org/x/net v0.22.0
 )
 

go.sum

@@ -3,8 +3,8 @@ github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg
 github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
-go.unistack.org/micro/v3 v3.10.42 h1:A0nA6WT6wNq5fyQyzliX70Bj5/SGj5kadLSOySX4hro=
-go.unistack.org/micro/v3 v3.10.42/go.mod h1:CSmEf5ddmft94MyKHnUSMM0W5dpmmTVbgImbgQWV5Ak=
+go.unistack.org/micro/v3 v3.10.51 h1:7JlgbJDXA4+9zyk5EJ5KqvRCeMA4htu0OofntiN+hFE=
+go.unistack.org/micro/v3 v3.10.51/go.mod h1:erMgt3Bl7vQQ0e9UpQyR5NlLiZ9pKeEJ9+1tfYFaqUg=
 golang.org/x/net v0.22.0 h1:9sGLhx7iRIHEiX0oAJ3MRZMUCElJgy7Br1nO+AMN3Tc=
 golang.org/x/net v0.22.0/go.mod h1:JKghWKKOSdJwpW2GEx0Ja7fmaKnMsbu+MWVZTokSYmg=
 golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=

tcp.go

@@ -4,6 +4,7 @@ package tcp // import "go.unistack.org/micro-server-tcp/v3"
 import (
 	"crypto/tls"
 	"fmt"
+	"golang.org/x/net/context"
 	"net"
 	"sort"
 	"sync"
@@ -402,6 +403,7 @@ func (h *tcpServer) Start() error {
 		}
 
 		h.gracefulStop()
+
 		ch <- ts.Close()
 
 		// deregister
@@ -424,19 +426,10 @@ func (h *tcpServer) Stop() error {
 }
 
 func (h *tcpServer) gracefulStop() {
-	tm := time.NewTimer(h.opts.GracefulTimeout)
-	defer tm.Stop()
+	ctx, cancel := context.WithTimeout(context.Background(), h.opts.GracefulTimeout)
+	defer cancel()
 
-	done := make(chan struct{})
-	go func() {
-		h.opts.Wait.Wait()
-		close(done)
-	}()
-
-	select {
-	case <-tm.C:
-	case <-done:
-	}
+	h.opts.Wait.WaitContext(ctx)
 }
 
 func (h *tcpServer) String() string {