diff --git a/handle_performance_test.go b/handle_performance_test.go new file mode 100644 index 0000000..33a430c --- /dev/null +++ b/handle_performance_test.go @@ -0,0 +1,49 @@ +package queue + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "github.com/golang-queue/queue/job" +) + +// Isolate handle overhead with fresh retry state for each completed invocation. +// A 1ns retry delay emphasizes timer bookkeeping rather than real backoff waits. +func BenchmarkHandleRetry(b *testing.B) { + for _, retries := range []int64{0, 1, 3, 10} { + b.Run(fmt.Sprintf("retries=%d", retries), func(b *testing.B) { + q, err := NewQueue(WithWorker(NewRing()), WithLogger(NewEmptyLogger())) + if err != nil { + b.Fatal(err) + } + b.Cleanup(q.Release) + retryErr := errors.New("retry") + var attempts int64 + m := job.Message{ + Timeout: time.Minute, RetryDelay: time.Nanosecond, + Task: func(context.Context) error { + attempts++ + if attempts <= retries { + return retryErr + } + return nil + }, + } + b.ReportAllocs() + b.ResetTimer() + for range b.N { + attempts = 0 + m.RetryCount = retries + if err := q.handle(&m); err != nil { + b.Fatal(err) + } + if attempts != retries+1 { + b.Fatalf("got %d attempts, want %d", attempts, retries+1) + } + } + }) + } +} diff --git a/handle_result_test.go b/handle_result_test.go new file mode 100644 index 0000000..505bb41 --- /dev/null +++ b/handle_result_test.go @@ -0,0 +1,123 @@ +package queue + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/golang-queue/queue/job" + + "github.com/stretchr/testify/require" +) + +func TestResultProtocolShutdown(t *testing.T) { + for _, kind := range []string{"success", "error", "panic"} { + t.Run(kind, func(t *testing.T) { + q, err := NewQueue(WithWorker(NewRing()), WithLogger(NewEmptyLogger())) + require.NoError(t, err) + t.Cleanup(q.Release) + entered := make(chan struct{}) + type outcome struct { + err error + p any + } + finished := make(chan outcome, 1) + taskErr := errors.New("task failed") + m := job.NewTask(func(ctx context.Context) error { + close(entered) + <-ctx.Done() + if kind == "panic" { + panic("shutdown panic") + } + if kind == "error" { + return taskErr + } + return nil + }, job.AllowOption{Timeout: job.Time(5 * time.Second)}) + go func() { + var result outcome + defer func() { result.p = recover(); finished <- result }() + result.err = q.handle(&m) + }() + receiveWithin(t, entered) + q.Shutdown() + result := receiveWithin(t, finished) + if kind == "panic" { + require.Equal(t, "shutdown panic", result.p) + } else { + require.Nil(t, result.p) + if kind == "error" { + require.ErrorIs(t, result.err, taskErr) + } else { + require.NoError(t, result.err) + } + } + }) + } +} + +// TestMain's goleak check detects a handler blocked on a late result send. +func TestResultProtocolLateCompletion(t *testing.T) { + for _, kind := range []string{"success", "error", "panic"} { + t.Run(kind, func(t *testing.T) { + q, err := NewQueue(WithWorker(NewRing()), WithLogger(NewEmptyLogger())) + require.NoError(t, err) + t.Cleanup(q.Release) + release := make(chan struct{}) + defer close(release) + entered := make(chan struct{}) + m := job.NewTask(func(context.Context) error { + close(entered) + <-release + if kind == "panic" { + panic("late panic") + } + if kind == "error" { + return errors.New("late failure") + } + return nil + }, job.AllowOption{Timeout: job.Time(5 * time.Millisecond)}) + require.ErrorIs(t, q.handle(&m), context.DeadlineExceeded) + receiveWithin(t, entered) + }) + } +} + +func TestResultProtocolRetryCancellation(t *testing.T) { + q, err := NewQueue(WithWorker(NewRing()), WithLogger(NewEmptyLogger())) + require.NoError(t, err) + t.Cleanup(q.Release) + taskErr := errors.New("retry") + task := func(context.Context) error { return taskErr } + m := job.NewTask(task, job.AllowOption{ + RetryCount: job.Int64(10), + RetryDelay: job.Time(time.Hour), + Timeout: job.Time(10 * time.Millisecond), + }) + require.ErrorIs(t, q.handle(&m), context.DeadlineExceeded) +} + +func TestResultProtocolRepeatedRetries(t *testing.T) { + for _, succeed := range []bool{false, true} { + q, err := NewQueue(WithWorker(NewRing()), WithLogger(NewEmptyLogger())) + require.NoError(t, err) + t.Cleanup(q.Release) + var attempts int + taskErr := errors.New("retry") + m := job.NewTask(func(context.Context) error { + attempts++ + if succeed && attempts == 11 { + return nil + } + return taskErr + }, job.AllowOption{RetryCount: job.Int64(10), RetryDelay: job.Time(time.Nanosecond)}) + err = q.handle(&m) + if succeed { + require.NoError(t, err) + } else { + require.ErrorIs(t, err, taskErr) + } + require.Equal(t, 11, attempts) + } +} diff --git a/options.go b/options.go index c45161f..6132f55 100644 --- a/options.go +++ b/options.go @@ -123,7 +123,8 @@ func WithFn(fn func(context.Context, core.TaskMessage) error) Option { // WithAfterFn sets a callback function that will be executed after each job completes. // This callback runs regardless of whether the job succeeded or failed. -// It executes after metrics are updated but before the worker picks up the next task. +// It executes after metrics are updated without occupying a worker slot. +// Callbacks may overlap; Release waits for them to finish. // Useful for cleanup, logging, or triggering post-processing workflows. // // Example: diff --git a/queue.go b/queue.go index 6115337..1c2ae74 100644 --- a/queue.go +++ b/queue.go @@ -35,6 +35,7 @@ type ( ready chan struct{} // Channel to signal worker readiness notify chan struct{} // Channel to notify workers of new jobs worker core.Worker // The worker implementation that processes jobs + startOnce sync.Once // Ensures only one dispatcher requests tasks stopOnce sync.Once // Ensures shutdown is only performed once stopFlag atomic.Int32 // Atomic flag indicating if shutdown has started afterFn func() // Optional callback after each job execution @@ -71,7 +72,7 @@ func NewQueue(opts ...Option) (*Queue, error) { return q, nil } -// Start launches all worker goroutines and begins processing jobs. +// Start launches the dispatcher and workers once and begins processing jobs. // If workerCount is zero, Start is a no-op. func (q *Queue) Start() { q.Lock() @@ -80,8 +81,8 @@ func (q *Queue) Start() { if count == 0 { return } - q.routineGroup.Run(func() { - q.start() + q.startOnce.Do(func() { + q.routineGroup.Run(q.start) }) } @@ -198,7 +199,10 @@ func (q *Queue) work(task core.TaskMessage) { q.metric.IncFailureTask() } if q.afterFn != nil { - q.afterFn() + // Completion callbacks have never occupied a worker slot. Keep them + // independent so a callback can wait for another queued task, while + // still making Release wait for all callbacks to finish. + q.routineGroup.Run(q.afterFn) } }() @@ -221,10 +225,12 @@ func (q *Queue) run(task core.TaskMessage) error { // handle executes a job.Message, supporting retries, timeouts, and panic recovery. // Returns an error if the job fails or times out. func (q *Queue) handle(m *job.Message) error { - // done: receives the result of the job execution - // panicChan: receives any panic that occurs in the job goroutine - done := make(chan error, 1) - panicChan := make(chan any, 1) + // One buffered result lets the handler finish even after timeout returns. + type result struct { + err error + panicValue any + } + done := make(chan result, 1) startTime := time.Now() ctx, cancel := context.WithTimeout(context.Background(), m.Timeout) defer func() { @@ -233,10 +239,10 @@ func (q *Queue) handle(m *job.Message) error { // Run the job in a separate goroutine to support timeout and panic recovery. go func() { - // Defer block to catch panics and send to panicChan + // Forward panic through the same result channel. defer func() { if p := recover(); p != nil { - panicChan <- p + done <- result{panicValue: p} } }() @@ -279,12 +285,10 @@ func (q *Queue) handle(m *job.Message) error { } } - done <- err + done <- result{err: err} }() select { - case p := <-panicChan: - panic(p) case <-ctx.Done(): // Timeout reached return ctx.Err() case <-q.quit: // Queue is shutting down @@ -294,13 +298,17 @@ func (q *Queue) handle(m *job.Message) error { select { case <-time.After(leftTime): return context.DeadlineExceeded - case err := <-done: // Job finished - return err - case p := <-panicChan: - panic(p) + case r := <-done: // Job finished + if r.panicValue != nil { + panic(r.panicValue) + } + return r.err } - case err := <-done: // Job finished - return err + case r := <-done: // Job finished + if r.panicValue != nil { + panic(r.panicValue) + } + return r.err } } @@ -313,86 +321,99 @@ func (q *Queue) UpdateWorkerCount(num int64) { q.schedule() } -// schedule checks if more workers can be started based on the current busy count. -// If so, it signals readiness to start a new worker. +// schedule wakes the dispatcher after a completion or worker-count update. +// Notifications may coalesce; the dispatcher always rechecks the current limit. func (q *Queue) schedule() { - q.Lock() - defer q.Unlock() - if q.BusyWorkers() >= q.workerCount { - return - } - select { case q.ready <- struct{}{}: default: } } -/* -start launches the main worker loop, which manages job scheduling and execution. +// consume reuses a worker goroutine until it is retired or dispatch has stopped. +// Retiring a busy worker does not interrupt its current task. +func (q *Queue) consume(tasks <-chan core.TaskMessage, retire <-chan struct{}) { + for { + // Prefer retirement before accepting another task. + select { + case <-retire: + return + default: + } + select { + case <-retire: + return + case task, ok := <-tasks: + if !ok { + return + } + q.work(task) + } + } +} -- It uses a ticker to periodically retry job requests if the queue is empty. -- For each available worker slot, it requests a new task from the worker. -- If a task is available, it is sent to the tasks channel and processed by a new goroutine. -- The loop exits when the quit channel is closed. -*/ +// start serializes backend requests and dispatches tasks to reusable workers. +// Only the dispatcher owns the worker set; completed tasks and count updates wake +// it to recheck capacity. The handler goroutine in handle still enforces timeout. func (q *Queue) start() { - tasks := make(chan core.TaskMessage, 1) + tasks := make(chan core.TaskMessage) + defer close(tasks) ticker := time.NewTicker(q.retryInterval) defer ticker.Stop() + var workers []chan struct{} for { - // Ensure the number of busy workers does not exceed the configured worker count. - q.schedule() + q.Lock() + count := q.workerCount + q.Unlock() + + // A nonpositive dynamic count pauses dispatch without interrupting work. + count = max(count, 0) + for int64(len(workers)) > count { + last := len(workers) - 1 + close(workers[last]) + workers[last] = nil + workers = workers[:last] + } + for int64(len(workers)) < count { + retire := make(chan struct{}) + workers = append(workers, retire) + q.routineGroup.Run(func() { q.consume(tasks, retire) }) + } + if q.BusyWorkers() >= count { + select { + case <-q.ready: + continue + case <-q.quit: + return + } + } select { - case <-q.ready: // Wait for a worker slot to become available - case <-q.quit: // Shutdown signal received + case <-q.quit: return + default: } - // Request a task from the worker in a background goroutine. - q.routineGroup.Run(func() { - for { - t, err := q.worker.Request() - if t == nil || err != nil { - if err != nil { - select { - case <-q.quit: - if !errors.Is(err, ErrNoTaskInQueue) { - close(tasks) - return - } - case <-ticker.C: - case <-q.notify: - } - } - } - if t != nil { - tasks <- t - return - } - - select { - case <-q.quit: - if !errors.Is(err, ErrNoTaskInQueue) { - close(tasks) - return - } - default: - } + // Request was already serialized by the old dispatcher waiting for each + // result. Calling it here removes a goroutine and channel hop per task. + task, err := q.worker.Request() + if task == nil || err != nil { + select { + case <-q.quit: + return + case <-ticker.C: + case <-q.notify: + case <-q.ready: + } + if task == nil { + continue } - }) - - task, ok := <-tasks - if !ok { - return } - // Start processing the new task in a separate goroutine. + // Once dequeued, the task must be handed off even during shutdown. + // The workers remain alive until this dispatcher closes tasks. q.metric.IncBusyWorker() - q.routineGroup.Run(func() { - q.work(task) - }) + tasks <- task } } diff --git a/queue_performance_test.go b/queue_performance_test.go new file mode 100644 index 0000000..eed5d19 --- /dev/null +++ b/queue_performance_test.go @@ -0,0 +1,55 @@ +package queue + +import ( + "context" + "fmt" + "sync" + "testing" +) + +// BenchmarkQueueEndToEnd measures public submission through handler completion +// and final Release. Batches bound the backlog independently of b.N. +func BenchmarkQueueEndToEnd(b *testing.B) { + for _, workers := range []int64{1, 4, 12} { + b.Run(fmt.Sprintf("workers=%d", workers), func(b *testing.B) { + benchmarkQueueCompletion(b, workers, false) + }) + } +} + +// BenchmarkQueueAfterFnEndToEnd also waits for every configured callback. +func BenchmarkQueueAfterFnEndToEnd(b *testing.B) { + benchmarkQueueCompletion(b, 4, true) +} + +func benchmarkQueueCompletion(b *testing.B, workers int64, withCallback bool) { + var completed sync.WaitGroup + opts := []Option{WithLogger(NewEmptyLogger())} + task := func(context.Context) error { + completed.Done() + return nil + } + if withCallback { + opts = append(opts, WithAfterFn(completed.Done)) + task = func(context.Context) error { return nil } + } + q := NewPool(workers, opts...) + b.ReportAllocs() + b.ResetTimer() + for submitted := 0; submitted < b.N; { + n := min(1024, b.N-submitted) + completed.Add(n) + for range n { + if err := q.QueueTask(task); err != nil { + b.Fatal(err) + } + } + completed.Wait() + submitted += n + } + q.Release() + b.StopTimer() + if got := q.SuccessTasks(); got != uint64(b.N) { + b.Fatalf("completed %d tasks, want %d", got, b.N) + } +} diff --git a/scheduling_test.go b/scheduling_test.go new file mode 100644 index 0000000..76762a6 --- /dev/null +++ b/scheduling_test.go @@ -0,0 +1,260 @@ +package queue + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/golang-queue/queue/core" + "github.com/golang-queue/queue/job" + + "github.com/stretchr/testify/require" +) + +func receiveWithin[T any](t *testing.T, ch <-chan T) T { + t.Helper() + select { + case value := <-ch: + return value + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for queue progress") + var zero T + return zero + } +} + +func TestQueueResizePauseAndDrain(t *testing.T) { + const total = 16 + started := make(chan int, total) + completed := make(chan struct{}, total) + release := make(chan struct{}) + var releaseOnce sync.Once + q := NewPool(2, WithLogger(NewEmptyLogger()), WithAfterFn(func() { completed <- struct{}{} })) + t.Cleanup(func() { + releaseOnce.Do(func() { close(release) }) + q.UpdateWorkerCount(2) + q.Release() + }) + for id := range total { + require.NoError(t, q.QueueTask(func(context.Context) error { + started <- id + <-release + return nil + })) + } + seen := make(map[int]bool) + receiveStarted := func(n int) { + for range n { + id := receiveWithin(t, started) + require.False(t, seen[id], "task executed twice") + seen[id] = true + } + } + receiveStarted(2) + q.UpdateWorkerCount(4) + receiveStarted(2) + require.EqualValues(t, 4, q.BusyWorkers()) + q.UpdateWorkerCount(1) + for range 4 { + select { + case release <- struct{}{}: + case <-time.After(5 * time.Second): + t.Fatal("workers did not finish during scale-down") + } + } + for range 4 { + receiveWithin(t, completed) + } + receiveStarted(1) + require.EqualValues(t, 1, q.BusyWorkers()) + q.UpdateWorkerCount(0) + release <- struct{}{} + receiveWithin(t, completed) + select { + case <-started: + t.Fatal("task started while paused") + case <-time.After(20 * time.Millisecond): + } + q.UpdateWorkerCount(2) + receiveStarted(2) + require.EqualValues(t, 2, q.BusyWorkers()) + releaseOnce.Do(func() { close(release) }) + q.Release() + receiveStarted(total - len(seen)) + require.Len(t, seen, total) + require.EqualValues(t, total, q.SuccessTasks()) + require.EqualValues(t, 0, q.BusyWorkers()) +} + +func TestReusableWorkerOutcomes(t *testing.T) { + completed := make(chan struct{}, 5) + lateRelease := make(chan struct{}) + lateDone := make(chan struct{}) + q := NewPool(1, WithLogger(NewEmptyLogger()), WithAfterFn(func() { completed <- struct{}{} })) + t.Cleanup(func() { + close(lateRelease) + q.Release() + }) + wantErr := errors.New("task failed") + require.NoError(t, q.QueueTask(func(context.Context) error { return wantErr })) + require.NoError(t, q.QueueTask(func(context.Context) error { panic("task panic") })) + var attempts atomic.Int32 + require.NoError(t, q.QueueTask(func(context.Context) error { + if attempts.Add(1) == 1 { + return wantErr + } + return nil + }, job.AllowOption{RetryCount: job.Int64(1), RetryDelay: job.Time(time.Millisecond)})) + require.NoError(t, q.QueueTask(func(context.Context) error { + defer close(lateDone) + <-lateRelease // Deliberately ignore cancellation to check timeout isolation. + return nil + }, job.AllowOption{Timeout: job.Time(10 * time.Millisecond)})) + require.NoError(t, q.QueueTask(func(context.Context) error { return nil })) + for range 5 { + receiveWithin(t, completed) + } + require.EqualValues(t, 2, attempts.Load()) + require.EqualValues(t, 2, q.SuccessTasks()) + require.EqualValues(t, 3, q.FailureTasks()) + require.EqualValues(t, 5, q.CompletedTasks()) + select { + case <-lateDone: + t.Fatal("timeout test handler finished before being released") + default: + } + // Release still completes even though the timed-out handler has not returned. + q.Release() +} + +func TestAfterFnCanWaitForNextTask(t *testing.T) { + var q *Queue + var callbacks atomic.Int32 + callbackDone := make(chan error, 1) + q = NewPool(1, WithLogger(NewEmptyLogger()), WithAfterFn(func() { + if callbacks.Add(1) != 1 { + return + } + nextTask := make(chan struct{}) + if err := q.QueueTask(func(context.Context) error { + close(nextTask) + return nil + }); err != nil { + callbackDone <- err + return + } + select { + case <-nextTask: + callbackDone <- nil + case <-time.After(time.Second): + callbackDone <- errors.New("afterFn blocked the next task") + } + })) + t.Cleanup(q.Release) + require.NoError(t, q.QueueTask(func(context.Context) error { return nil })) + require.NoError(t, receiveWithin(t, callbackDone)) + q.Release() + require.EqualValues(t, 2, q.SuccessTasks()) + require.EqualValues(t, 2, callbacks.Load()) +} + +func TestReleaseWaitsForAfterFn(t *testing.T) { + entered := make(chan struct{}) + unblock := make(chan struct{}) + var once sync.Once + q := NewPool(1, WithLogger(NewEmptyLogger()), WithAfterFn(func() { + close(entered) + <-unblock + })) + t.Cleanup(func() { + once.Do(func() { close(unblock) }) + q.Release() + }) + require.NoError(t, q.QueueTask(func(context.Context) error { return nil })) + receiveWithin(t, entered) + released := make(chan struct{}) + go func() { + q.Release() + close(released) + }() + select { + case <-released: + t.Fatal("Release returned before afterFn finished") + case <-time.After(20 * time.Millisecond): + } + once.Do(func() { close(unblock) }) + receiveWithin(t, released) +} + +// blockingWorker models backends whose Request blocks until work or shutdown. +// Counting active calls detects accidental parallel Request calls. +type blockingWorker struct { + requests chan core.TaskMessage + closed chan struct{} + entered chan struct{} + once sync.Once + active atomic.Int32 + parallel atomic.Bool +} + +func (w *blockingWorker) Queue(m core.TaskMessage) error { + w.requests <- m + return nil +} + +func (w *blockingWorker) Request() (core.TaskMessage, error) { + if w.active.Add(1) != 1 { + w.parallel.Store(true) + } + defer w.active.Add(-1) + select { + case w.entered <- struct{}{}: + default: + } + select { + case m := <-w.requests: + return m, nil + case <-w.closed: + return nil, ErrQueueHasBeenClosed + } +} + +func (w *blockingWorker) Run(context.Context, core.TaskMessage) error { return nil } + +func (w *blockingWorker) Shutdown() error { + w.once.Do(func() { close(w.closed) }) + return nil +} + +func TestBlockingBackendAndRepeatedStart(t *testing.T) { + w := &blockingWorker{ + requests: make(chan core.TaskMessage, 64), + closed: make(chan struct{}), + entered: make(chan struct{}, 128), + } + completed := make(chan struct{}, 64) + q, err := NewQueue(WithWorker(w), WithWorkerCount(4), WithLogger(NewEmptyLogger()), + WithAfterFn(func() { completed <- struct{}{} })) + require.NoError(t, err) + t.Cleanup(q.Release) + q.Start() + receiveWithin(t, w.entered) + q.Start() + for range 64 { + require.NoError(t, q.QueueTask(func(context.Context) error { return nil })) + } + for range 64 { + receiveWithin(t, completed) + } + released := make(chan struct{}) + go func() { + q.Release() + close(released) + }() + receiveWithin(t, released) + require.False(t, w.parallel.Load(), "backend Request must remain serialized") + require.EqualValues(t, 64, q.SuccessTasks()) +}