graham/towerops-agent
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add worker pool with bounded concurrency for job dispatch

Browse source 
Replace unbounded goroutine spawning with fixed-size worker pools:
100 workers for SNMP, 20 for MikroTik, 50 for ping. Prevents file
descriptor exhaustion and device overload when handling thousands
of concurrent jobs. Pool provides backpressure when saturated.
This commit is contained in:
Graham McIntire 2026-02-12 09:50:17 -06:00
parent 07e7ed2745
commit 8b866ae946
No known key found for this signature in database
4 changed files with 152 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
agent.go
View file

@ -79,6 +79,13 @@ func runSession(ctx context.Context, baseURL, token string) error {
// Channel for serializing WebSocket writes
writeCh := make(chan []byte, 500)
// Worker pools — bounded concurrency for each job type
pools := &jobPools{
snmp: newWorkerPool(100),
mikrotik: newWorkerPool(20),
ping: newWorkerPool(50),
}
// Result channels
snmpResultCh := make(chan *pb.SnmpResult, 1000)
mikrotikResultCh := make(chan *pb.MikrotikResult, 1000)
@ -169,6 +176,9 @@ func runSession(ctx context.Context, baseURL, token string) error {
startTime := time.Now()
defer func() {
pools.snmp.stop()
pools.mikrotik.stop()
pools.ping.stop()
close(writeCh)
writerWg.Wait()
}()
@ -188,7 +198,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
slog.Warn("invalid message", "error", err)
continue
}
handleMessage(msg, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh)
handleMessage(msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh)
case result := <-snmpResultCh:
sendBinaryResult("result", result)
@ -236,6 +246,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
// handleMessage dispatches incoming channel messages.
func handleMessage(
msg channelMsg,
pools *jobPools,
snmpResultCh chan<- *pb.SnmpResult,
mikrotikResultCh chan<- *pb.MikrotikResult,
credTestResultCh chan<- *pb.CredentialTestResult,
@ -265,7 +276,7 @@ func handleMessage(
}
slog.Info("received jobs", "count", len(jobList.Jobs))
for _, job := range jobList.Jobs {
dispatchJob(job, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh)
dispatchJob(job, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh)
}
case "restart":
@ -291,9 +302,17 @@ func handleMessage(
}
}
// dispatchJob routes a job to the appropriate handler goroutine.
// jobPools holds the worker pools for each job type.
type jobPools struct {
snmp *workerPool
mikrotik *workerPool
ping *workerPool
}
// dispatchJob routes a job to the appropriate worker pool.
func dispatchJob(
job *pb.AgentJob,
pools *jobPools,
snmpResultCh chan<- *pb.SnmpResult,
mikrotikResultCh chan<- *pb.MikrotikResult,
credTestResultCh chan<- *pb.CredentialTestResult,
@ -303,14 +322,13 @@ func dispatchJob(
switch job.JobType {
case pb.JobType_MIKROTIK:
go executeMikrotikJob(job, mikrotikResultCh)
pools.mikrotik.submit(func() { executeMikrotikJob(job, mikrotikResultCh) })
case pb.JobType_TEST_CREDENTIALS:
go executeCredentialTest(job, credTestResultCh)
pools.snmp.submit(func() { executeCredentialTest(job, credTestResultCh) })
case pb.JobType_PING:
go executePingJob(job, monitoringCheckCh)
pools.ping.submit(func() { executePingJob(job, monitoringCheckCh) })
default:
// DISCOVER, POLL
go executeSnmpJob(job, snmpResultCh)
pools.snmp.submit(func() { executeSnmpJob(job, snmpResultCh) })
}
}

43
agent_test.go
View file

@ -75,6 +75,17 @@ func searchString(s, substr string) bool {
return false
}
func testPools(t *testing.T) *jobPools {
t.Helper()
p := &jobPools{
snmp: newWorkerPool(4),
mikrotik: newWorkerPool(4),
ping: newWorkerPool(4),
}
t.Cleanup(func() { p.snmp.stop(); p.mikrotik.stop(); p.ping.stop() })
return p
}
// makeJobPayload creates a base64-encoded protobuf job list payload.
func makeJobPayload(jobs ...*pb.AgentJob) json.RawMessage {
list := &pb.AgentJobList{Jobs: jobs}
@ -89,7 +100,7 @@ func TestHandleMessage(t *testing.T) {
mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
handleMessage(channelMsg{Event: "phx_reply", Payload: json.RawMessage(`{}`)}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "phx_reply", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Just verify it doesn't panic
})
@ -116,7 +127,7 @@ func TestHandleMessage(t *testing.T) {
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161},
})
handleMessage(channelMsg{Event: "jobs", Payload: payload}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Wait for goroutine to finish
select {
case <-snmpCh:
@ -130,7 +141,7 @@ func TestHandleMessage(t *testing.T) {
mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
handleMessage(channelMsg{Event: "jobs", Payload: json.RawMessage(`not json`)}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Should log error but not panic
})
@ -140,7 +151,7 @@ func TestHandleMessage(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
payload, _ := json.Marshal(map[string]string{"binary": "not-base64!!!"})
handleMessage(channelMsg{Event: "jobs", Payload: payload}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
})
t.Run("invalid protobuf", func(t *testing.T) {
@ -149,7 +160,7 @@ func TestHandleMessage(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})})
handleMessage(channelMsg{Event: "jobs", Payload: payload}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
})
t.Run("restart", func(t *testing.T) {
@ -163,7 +174,7 @@ func TestHandleMessage(t *testing.T) {
mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
handleMessage(channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
if exitCode != 0 {
t.Errorf("expected exit code 0, got %d", exitCode)
@ -185,7 +196,7 @@ func TestHandleMessage(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"})
handleMessage(channelMsg{Event: "update", Payload: payload}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
if calledURL != "https://example.com/agent" {
t.Errorf("expected update URL %q, got %q", "https://example.com/agent", calledURL)
@ -208,7 +219,7 @@ func TestHandleMessage(t *testing.T) {
monCh := make(chan *pb.MonitoringCheck, 1)
// Missing URL field
payload, _ := json.Marshal(map[string]string{"checksum": "abc123"})
handleMessage(channelMsg{Event: "update", Payload: payload}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
if called {
t.Error("selfUpdate should not be called with empty URL")
@ -228,7 +239,7 @@ func TestHandleMessage(t *testing.T) {
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent"})
handleMessage(channelMsg{Event: "update", Payload: payload}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Should log error but not panic
})
@ -237,7 +248,7 @@ func TestHandleMessage(t *testing.T) {
mtCh := make(chan *pb.MikrotikResult, 1)
credCh := make(chan *pb.CredentialTestResult, 1)
monCh := make(chan *pb.MonitoringCheck, 1)
handleMessage(channelMsg{Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh)
// Should just log and not panic
})
@ -263,7 +274,7 @@ func TestHandleMessage(t *testing.T) {
JobType: pb.JobType_DISCOVER,
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
})
handleMessage(channelMsg{Event: "discovery_job", Payload: payload}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "discovery_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
select {
case <-snmpCh:
case <-time.After(2 * time.Second):
@ -290,7 +301,7 @@ func TestHandleMessage(t *testing.T) {
JobType: pb.JobType_MIKROTIK,
MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"},
})
handleMessage(channelMsg{Event: "backup_job", Payload: payload}, snmpCh, mtCh, credCh, monCh)
handleMessage(channelMsg{Event: "backup_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh)
select {
case result := <-mtCh:
if result.Error != "" {
@ -319,7 +330,7 @@ func TestDispatchJob(t *testing.T) {
JobId: "mt1",
JobType: pb.JobType_MIKROTIK,
MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728},
}, snmpCh, mtCh, credCh, monCh)
}, testPools(t), snmpCh, mtCh, credCh, monCh)
select {
case result := <-mtCh:
@ -347,7 +358,7 @@ func TestDispatchJob(t *testing.T) {
JobId: "tc1",
JobType: pb.JobType_TEST_CREDENTIALS,
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
}, snmpCh, mtCh, credCh, monCh)
}, testPools(t), snmpCh, mtCh, credCh, monCh)
select {
case result := <-credCh:
@ -375,7 +386,7 @@ func TestDispatchJob(t *testing.T) {
JobId: "p1",
JobType: pb.JobType_PING,
SnmpDevice: &pb.SnmpDevice{Ip: "127.0.0.1"},
}, snmpCh, mtCh, credCh, monCh)
}, testPools(t), snmpCh, mtCh, credCh, monCh)
select {
case result := <-monCh:
@ -407,7 +418,7 @@ func TestDispatchJob(t *testing.T) {
JobId: "s1",
JobType: pb.JobType_POLL,
SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
}, snmpCh, mtCh, credCh, monCh)
}, testPools(t), snmpCh, mtCh, credCh, monCh)
select {
case <-snmpCh:

40
workerpool.go Normal file
View file

@ -0,0 +1,40 @@
package main
import "sync"
// workerPool is a fixed-size goroutine pool for executing tasks.
type workerPool struct {
tasks chan func()
wg sync.WaitGroup
once sync.Once
}
// newWorkerPool creates a pool with n worker goroutines.
func newWorkerPool(n int) *workerPool {
p := &workerPool{
tasks: make(chan func(), n*4),
}
p.wg.Add(n)
for range n {
go func() {
defer p.wg.Done()
for fn := range p.tasks {
fn()
}
}()
}
return p
}
// submit enqueues a task. Blocks if all workers are busy and the queue is full.
func (p *workerPool) submit(fn func()) {
p.tasks <- fn
}
// stop closes the task channel and waits for all workers to finish.
func (p *workerPool) stop() {
p.once.Do(func() {
close(p.tasks)
p.wg.Wait()
})
}

59
workerpool_test.go Normal file
View file

@ -0,0 +1,59 @@
package main
import (
"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(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(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
})
}