worker/internal/distworker/runner_test.go

package distworker

import (
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Jeffail/tunny"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"rocketgit.ru/rsmon/worker/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")
}