Merge pull request #31 from towerops-app/fix/pool-shutdown-timeout

fix: add timeout to worker pool shutdown to prevent blocking reconnec…

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()
 	}()

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
+	}
+}

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) {