package distworker import ( "strconv" "sync" "sync/atomic" "testing" "time" "github.com/Jeffail/tunny" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "rsgit.ru/rsmon/rsmon/internal/wire" ) // newTestRunner builds a runner with a deterministic executor and the // fixed dispatcher/pool layout. It registers a cleanup hook that drains // the runner so tests can leak-free. func newTestRunner(t *testing.T, maxConc, poolSize int, fn func(payload interface{}) interface{}) *Runner { t.Helper() r := NewRunner(&Config{MaxConcurrency: maxConc}) r.executor = fn require.NotNil(t, r.executor) atomic.StoreInt64(&r.concurrency, int64(poolSize)) r.jobQueue = make(chan wire.CheckJob, r.queueCapacity()) r.results = make(chan resultEnvelope, r.queueCapacity()) r.pool = tunny.NewFunc(poolSize, r.executor) for i := 0; i < maxConc; i++ { r.wg.Add(1) go r.dispatcher() } t.Cleanup(func() { r.Stop() r.wg.Wait() r.pool.Close() close(r.results) }) return r } func TestQueueCapacityScalesWithMaxConcurrency(t *testing.T) { cases := []struct { maxConc int wantCapacity int }{ // 2*maxConc, floored at minQueueCapacity so a small pool still // has backpressure headroom. {maxConc: 1, wantCapacity: minQueueCapacity}, {maxConc: 4, wantCapacity: minQueueCapacity}, {maxConc: 16, wantCapacity: 2 * 16}, {maxConc: 64, wantCapacity: 2 * 64}, } for _, tc := range cases { t.Run("max="+strconv.Itoa(tc.maxConc), func(t *testing.T) { r := NewRunner(&Config{MaxConcurrency: tc.maxConc}) assert.Equal(t, tc.wantCapacity, r.QueueCapacity(), "queue capacity should scale with max concurrency") }) } } func TestDispatcherRunsJobsConcurrently(t *testing.T) { const ( maxConc = 8 poolSize = 4 jobCount = 12 hold = 80 * time.Millisecond ) var ( inFlight atomic.Int64 peak atomic.Int64 ) executor := func(payload interface{}) interface{} { cur := inFlight.Add(1) for { p := peak.Load() if cur <= p || peak.CompareAndSwap(p, cur) { break } } time.Sleep(hold) inFlight.Add(-1) job := payload.(wire.CheckJob) return []wire.CheckResultReport{{ JobID: job.JobID, CheckID: job.CheckID, State: "OK", }} } r := newTestRunner(t, maxConc, poolSize, executor) // Drain the results channel so dispatchers do not block. var drainWG sync.WaitGroup drainWG.Add(1) go func() { defer drainWG.Done() for i := 0; i < jobCount; i++ { select { case <-r.results: case <-r.stopCh: return } } }() for i := 0; i < jobCount; i++ { job := wire.CheckJob{ JobID: "job-" + strconv.Itoa(i), CheckID: int64(i + 1), Kind: "http", Host: "example.com", } require.True(t, r.Enqueue(job)) } drainWG.Wait() // The tunny.Pool size (poolSize) limits how many jobs run in // parallel, so the peak should be at most poolSize and at least 2 // (otherwise the test would pass on a serial pool). observed := peak.Load() assert.GreaterOrEqual(t, observed, int64(2), "expected concurrent execution, observed peak=%d", observed) assert.LessOrEqual(t, observed, int64(poolSize), "peak should be bounded by pool size, observed peak=%d", observed) } func TestEnqueueRespectsBackpressure(t *testing.T) { // Use a small queue with no dispatchers consuming it, so the bounded // channel is the only source of backpressure. This is the cleanest // way to assert that Enqueue parks when the buffer is full. const queueCap = 4 r := NewRunner(&Config{MaxConcurrency: 2}) r.jobQueue = make(chan wire.CheckJob, queueCap) r.results = make(chan resultEnvelope, queueCap) t.Cleanup(func() { r.Stop() close(r.results) }) // Fill the bounded queue. for i := 0; i < queueCap; i++ { require.True(t, r.Enqueue(wire.CheckJob{JobID: "prefill-" + strconv.Itoa(i)})) } assert.Equal(t, queueCap, r.QueueDepth(), "queue should be full after %d enqueues", queueCap) // The next Enqueue must block because the queue is full. enqueueDone := make(chan bool, 1) go func() { enqueueDone <- r.Enqueue(wire.CheckJob{JobID: "blocking"}) }() select { case got := <-enqueueDone: t.Fatalf("Enqueue returned %v while the queue was full; expected backpressure", got) case <-time.After(50 * time.Millisecond): // expected: still parked } // Free a slot and confirm the parked Enqueue unblocks. select { case <-r.jobQueue: case <-r.stopCh: t.Fatal("runner stopped unexpectedly") } select { case ok := <-enqueueDone: assert.True(t, ok, "Enqueue should succeed once a slot is free") case <-time.After(time.Second): t.Fatal("Enqueue did not unblock after slot was freed") } } func TestApplyInitResizesPool(t *testing.T) { executor := func(payload interface{}) interface{} { return []wire.CheckResultReport{} } r := newTestRunner(t, 8, 1, executor) // Replace the pool with a proxy that records SetSize calls. The // proxy delegates Close back to the underlying pool so the test // runner cleanup only closes the pool once. original := r.pool proxy := newSizeProbe(original) r.pool = proxy r.applyInit(&wire.WorkerInit{Concurrency: 5, WorkerID: "w-1"}) assert.Equal(t, 5, r.Concurrency(), "applyInit should update pool size") require.NotEmpty(t, proxy.sizes, "expected pool.SetSize to be called") assert.Equal(t, 5, proxy.sizes[len(proxy.sizes)-1]) // Concurrency above maxConcurrency should be clamped. r.applyInit(&wire.WorkerInit{Concurrency: 999, WorkerID: "w-1"}) assert.Equal(t, r.MaxConcurrency(), r.Concurrency(), "applyInit should clamp concurrency to maxConcurrency") } // sizeProbePool wraps a jobPool and records SetSize calls. Close is // forwarded to the wrapped pool so cleanup happens exactly once. type sizeProbePool struct { jobPool sizes []int } func newSizeProbe(p jobPool) *sizeProbePool { return &sizeProbePool{jobPool: p} } func (s *sizeProbePool) SetSize(n int) { s.sizes = append(s.sizes, n) s.jobPool.SetSize(n) } func TestApplyInitIgnoresNonPositive(t *testing.T) { executor := func(payload interface{}) interface{} { return []wire.CheckResultReport{} } r := newTestRunner(t, 4, 1, executor) // newTestRunner mirrors Start()'s initial concurrency of 1. require.Equal(t, 1, r.Concurrency()) // Concurrency = 0 must not break the runner or change its size. r.applyInit(&wire.WorkerInit{Concurrency: 0}) assert.Equal(t, 1, r.Concurrency(), "non-positive concurrency should be ignored") // Concurrency = -5 should also be ignored. r.applyInit(&wire.WorkerInit{Concurrency: -5}) assert.Equal(t, 1, r.Concurrency()) } // TestApplyInitStoresCredentialsInMemory verifies that WorkerInit.Credentials // is stored on the Runner via applyInit and is retrievable through // Credentials(). It also confirms that a subsequent applyInit with nil // credentials replaces the previous value. func TestApplyInitStoresCredentialsInMemory(t *testing.T) { executor := func(payload interface{}) interface{} { return []wire.CheckResultReport{} } r := newTestRunner(t, 4, 1, executor) // Before any applyInit, Credentials() returns nil. assert.Nil(t, r.Credentials(), "Credentials() must return nil before any init") want := &wire.NotificationCredentials{ SMTP: []wire.SMTPCredential{{ ID: 11, Name: "primary", Server: "smtp.example.com", Port: 587, Login: "alerts@example.com", Password: "smtp-password-xyz", }}, Telegram: []wire.TelegramCredential{{ ID: 22, Name: "main-bot", Token: "bot-token-9876543210:ABCDEFG", }}, } r.applyInit(&wire.WorkerInit{ WorkerID: "w-1", Concurrency: 2, Credentials: want, }) got := r.Credentials() require.NotNil(t, got, "Credentials() must return non-nil after applyInit") require.Len(t, got.SMTP, 1) require.Len(t, got.Telegram, 1) assert.Equal(t, "primary", got.SMTP[0].Name) assert.Equal(t, "smtp-password-xyz", got.SMTP[0].Password) assert.Equal(t, "main-bot", got.Telegram[0].Name) assert.Equal(t, "bot-token-9876543210:ABCDEFG", got.Telegram[0].Token) // A subsequent applyInit with nil Credentials replaces the value. r.applyInit(&wire.WorkerInit{WorkerID: "w-1", Concurrency: 2}) assert.Nil(t, r.Credentials(), "Credentials() must return nil after applyInit with nil Credentials") } func TestNotificationHeartbeatCounters(t *testing.T) { r := NewRunner(&Config{MaxConcurrency: 1}) atomic.StoreInt64(&r.notifyDepth, 2) atomic.StoreInt64(&r.notifyActive, 3) assert.Equal(t, 5, r.ActiveNotifications()) assert.Equal(t, 2, r.NotificationQueueDepth()) } // TestApplyInitStoresURLInMemory verifies that the URL field added in // Task 2 is stored on the Runner via applyInit and exposed through // URL(). A subsequent applyInit with empty URL replaces the previous // value (consistent with the Credentials contract). func TestApplyInitStoresURLInMemory(t *testing.T) { executor := func(payload interface{}) interface{} { return []wire.CheckResultReport{} } r := newTestRunner(t, 4, 1, executor) // Before any applyInit, URL() returns empty. assert.Equal(t, "", r.URL(), "URL() must return empty before any init") r.applyInit(&wire.WorkerInit{ WorkerID: "w-1", Concurrency: 2, URL: "https://worker-eu.example.com", }) assert.Equal(t, "https://worker-eu.example.com", r.URL(), "URL() must return the value pushed by applyInit") // Empty URL in a subsequent applyInit must clear the stored value. r.applyInit(&wire.WorkerInit{WorkerID: "w-1", Concurrency: 2}) assert.Equal(t, "", r.URL(), "URL() must return empty after applyInit with empty URL") }