From 51e9345f1f9c5e9a46a00a5cfceed68885486e28 Mon Sep 17 00:00:00 2001
From: Graham McIntire <graham@mcintire.me>
Date: Sun, 22 Mar 2026 16:51:45 -0500
Subject: [PATCH] fix: add timeout to worker pool shutdown to prevent blocking
 reconnection

When the WebSocket connection drops, session cleanup waits for all
worker pools to drain. SNMP/MikroTik operations against unresponsive
devices can block for 30+ seconds per worker, preventing the agent
from reconnecting during that time.

Add stopWithTimeout(5s) so pool shutdown is bounded. If workers don't
finish within the timeout, they're abandoned and reconnection proceeds
immediately.
---
 agent.go           | 17 +++++++++++++----
 workerpool.go      | 18 ++++++++++++++++++
 workerpool_test.go | 21 +++++++++++++++++++++
 3 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/agent.go b/agent.go
index c0f37a0..8ae40e9 100644
--- a/agent.go
+++ b/agent.go
@@ -251,10 +251,19 @@ func runSession(ctx context.Context, baseURL, token string) error {
 	startTime := time.Now()
 
 	defer func() {
-		pools.snmp.stop()
-		pools.mikrotik.stop()
-		pools.ping.stop()
-		pools.checks.stop()
+		const poolShutdownTimeout = 5 * time.Second
+		if !pools.snmp.stopWithTimeout(poolShutdownTimeout) {
+			slog.Warn("snmp pool shutdown timed out, abandoning in-flight jobs")
+		}
+		if !pools.mikrotik.stopWithTimeout(poolShutdownTimeout) {
+			slog.Warn("mikrotik pool shutdown timed out, abandoning in-flight jobs")
+		}
+		if !pools.ping.stopWithTimeout(poolShutdownTimeout) {
+			slog.Warn("ping pool shutdown timed out, abandoning in-flight jobs")
+		}
+		if !pools.checks.stopWithTimeout(poolShutdownTimeout) {
+			slog.Warn("checks pool shutdown timed out, abandoning in-flight jobs")
+		}
 		close(writeCh)
 		writerWg.Wait()
 	}()
diff --git a/workerpool.go b/workerpool.go
index 957ef8a..e6ed4a4 100644
--- a/workerpool.go
+++ b/workerpool.go
@@ -5,6 +5,7 @@ import (
 	"log/slog"
 	"runtime/debug"
 	"sync"
+	"time"
 )
 
 // workerPool is a fixed-size goroutine pool for executing tasks.
@@ -55,3 +56,20 @@ func (p *workerPool) stop() {
 		p.wg.Wait()
 	})
 }
+
+// stopWithTimeout closes the task channel and waits up to timeout for workers
+// to finish. Returns true if all workers completed, false if the timeout was
+// reached and some workers were abandoned.
+func (p *workerPool) stopWithTimeout(timeout time.Duration) bool {
+	done := make(chan struct{})
+	go func() {
+		p.stop()
+		close(done)
+	}()
+	select {
+	case <-done:
+		return true
+	case <-time.After(timeout):
+		return false
+	}
+}
diff --git a/workerpool_test.go b/workerpool_test.go
index 9deb518..3e61b3e 100644
--- a/workerpool_test.go
+++ b/workerpool_test.go
@@ -57,6 +57,27 @@ func TestWorkerPool(t *testing.T) {
 		pool.stop()
 		pool.stop() // should not panic
 	})
+
+	t.Run("stopWithTimeout returns true when workers finish in time", func(t *testing.T) {
+		pool := newWorkerPool(2)
+		pool.submit(context.Background(), func() {
+			time.Sleep(10 * time.Millisecond)
+		})
+		if !pool.stopWithTimeout(time.Second) {
+			t.Error("expected stopWithTimeout to return true")
+		}
+	})
+
+	t.Run("stopWithTimeout returns false when workers exceed timeout", func(t *testing.T) {
+		pool := newWorkerPool(1)
+		blocker := make(chan struct{})
+		pool.submit(context.Background(), func() { <-blocker })
+
+		if pool.stopWithTimeout(50 * time.Millisecond) {
+			t.Error("expected stopWithTimeout to return false")
+		}
+		close(blocker) // unblock for cleanup
+	})
 }
 
 func TestWorkerPoolRecoversPanic(t *testing.T) {