fix: resolve goroutine hang/leak on shutdown
fatheradvisor committed
Mar 6, 2026 at 01:27 UTC
0a9772617622aaf8b3fc34500e8d12cb173ec165
5 files changed
+60
-1
cmd/portal-tunnel/main.go
+4
@@ -243,6 +243,10 @@ func proxyConnection(ctx context.Context, localAddr string, relayConn net.Conn)
243
if err != nil {
244
log.Debug().Err(err).Msg("local->relay copy ended")
245
}
246
+ // TLS does not support half-close; full close unblocks the relay→local goroutine.
247
+ if closeErr := relayConn.Close(); closeErr != nil {
248
+ log.Debug().Err(closeErr).Msg("failed to close relay conn after local->relay copy")
249
+ }
250
errCh <- err
251
}()
252
cmd/relay-server/main.go
+4
-1
@@ -176,13 +176,16 @@ func runServer(cfg relayServerConfig) error {
176
if err := serv.Start(); err != nil {
177
return fmt.Errorf("start relay server: %w", err)
178
}
179
- defer serv.Stop()
179
180
apiServ := serveAPI(fmt.Sprintf(":%d", cfg.AdminPort), serv, admin, frontend, cfg, stop)
181
182
<-ctx.Done()
183
log.Info().Msg("[server] shutting down...")
184
185
+ // Stop relay first: drain idle reverse conns + close active SNI conns
186
+ // so that all HTTP handlers blocked on HandleConnect.Wait() can return.
187
+ serv.Stop()
188
+
189
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
190
defer cancel()
191
if apiServ != nil {
portal/relay.go
+3
@@ -215,6 +215,9 @@ func (g *RelayServer) Stop() {
215
if g.leaseManager != nil {
216
g.leaseManager.Stop()
217
}
218
+ if g.reverseHub != nil {
219
+ g.reverseHub.Shutdown()
220
+ }
221
if g.sniRouter != nil {
222
if err := g.sniRouter.Stop(); err != nil {
223
log.Warn().Err(err).Msg("[RelayServer] Failed to stop SNI router")
portal/reverse_hub.go
+36
@@ -105,7 +105,9 @@ type ReverseHub struct {
105
authorizer func(leaseID, token string) bool
106
ipBanChecker func(ip string) bool
107
onAccepted func(leaseID, ip string)
108
+ stopCh chan struct{}
109
mu sync.RWMutex
110
+ stopOnce sync.Once
111
}
112
113
// NewReverseHub creates a new reverse connection hub.
@@ -113,6 +115,7 @@ func NewReverseHub() *ReverseHub {
115
return &ReverseHub{
116
pools: make(map[string]chan *ReverseConn),
117
dropped: make(map[string]struct{}),
118
+ stopCh: make(chan struct{}),
119
}
120
}
121
@@ -377,12 +380,45 @@ func (h *ReverseHub) closeConn(conn net.Conn, debugCloseMessage string) {
380
}
381
}
382
383
+// Shutdown closes the stop channel and drains all pools, causing idle
384
+// HandleConnect goroutines to unblock and return.
385
+func (h *ReverseHub) Shutdown() {
386
+ h.stopOnce.Do(func() {
387
+ close(h.stopCh)
388
+
389
+ h.mu.Lock()
390
+ pools := h.pools
391
+ h.pools = make(map[string]chan *ReverseConn)
392
+ h.mu.Unlock()
393
+
394
+ for _, pool := range pools {
395
+ drainPool(pool)
396
+ }
397
+ })
398
+}
399
+
400
+func drainPool(pool chan *ReverseConn) {
401
+ for {
402
+ select {
403
+ case conn := <-pool:
404
+ if conn != nil {
405
+ conn.Close()
406
+ }
407
+ default:
408
+ return
409
+ }
410
+ }
411
+}
412
+
413
func (h *ReverseHub) keepAliveWhileIdle(conn *ReverseConn, leaseID string) {
414
ticker := time.NewTicker(ReverseIdleKeepaliveInterval)
415
defer ticker.Stop()
416
417
for {
418
select {
419
+ case <-h.stopCh:
420
+ conn.Close()
421
+ return
422
case <-conn.done:
423
return
424
case <-conn.active:
portal/sni/router.go
+13
@@ -39,6 +39,7 @@ type Router struct {
39
listener net.Listener
40
routes map[string]*Route
41
leases map[string]*Route
42
+ conns map[net.Conn]struct{}
43
onConnection func(conn net.Conn, route *Route)
44
onNoRoute func(conn net.Conn, sni string) bool
45
stopCh chan struct{}
@@ -54,6 +55,7 @@ func NewRouter(addr string) *Router {
55
addr: addr,
56
routes: make(map[string]*Route),
57
leases: make(map[string]*Route),
58
+ conns: make(map[net.Conn]struct{}),
59
stopCh: make(chan struct{}),
60
}
61
}
@@ -237,6 +239,9 @@ func (r *Router) Stop() error {
239
if r.listener != nil {
240
closeWithDebugLog(r.listener, "[SNI] failed to close listener")
241
}
242
+ for conn := range r.conns {
243
+ closeWithDebugLog(conn, "[SNI] failed to close tracked connection")
244
+ }
245
r.mu.Unlock()
246
})
247
@@ -334,7 +339,15 @@ func (r *Router) handleConnection(clientConn net.Conn) {
339
onConnection := r.getConnectionHandler()
340
341
if onConnection != nil {
342
+ r.mu.Lock()
343
+ r.conns[wrappedConn] = struct{}{}
344
+ r.mu.Unlock()
345
+
346
onConnection(wrappedConn, route)
347
+
348
+ r.mu.Lock()
349
+ delete(r.conns, wrappedConn)
350
+ r.mu.Unlock()
351
return
352
}
353