transport: fixed a goroutine leak issue in tcp.go (#323)

#322
lni · Sep 25, 2023 · a6b5353 · a6b5353
1 parent d16220b
commit a6b5353
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/internal/transport/tcp.go b/internal/transport/tcp.go
@@ -473,15 +473,23 @@ func (t *TCP) Start() error {
 				panic(err)
 			}
 			var once sync.Once
+			connCloseCh := make(chan struct{})
 			closeFn := func() {
 				once.Do(func() {
+					select {
+						case connCloseCh <- struct{}{}:
+						default:
+					}
 					if err := conn.Close(); err != nil {
 						plog.Errorf("failed to close the connection %v", err)
 					}
 				})
 			}
 			t.connStopper.RunWorker(func() {
-				<-t.stopper.ShouldStop()
+				select {
+				case <-t.stopper.ShouldStop():
+				case <-connCloseCh:
+				}
 				closeFn()
 			})
 			t.connStopper.RunWorker(func() {