counters/internal/counting/queue_test.go

package counting

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type recorder struct {
	mu     sync.Mutex
	jobs   []Job
	waits  []time.Duration
	depths []int
	panics int
}

func (r *recorder) handle(ctx context.Context, job Job) {
	r.mu.Lock()
	r.jobs = append(r.jobs, job)
	r.mu.Unlock()
}

func (r *recorder) observer() Observer {
	return Observer{
		Depth: func(n int) { r.mu.Lock(); r.depths = append(r.depths, n); r.mu.Unlock() },
		Wait:  func(d time.Duration) { r.mu.Lock(); r.waits = append(r.waits, d); r.mu.Unlock() },
		Panic: func(job Job, p any, stack []byte) { r.mu.Lock(); r.panics++; r.mu.Unlock() },
	}
}

func (r *recorder) count() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return len(r.jobs)
}

func settle(t *testing.T, q *Queue) {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for q.Pending() > 0 {
		if time.Now().After(deadline) {
			t.Fatalf("queue did not settle: pending=%d", q.Pending())
		}
		time.Sleep(time.Millisecond)
	}
}

func TestQueue_ProcessesJobs(t *testing.T) {
	r := &recorder{}
	q := NewQueue(QueueConfig{Size: 10, Workers: 2, Timeout: time.Second}, r.handle, r.observer())
	defer q.Drain(context.Background())
	for i := 0; i < 5; i++ {
		if !q.Enqueue(Job{Route: RouteIn, TraceID: "t", Recipient: "y@example.com", Received: time.Now()}) {
			t.Fatalf("enqueue %d refused", i)
		}
	}
	settle(t, q)
	if r.count() != 5 {
		t.Errorf("handled %d jobs, want 5", r.count())
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.waits) != 5 || len(r.depths) == 0 {
		t.Errorf("observer: waits=%d depths=%d", len(r.waits), len(r.depths))
	}
	for _, w := range r.waits {
		if w < 0 || w > time.Second {
			t.Errorf("queue wait %v out of range", w)
		}
	}
}

// Workers run the handler with their own bounded context, never a request's.
func TestQueue_HandlerContextHasTimeout(t *testing.T) {
	var got atomic.Value
	q := NewQueue(QueueConfig{Size: 1, Workers: 1, Timeout: 250 * time.Millisecond}, func(ctx context.Context, job Job) {
		dl, ok := ctx.Deadline()
		got.Store([2]any{ok, time.Until(dl)})
	}, Observer{})
	defer q.Drain(context.Background())
	q.Enqueue(Job{})
	settle(t, q)
	v := got.Load().([2]any)
	if !v[0].(bool) || v[1].(time.Duration) > 250*time.Millisecond || v[1].(time.Duration) < 100*time.Millisecond {
		t.Errorf("deadline = %v", v)
	}
}

// A full queue refuses the job instead of blocking the HTTP handler.
func TestQueue_FullRefuses(t *testing.T) {
	release := make(chan struct{})
	q := NewQueue(QueueConfig{Size: 2, Workers: 1, Timeout: time.Second}, func(ctx context.Context, job Job) { <-release }, Observer{})
	// One job occupies the worker ...
	if !q.Enqueue(Job{}) {
		t.Fatal("first enqueue refused")
	}
	deadline := time.Now().Add(time.Second)
	for q.Depth() != 0 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	// ... two fill the channel.
	for i := 0; i < 2; i++ {
		if !q.Enqueue(Job{}) {
			t.Fatalf("enqueue %d should be accepted", i)
		}
	}
	if q.Enqueue(Job{}) {
		t.Error("enqueue on a full queue should be refused")
	}
	if q.Depth() != 2 || q.Pending() != 3 {
		t.Errorf("depth=%d pending=%d, want 2/3", q.Depth(), q.Pending())
	}
	close(release)
	if err := q.Drain(context.Background()); err != nil {
		t.Fatal(err)
	}
	if q.Enqueue(Job{}) {
		t.Error("enqueue after Drain must be refused, not panic")
	}
}

// Drain waits for queued and in-flight jobs; a context that expires first
// reports what was left behind.
func TestQueue_Drain(t *testing.T) {
	r := &recorder{}
	slow := func(ctx context.Context, job Job) { time.Sleep(50 * time.Millisecond); r.handle(ctx, job) }
	q := NewQueue(QueueConfig{Size: 10, Workers: 1, Timeout: time.Second}, slow, Observer{})
	for i := 0; i < 4; i++ {
		q.Enqueue(Job{})
	}
	start := time.Now()
	if err := q.Drain(context.Background()); err != nil {
		t.Fatal(err)
	}
	if r.count() != 4 || time.Since(start) < 150*time.Millisecond {
		t.Errorf("Drain returned after %v with %d jobs handled, want all 4", time.Since(start), r.count())
	}

	block := make(chan struct{})
	q2 := NewQueue(QueueConfig{Size: 10, Workers: 1, Timeout: time.Second}, func(ctx context.Context, job Job) { <-block }, Observer{})
	q2.Enqueue(Job{})
	q2.Enqueue(Job{})
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	err := q2.Drain(ctx)
	if err == nil || !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("Drain with a stuck worker should report the deadline: %v", err)
	}
	close(block)
}

// A panicking handler is recovered and reported; the worker keeps serving.
func TestQueue_RecoversPanic(t *testing.T) {
	r := &recorder{}
	q := NewQueue(QueueConfig{Size: 10, Workers: 1, Timeout: time.Second}, func(ctx context.Context, job Job) {
		if job.TraceID == "boom" {
			panic("boom")
		}
		r.handle(ctx, job)
	}, r.observer())
	defer q.Drain(context.Background())
	q.Enqueue(Job{TraceID: "boom"})
	q.Enqueue(Job{TraceID: "ok"})
	settle(t, q)
	if r.count() != 1 || r.panics != 1 {
		t.Errorf("handled=%d panics=%d, want 1/1", r.count(), r.panics)
	}
}

func TestQueue_Defaults(t *testing.T) {
	q := NewQueue(QueueConfig{}, func(context.Context, Job) {}, Observer{})
	defer q.Drain(context.Background())
	if q.cfg.Size != 10000 || q.cfg.Workers != 8 || q.cfg.Timeout != 4*time.Second {
		t.Errorf("defaults = %+v", q.cfg)
	}
}