Graham McIntire
87606bb84a
Fix 1: Limit inbound job payload size to 4MB to prevent OOM from malicious server Fix 2: Worker pool submit accepts context and returns bool, prevents blocking event loop Fix 3: Panic recovery in worker pool goroutines prevents permanent worker death Fix 4: WebSocket handshake timeout (30s) prevents indefinite blocking on slow server Fix 5: Writer goroutine propagates errors to main loop via writeErrCh
113 lines
2.5 KiB
Go
113 lines
2.5 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestWorkerPool(t *testing.T) {
|
|
t.Run("executes all tasks", func(t *testing.T) {
|
|
pool := newWorkerPool(4)
|
|
defer pool.stop()
|
|
|
|
var count atomic.Int32
|
|
for i := 0; i < 100; i++ {
|
|
pool.submit(context.Background(), func() {
|
|
count.Add(1)
|
|
})
|
|
}
|
|
|
|
pool.stop()
|
|
if got := count.Load(); got != 100 {
|
|
t.Errorf("got %d completions, want 100", got)
|
|
}
|
|
})
|
|
|
|
t.Run("limits concurrency", func(t *testing.T) {
|
|
pool := newWorkerPool(2)
|
|
defer pool.stop()
|
|
|
|
var concurrent atomic.Int32
|
|
var maxConcurrent atomic.Int32
|
|
|
|
for i := 0; i < 20; i++ {
|
|
pool.submit(context.Background(), func() {
|
|
cur := concurrent.Add(1)
|
|
for {
|
|
old := maxConcurrent.Load()
|
|
if cur <= old || maxConcurrent.CompareAndSwap(old, cur) {
|
|
break
|
|
}
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
concurrent.Add(-1)
|
|
})
|
|
}
|
|
|
|
pool.stop()
|
|
if max := maxConcurrent.Load(); max > 2 {
|
|
t.Errorf("max concurrent was %d, want <= 2", max)
|
|
}
|
|
})
|
|
|
|
t.Run("stop is idempotent", func(t *testing.T) {
|
|
pool := newWorkerPool(2)
|
|
pool.stop()
|
|
pool.stop() // should not panic
|
|
})
|
|
}
|
|
|
|
func TestWorkerPoolRecoversPanic(t *testing.T) {
|
|
pool := newWorkerPool(1)
|
|
defer pool.stop()
|
|
|
|
// Submit a function that panics
|
|
pool.submit(context.Background(), func() { panic("boom") })
|
|
|
|
// Give the panic time to be processed
|
|
time.Sleep(50 * time.Millisecond)
|
|
|
|
// Submit a normal function — the worker should still be alive
|
|
done := make(chan struct{})
|
|
ok := pool.submit(context.Background(), func() { close(done) })
|
|
if !ok {
|
|
t.Fatal("expected submit to succeed after panic recovery")
|
|
}
|
|
|
|
select {
|
|
case <-done:
|
|
// Worker survived the panic
|
|
case <-time.After(2 * time.Second):
|
|
t.Error("timed out — worker did not survive panic")
|
|
}
|
|
}
|
|
|
|
func TestWorkerPoolSubmitRespectsContext(t *testing.T) {
|
|
pool := newWorkerPool(1) // 1 worker, queue capacity 4
|
|
defer pool.stop()
|
|
|
|
blocker := make(chan struct{})
|
|
|
|
// Occupy the single worker
|
|
pool.submit(context.Background(), func() { <-blocker })
|
|
|
|
// Fill the buffered queue (capacity = n*4 = 4)
|
|
for i := 0; i < 4; i++ {
|
|
pool.submit(context.Background(), func() { <-blocker })
|
|
}
|
|
|
|
// Now the queue is full and the worker is busy.
|
|
// Submit with a cancelled context should return false immediately.
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
ok := pool.submit(ctx, func() { t.Error("should not execute") })
|
|
if ok {
|
|
t.Error("expected submit to return false with cancelled context")
|
|
}
|
|
|
|
// Unblock everything for cleanup
|
|
close(blocker)
|
|
}
|