diff --git a/cmd/pulse-sensor-proxy/http_server.go b/cmd/pulse-sensor-proxy/http_server.go index 78223ea..e3163d3 100644 --- a/cmd/pulse-sensor-proxy/http_server.go +++ b/cmd/pulse-sensor-proxy/http_server.go @@ -304,8 +304,13 @@ func (h *HTTPServer) handleTemperature(w http.ResponseWriter, r *http.Request) { } } - // Acquire per-node concurrency lock - releaseNode := h.proxy.nodeGate.acquire(nodeName) + // Acquire per-node concurrency lock (context-aware to prevent goroutine leaks) + releaseNode, err := h.proxy.nodeGate.acquireContext(ctx, nodeName) + if err != nil { + log.Warn().Err(err).Str("node", nodeName).Msg("Request cancelled while waiting for node lock") + h.sendJSONError(w, http.StatusServiceUnavailable, "request cancelled while waiting for node") + return + } defer releaseNode() // Fetch temperature data via SSH with context timeout diff --git a/cmd/pulse-sensor-proxy/main.go b/cmd/pulse-sensor-proxy/main.go index 7fec79f..adbd9b5 100644 --- a/cmd/pulse-sensor-proxy/main.go +++ b/cmd/pulse-sensor-proxy/main.go @@ -1188,16 +1188,20 @@ func (p *Proxy) handleGetTemperatureV2(ctx context.Context, req *RPCRequest, log } } - // Acquire per-node concurrency lock (prevents multiple simultaneous requests to same node) - releaseNode := p.nodeGate.acquire(node) + // Acquire per-node concurrency lock (context-aware to prevent goroutine leaks) + releaseNode, err := p.nodeGate.acquireContext(ctx, node) + if err != nil { + logger.Warn().Err(err).Str("node", node).Msg("Request cancelled while waiting for node lock") + return nil, fmt.Errorf("request cancelled while waiting for node") + } defer releaseNode() logger.Debug().Str("node", node).Msg("Fetching temperature via SSH") // Fetch temperature data with timeout - ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + sshCtx, cancel := context.WithTimeout(ctx, 15*time.Second) defer cancel() - tempData, err := p.getTemperatureViaSSH(ctx, node) + tempData, err := p.getTemperatureViaSSH(sshCtx, node) if err != nil { logger.Warn().Err(err).Str("node", node).Msg("Failed to get temperatures") return nil, fmt.Errorf("failed to get temperatures: %w", err) diff --git a/cmd/pulse-sensor-proxy/throttle.go b/cmd/pulse-sensor-proxy/throttle.go index 8d39d73..807ef27 100644 --- a/cmd/pulse-sensor-proxy/throttle.go +++ b/cmd/pulse-sensor-proxy/throttle.go @@ -1,6 +1,7 @@ package main import ( + "context" "fmt" "sync" "time" @@ -253,6 +254,7 @@ func newNodeGate() *nodeGate { // acquire gets exclusive access to make requests to a node // Returns a release function that must be called when done +// Deprecated: Use acquireContext for context-aware acquisition func (g *nodeGate) acquire(node string) func() { g.mu.Lock() lock := g.inFlight[node] @@ -279,3 +281,42 @@ func (g *nodeGate) acquire(node string) func() { g.mu.Unlock() } } + +// acquireContext gets exclusive access to make requests to a node with context cancellation support. +// Returns a release function and nil error on success. +// Returns nil and context error if context is cancelled while waiting. +func (g *nodeGate) acquireContext(ctx context.Context, node string) (func(), error) { + g.mu.Lock() + lock := g.inFlight[node] + if lock == nil { + lock = &nodeLock{ + guard: make(chan struct{}, 1), + } + g.inFlight[node] = lock + } + lock.refCount++ + g.mu.Unlock() + + // Wait for exclusive access OR context cancellation + select { + case lock.guard <- struct{}{}: + return func() { + <-lock.guard + g.mu.Lock() + lock.refCount-- + if lock.refCount == 0 { + delete(g.inFlight, node) + } + g.mu.Unlock() + }, nil + case <-ctx.Done(): + // Clean up refCount since we're not proceeding + g.mu.Lock() + lock.refCount-- + if lock.refCount == 0 { + delete(g.inFlight, node) + } + g.mu.Unlock() + return nil, ctx.Err() + } +}