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