fix: ensure graceful shutdown by awaiting all worker goroutines
- Reorder RelayClient.Close() to signal workers first, close session to unblock calls, then wait for workers - Add waitGroup to RDClient for tracking listener workers and wait for them in Close() method - Prevents potential races and resource leaks during shutdown
lemon-mint committed
Oct 28, 2025 at 15:33 UTC
6f78a5a4c33808bf6ffb787a4d84cfafeabd1ff9
2 files changed
+13
-3
relaydns/client.go
+6
-2
@@ -92,12 +92,13 @@ func NewRelayClient(conn io.ReadWriteCloser) *RelayClient {
92
// Close는 서버와의 연결을 종료합니다.
93
func (g *RelayClient) Close() error {
94
log.Debug().Msg("[RelayClient] Closing relay client")
95
+
96
+ // Signal workers to stop
97
close(g.stopCh)
96
- g.waitGroup.Wait()
98
99
var errs []error
100
100
- // Close the session first
101
+ // Close the session first to unblock AcceptStream() calls
102
if g.sess != nil {
103
if err := g.sess.Close(); err != nil {
104
log.Error().Err(err).Msg("[RelayClient] Error closing yamux session")
@@ -105,6 +106,9 @@ func (g *RelayClient) Close() error {
106
}
107
}
108
109
+ // Wait for workers to finish after unblocking them
110
+ g.waitGroup.Wait()
111
+
112
// Then close the underlying connection
113
if g.conn != nil {
114
if err := g.conn.Close(); err != nil {
sdk/sdk.go
+7
-1
@@ -116,7 +116,8 @@ type RDClient struct {
116
relays map[string]*rdRelay
117
listeners map[string]*RDListener
118
119
- stopch chan struct{}
119
+ stopch chan struct{}
120
+ waitGroup sync.WaitGroup // Track all listener workers
121
}
122
123
var (
@@ -308,6 +309,7 @@ func (g *RDClient) Listen(cred *cryptoops.Credential, name string, alpns []strin
309
310
// Start listener worker for each relay
311
for _, relay := range g.relays {
312
+ g.waitGroup.Add(1)
313
go g.listenerWorker(relay)
314
}
315
@@ -316,6 +318,7 @@ func (g *RDClient) Listen(cred *cryptoops.Credential, name string, alpns []strin
318
}
319
320
func (g *RDClient) listenerWorker(server *rdRelay) {
321
+ defer g.waitGroup.Done()
322
log.Debug().Str("relay", server.addr).Msg("[SDK] Listener worker started")
323
324
for {
@@ -406,6 +409,9 @@ func (g *RDClient) Close() error {
409
g.relays = make(map[string]*rdRelay)
410
g.mu.Unlock()
411
412
+ // Wait for all listener workers to finish
413
+ g.waitGroup.Wait()
414
+
415
if len(errs) > 0 {
416
return errs[0]
417
}