From 08e41e9898a56486cf5a52df768231171b7f1cec Mon Sep 17 00:00:00 2001 From: Atsika Date: Sat, 15 Aug 2026 19:01:49 +0200 Subject: [PATCH] Make listener teardown exactly-once and bound the accept loop Two problems in the same lifecycle. A listener whose Accept keeps failing spun forever: it logged on every iteration with no delay while `listener list` still reported Running:true and no agent could ever arrive -- the zombie pattern already fixed in ReceiveLoop. The accept loop now backs off between failures and, after a bounded run of them, tears the listener down and marks it stopped. That gave a second goroutine the ability to stop a listener, and Running, Listener, and CancelFunc were plain fields mutated by StopListener while the CLI read them, with nothing making teardown happen once. Running is now an atomic.Bool claimed with CompareAndSwap, so exactly one of the CLI and the accept loop owns the shutdown, and listener/cancel move behind a mutex with the teardown path factored into teardownListener. Clearing the selected agent and listener stays in StopListener rather than moving into the shared helper: selection is CLI-owned state, and the accept goroutine writing it would trade one race for another. monitorAgent already drops agents without touching selection, so this matches. --- cmd/proxy/main.go | 200 +++++++++++++++++++++++++++++++++------------- 1 file changed, 146 insertions(+), 54 deletions(-) diff --git a/cmd/proxy/main.go b/cmd/proxy/main.go index 9e4bcb7..6032de2 100644 --- a/cmd/proxy/main.go +++ b/cmd/proxy/main.go @@ -57,13 +57,53 @@ type Config struct { } // ListenerState tracks the state of a listener. +// +// Two goroutines can decide to shut a listener down: the CLI, via StopListener, +// and the accept loop, when Accept has failed too many times in a row. running +// is the gate that makes exactly one of them the owner of the teardown. type ListenerState struct { - Config *ListenerConfig // listener configuration - Listener *aznet.Listener // aznet listener - ListenAddr string // address where listener is bound - Running bool // whether listener is running - StartedAt time.Time // when listener was started - CancelFunc context.CancelFunc // function to cancel the accept loop context + Config *ListenerConfig // listener configuration + ListenAddr string // address where listener is bound + StartedAt time.Time // when listener was started + + running atomic.Bool // whether listener is running + + mu sync.Mutex // guards listener and cancel + listener *aznet.Listener // aznet listener, nil once torn down + cancel context.CancelFunc // cancels the accept loop context, nil once torn down +} + +// IsRunning reports whether the listener is accepting connections. +func (s *ListenerState) IsRunning() bool { + return s.running.Load() +} + +// claimStop returns true for the single caller that wins the right to tear this +// listener down; every later caller gets false and must not touch it. +func (s *ListenerState) claimStop() bool { + return s.running.CompareAndSwap(true, false) +} + +// transport returns the live listener, or nil once it has been torn down. +func (s *ListenerState) transport() *aznet.Listener { + s.mu.Lock() + defer s.mu.Unlock() + return s.listener +} + +// closeTransport cancels the accept context and closes the listener. Idempotent. +func (s *ListenerState) closeTransport() { + s.mu.Lock() + cancel, listener := s.cancel, s.listener + s.cancel, s.listener = nil, nil + s.mu.Unlock() + + if cancel != nil { + cancel() + } + if listener != nil { + listener.Close() + } } // AgentConnection tracks a connected agent. @@ -200,7 +240,7 @@ func StartListener(listenerID string) error { // Check if listener already exists if val, ok := listeners.Load(listenerID); ok { state := val.(*ListenerState) - if state.Running { + if state.IsRunning() { return fmt.Errorf("listener '%s' is already running", listenerID) } } @@ -257,12 +297,12 @@ func StartListener(listenerID string) error { // Store listener state state := &ListenerState{ Config: listenerConfig, - Listener: listener, ListenAddr: listenAddr, - Running: true, StartedAt: time.Now(), - CancelFunc: cancel, + listener: listener, + cancel: cancel, } + state.running.Store(true) listeners.Store(listenerID, state) // Set this listener as the default when started @@ -284,7 +324,7 @@ func StopListener(listenerID string) error { } state := val.(*ListenerState) - if !state.Running { + if !state.IsRunning() { return fmt.Errorf("listener '%s' is not running", listenerID) } @@ -308,20 +348,16 @@ func StopListener(listenerID string) error { // This gives agents time to receive the shutdown notification time.Sleep(200 * time.Millisecond) - // Now cancel the accept loop context - if state.CancelFunc != nil { - state.CancelFunc() - } - - // Close the listener after FIN messages have been sent - if state.Listener != nil { - state.Listener.Close() + // Claim the teardown only now. Claiming before the FIN window would report the + // listener as stopped while it is still up, letting a restart bind a new + // listener under this ID whose agents this teardown would then disconnect. + if !state.claimStop() { + return fmt.Errorf("listener '%s' is not running", listenerID) } - // Update state - state.Running = false - state.Listener = nil - state.CancelFunc = nil + // Close the transport after FIN messages have been sent, then drop the agents + // that came in through it. + dropped := teardownListener(listenerID, state, "listener stopped") // Clear default selection if this was the selected listener if selectedListener == listenerID { @@ -329,41 +365,53 @@ func StopListener(listenerID string) error { log.Info().Msg("Default listener cleared (listener was stopped)") } - // Now fully close all connections and clean up - connectedAgents.Range(func(key, value interface{}) bool { - agent := value.(*AgentConnection) - if agent.ListenerID == listenerID { - agentID := key.(string) - - // Stop proxy if running - if val, ok := runningProxies.LoadAndDelete(agentID); ok { - if server, ok := val.(*proxy.ProxyServer); ok { - server.Stop() - } + // Selection is CLI-owned state, so it is cleared here rather than in + // teardownListener, which also runs on the accept goroutine. + for _, agentID := range dropped { + if selectedAgent == agentID { + selectedAgent = "" + // Update prompt to default if app is available + if app != nil { + app.SetPrompt("proxyblob » ") } + } + } - // Close connection (this will also send FIN if not already sent) - agent.Conn.Close() + log.Info().Str("listener_id", listenerID).Msg("Listener stopped") + return nil +} - // Remove from map - connectedAgents.Delete(key) +// teardownListener closes the listener's transport and disconnects every agent +// that arrived through it, returning their IDs. The caller must have won +// claimStop first, which is what keeps this to one execution per listener. +func teardownListener(listenerID string, state *ListenerState, reason string) []string { + state.closeTransport() - // Clear selected agent if it was this one - if selectedAgent == agentID { - selectedAgent = "" - // Update prompt to default if app is available - if app != nil { - app.SetPrompt("proxyblob » ") - } - } + var dropped []string + connectedAgents.Range(func(key, value interface{}) bool { + agent := value.(*AgentConnection) + if agent.ListenerID != listenerID { + return true + } + agentID := key.(string) - log.Info().Str("agent_id", agentID).Str("listener_id", listenerID).Msg("Agent disconnected (listener stopped)") + // Stop proxy if running + if val, ok := runningProxies.LoadAndDelete(agentID); ok { + if server, ok := val.(*proxy.ProxyServer); ok { + server.Stop() + } } + + // Close connection (this will also send FIN if not already sent) + agent.Conn.Close() + connectedAgents.Delete(key) + dropped = append(dropped, agentID) + + log.Info().Str("agent_id", agentID).Str("listener_id", listenerID).Str("reason", reason).Msg("Agent disconnected") return true }) - log.Info().Str("listener_id", listenerID).Msg("Listener stopped") - return nil + return dropped } // GenerateConnectionString creates a connection string for agents using the specified listener. @@ -374,13 +422,14 @@ func GenerateConnectionString(listenerID string, expiry time.Duration) (string, } state := val.(*ListenerState) - if !state.Running || state.Listener == nil { + listener := state.transport() + if !state.IsRunning() || listener == nil { return "", fmt.Errorf("listener '%s' is not running", listenerID) } // The expiry parameter is currently ignored by the library's ConnectionString method. // It uses the DefaultSASExpiry (24h) or the one set via WithSASExpiry during Listen. - connStr, err := state.Listener.ConnectionString() + connStr, err := listener.ConnectionString() if err != nil { return "", err } @@ -395,6 +444,17 @@ func GenerateConnectionString(listenerID string, expiry time.Duration) (string, // AcceptAgentLoopForListener accepts incoming agent connections for a specific listener and stores them. func AcceptAgentLoopForListener(ctx context.Context, listenerID string, listener net.Listener) { + // An Accept failure is usually transient -- a throttled poll, a dropped + // request -- so back off and retry. But retrying forever turns a listener + // whose transport is gone into a zombie: it logs on every iteration while + // `listener list` still reports it as running and no agent can ever arrive. + // Give up after a bounded run of failures and tear it down instead. + const maxAcceptFailures = 20 + const maxAcceptBackoff = 5 * time.Second + const maxAcceptBackoffShift = 6 + + failures := 0 + for { select { case <-ctx.Done(): @@ -413,13 +473,45 @@ func AcceptAgentLoopForListener(ctx context.Context, listenerID string, listener return } state := val.(*ListenerState) - if !state.Running { + if !state.IsRunning() { + return + } + + failures++ + if failures >= maxAcceptFailures { + log.Error(). + Err(err). + Int("consecutive", failures). + Str("listener_id", listenerID). + Msg("Accept failing persistently, stopping listener") + // Lost the race against a concurrent CLI stop: that caller owns + // the teardown, so just exit. + if state.claimStop() { + teardownListener(listenerID, state, "accept loop gave up") + } return } + log.Error().Err(err).Str("listener_id", listenerID).Msg("Failed to accept connection") + + shift := failures - 1 + if shift > maxAcceptBackoffShift { + shift = maxAcceptBackoffShift + } + backoff := time.Duration(100< maxAcceptBackoff { + backoff = maxAcceptBackoff + } + select { + case <-ctx.Done(): + return + case <-time.After(backoff): + } continue } + failures = 0 + // Generate a unique ID for this agent connection agentID := uuid.New().String() @@ -538,7 +630,7 @@ func ListListeners() []ListenerInfo { // Check if listener is running if val, ok := listeners.Load(listenerID); ok { state := val.(*ListenerState) - if state.Running { + if state.IsRunning() { info.Status = "running" info.StartedAt = state.StartedAt info.Protocol = state.Config.Driver @@ -832,7 +924,7 @@ func AddCommands(app *grumble.App) { } state := val.(*ListenerState) - if !state.Running { + if !state.IsRunning() { log.Error().Str("listener_id", listenerID).Msg("Listener is not running") return nil }