From 8004ba6420f0b81e163ba796765d24fbb2b6937d Mon Sep 17 00:00:00 2001 From: Chris Miles Date: Sun, 16 Aug 2026 23:40:50 +0000 Subject: [PATCH 1/4] feat: add worker lifecycle controls --- fake_queue.go | 8 ++ queue.go | 242 ++++++++++++++++++++++++++++++++++++- queue_local.go | 128 +++++++++++++++++--- queue_local_test.go | 87 +++++++++++++ queue_runtime_unit_test.go | 48 ++++++++ runtime.go | 12 ++ 6 files changed, 507 insertions(+), 18 deletions(-) diff --git a/fake_queue.go b/fake_queue.go index 70afd17..a02d2d4 100644 --- a/fake_queue.go +++ b/fake_queue.go @@ -157,6 +157,14 @@ func (f *FakeQueue) Register(string, Handler) {} // _ = err func (f *FakeQueue) StartWorkers(context.Context) error { return nil } +// PauseWorkers stops fake worker intake. +// @group Testing +func (f *FakeQueue) PauseWorkers(context.Context) error { return nil } + +// ResumeWorkers restarts fake worker intake. +// @group Testing +func (f *FakeQueue) ResumeWorkers(context.Context) error { return nil } + // Workers preserves fluent lifecycle compatibility without creating workers. // @group Testing // diff --git a/queue.go b/queue.go index 620fa7c..e6ed5ee 100644 --- a/queue.go +++ b/queue.go @@ -33,6 +33,14 @@ type queueRuntime interface { // @group Driver Integration StartWorkers(ctx context.Context) error + // PauseWorkers stops new worker intake after active handlers finish. + // @group Driver Integration + PauseWorkers(ctx context.Context) error + + // ResumeWorkers restarts worker intake after a pause. + // @group Driver Integration + ResumeWorkers(ctx context.Context) error + // Workers sets desired worker concurrency before StartWorkers. // @group Driver Integration Workers(count int) queueRuntime @@ -90,6 +98,12 @@ type runtimeQueueBackend interface { DrainWorkers(ctx context.Context) error } +// workerLifecycleBackend supports pausing consumers without closing producer resources. +type workerLifecycleBackend interface { + PauseWorkers(ctx context.Context) error + ResumeWorkers(ctx context.Context) error +} + func newSyncQueue() queueBackend { return newLocalQueueWithConfig(DriverSync, WorkerpoolConfig{}) } @@ -217,9 +231,11 @@ type nativeQueueRuntimeState struct { handlerSlots map[string]*runtimeHandlerSlot runtimeRegistrations map[string]struct{} started bool + paused bool draining bool closed bool start *runtimeStartAttempt + pause *runtimePauseAttempt shutdown *runtimeShutdownAttempt operations runtimeOperationState continuation *busruntime.ContinuationScope @@ -240,9 +256,11 @@ type externalQueueRuntimeState struct { worker runtimeWorkerBackend workerRegistrations map[string]struct{} started bool + paused bool draining bool closed bool start *runtimeStartAttempt + pause *runtimePauseAttempt shutdown *runtimeShutdownAttempt operations runtimeOperationState continuation *busruntime.ContinuationScope @@ -284,6 +302,12 @@ type runtimeStartAttempt struct { err error } +// runtimePauseAttempt lets concurrent lifecycle callers observe one consumer pause. +type runtimePauseAttempt struct { + done chan struct{} + err error +} + type runtimeHandlerSlot struct { mu sync.RWMutex handler Handler @@ -756,7 +780,7 @@ func (q *nativeQueueRuntime) StartWorkers(ctx context.Context) error { q.mu.Unlock() return ErrQueuerShuttingDown } - if q.started { + if q.started || q.paused { q.mu.Unlock() return nil } @@ -796,7 +820,7 @@ func (q *externalQueueRuntime) StartWorkers(ctx context.Context) error { q.mu.Unlock() return ErrQueuerShuttingDown } - if q.started { + if q.started || q.paused { q.mu.Unlock() return nil } @@ -853,6 +877,194 @@ func (q *externalQueueRuntime) StartWorkers(ctx context.Context) error { return err } +// PauseWorkers stops native worker intake while retaining dispatch resources. +func (q *nativeQueueRuntime) PauseWorkers(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } + q.mu.Lock() + if q.closed || q.draining { + q.mu.Unlock() + return ErrQueuerShuttingDown + } + if q.paused { + q.mu.Unlock() + return nil + } + if q.start != nil { + start := q.start + q.mu.Unlock() + if err := waitForRuntimeStartCompletion(ctx, start); err != nil { + return err + } + return q.PauseWorkers(ctx) + } + if q.pause != nil { + pause := q.pause + q.mu.Unlock() + return waitForRuntimePause(ctx, pause) + } + attempt := &runtimePauseAttempt{done: make(chan struct{})} + q.pause = attempt + started := q.started + q.mu.Unlock() + + var err error + if started { + if lifecycle, ok := q.runtime.(workerLifecycleBackend); ok { + err = lifecycle.PauseWorkers(ctx) + } + } + q.mu.Lock() + if err == nil { + q.paused = true + } + attempt.err = err + q.pause = nil + close(attempt.done) + q.mu.Unlock() + return err +} + +// ResumeWorkers restarts native worker intake after a pause. +func (q *nativeQueueRuntime) ResumeWorkers(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } + q.mu.Lock() + if q.closed || q.draining { + q.mu.Unlock() + return ErrQueuerShuttingDown + } + if q.pause != nil { + pause := q.pause + q.mu.Unlock() + if err := waitForRuntimePause(ctx, pause); err != nil { + return err + } + return q.ResumeWorkers(ctx) + } + if !q.paused { + started := q.started + q.mu.Unlock() + if started { + return nil + } + return q.StartWorkers(ctx) + } + started := q.started + q.mu.Unlock() + + if started { + if lifecycle, ok := q.runtime.(workerLifecycleBackend); ok { + if err := lifecycle.ResumeWorkers(ctx); err != nil { + return err + } + } + q.mu.Lock() + q.paused = false + q.mu.Unlock() + return nil + } + q.mu.Lock() + q.paused = false + q.mu.Unlock() + if err := q.StartWorkers(ctx); err != nil { + q.mu.Lock() + q.paused = true + q.mu.Unlock() + return err + } + return nil +} + +// PauseWorkers stops the external consumer while retaining producer resources. +func (q *externalQueueRuntime) PauseWorkers(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } + q.mu.Lock() + if q.closed || q.draining { + q.mu.Unlock() + return ErrQueuerShuttingDown + } + if q.paused { + q.mu.Unlock() + return nil + } + if q.start != nil { + start := q.start + q.mu.Unlock() + if err := waitForRuntimeStartCompletion(ctx, start); err != nil { + return err + } + return q.PauseWorkers(ctx) + } + if q.pause != nil { + pause := q.pause + q.mu.Unlock() + return waitForRuntimePause(ctx, pause) + } + attempt := &runtimePauseAttempt{done: make(chan struct{})} + q.pause = attempt + w := q.worker + q.mu.Unlock() + + var err error + if w != nil { + err = w.Shutdown(ctx) + } + q.mu.Lock() + if err == nil { + q.worker = nil + q.workerRegistrations = nil + q.started = false + q.paused = true + } + attempt.err = err + q.pause = nil + close(attempt.done) + q.mu.Unlock() + return err +} + +// ResumeWorkers creates a fresh external consumer after a pause. +func (q *externalQueueRuntime) ResumeWorkers(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } + q.mu.Lock() + if q.closed || q.draining { + q.mu.Unlock() + return ErrQueuerShuttingDown + } + if q.pause != nil { + pause := q.pause + q.mu.Unlock() + if err := waitForRuntimePause(ctx, pause); err != nil { + return err + } + return q.ResumeWorkers(ctx) + } + if !q.paused { + started := q.started + q.mu.Unlock() + if started { + return nil + } + return q.StartWorkers(ctx) + } + q.paused = false + q.mu.Unlock() + if err := q.StartWorkers(ctx); err != nil { + q.mu.Lock() + q.paused = true + q.mu.Unlock() + return err + } + return nil +} + func (q *nativeQueueRuntime) Workers(count int) queueRuntime { q.mu.Lock() defer q.mu.Unlock() @@ -880,6 +1092,14 @@ func (q *nativeQueueRuntime) Shutdown(ctx context.Context) error { ctx = context.Background() } q.mu.Lock() + if q.pause != nil { + pause := q.pause + q.mu.Unlock() + if err := waitForRuntimePause(ctx, pause); err != nil { + return err + } + return q.Shutdown(ctx) + } if q.start != nil { q.draining = true attempt := q.start @@ -938,6 +1158,14 @@ func (q *externalQueueRuntime) Shutdown(ctx context.Context) error { ctx = context.Background() } q.mu.Lock() + if q.pause != nil { + pause := q.pause + q.mu.Unlock() + if err := waitForRuntimePause(ctx, pause); err != nil { + return err + } + return q.Shutdown(ctx) + } if q.start != nil { q.draining = true attempt := q.start @@ -1021,6 +1249,16 @@ func waitForRuntimeShutdown(ctx context.Context, attempt *runtimeShutdownAttempt } } +// waitForRuntimePause lets concurrent lifecycle callers share one pause attempt while honoring their own deadline. +func waitForRuntimePause(ctx context.Context, attempt *runtimePauseAttempt) error { + select { + case <-attempt.done: + return attempt.err + case <-ctx.Done(): + return ctx.Err() + } +} + // waitForRuntimeStart lets concurrent callers share one startup attempt while honoring their own deadline. func waitForRuntimeStart(ctx context.Context, attempt *runtimeStartAttempt) error { select { diff --git a/queue_local.go b/queue_local.go index 251081f..c792787 100644 --- a/queue_local.go +++ b/queue_local.go @@ -13,22 +13,28 @@ import ( // localQueue is an in-memory queue implementation supporting sync and workerpool drivers. type localQueue struct { - driver Driver - cfg WorkerpoolConfig - mu sync.RWMutex - metricsMu sync.RWMutex - queueMu sync.RWMutex - handlers map[string]Handler - unique uniqueness.MemoryStore - metrics map[string]*localQueueMetrics - pausedQueues map[string]bool - workQueue chan queuedJob - workPending int - workIdle chan struct{} - continuation *busruntime.ContinuationScope - resizeBuffer bool - shutdownOnce sync.Once - workerWG sync.WaitGroup + driver Driver + cfg WorkerpoolConfig + mu sync.RWMutex + metricsMu sync.RWMutex + queueMu sync.RWMutex + handlers map[string]Handler + unique uniqueness.MemoryStore + metrics map[string]*localQueueMetrics + pausedQueues map[string]bool + workQueue chan queuedJob + workPending int + workIdle chan struct{} + continuation *busruntime.ContinuationScope + resizeBuffer bool + shutdownOnce sync.Once + workerWG sync.WaitGroup + workerStateMu sync.Mutex + workerPaused bool + workerStopping bool + workerActive int + workerResume chan struct{} + workerDrained chan struct{} syncWorkMu sync.Mutex syncWorkPending int @@ -138,6 +144,47 @@ func (d *localQueue) StartWorkers(_ context.Context) error { return nil } +// PauseWorkers prevents the in-memory pool from starting queued jobs and waits for active handlers. +func (d *localQueue) PauseWorkers(ctx context.Context) error { + if d.driver != DriverWorkerpool { + return nil + } + if ctx == nil { + ctx = context.Background() + } + d.workerStateMu.Lock() + if !d.workerPaused { + d.workerPaused = true + d.workerResume = make(chan struct{}) + } + if d.workerActive == 0 { + d.workerStateMu.Unlock() + return nil + } + if d.workerDrained == nil { + d.workerDrained = make(chan struct{}) + } + drained := d.workerDrained + d.workerStateMu.Unlock() + select { + case <-drained: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// ResumeWorkers allows the in-memory pool to start queued jobs again. +func (d *localQueue) ResumeWorkers(_ context.Context) error { + if d.driver != DriverWorkerpool { + return nil + } + d.workerStateMu.Lock() + d.resumeWorkersLocked() + d.workerStateMu.Unlock() + return nil +} + // Shutdown drains running work and releases resources. // @group Queue // @@ -162,6 +209,10 @@ func (d *localQueue) DrainWorkers(ctx context.Context) error { d.shutdownOnce.Do(func() { d.shuttingDown.Store(true) + d.workerStateMu.Lock() + d.workerStopping = d.workerPaused + d.resumeWorkersLocked() + d.workerStateMu.Unlock() }) if d.driver == DriverSync { @@ -526,8 +577,13 @@ func (d *localQueue) worker(workQueue <-chan queuedJob) { defer d.workerWG.Done() jobTimeout := d.cfg.DefaultJobTimeout for job := range workQueue { + if !d.beginWorkerExecution() { + d.finishQueuedWork() + continue + } func() { defer d.finishQueuedWork() + defer d.endWorkerExecution() if job.ready != nil { <-job.ready } @@ -569,6 +625,46 @@ func (d *localQueue) worker(workQueue <-chan queuedJob) { } } +// beginWorkerExecution waits outside the backend claim path and rejects pending work during a paused shutdown. +func (d *localQueue) beginWorkerExecution() bool { + for { + d.workerStateMu.Lock() + if d.workerStopping { + d.workerStateMu.Unlock() + return false + } + if !d.workerPaused { + d.workerActive++ + d.workerStateMu.Unlock() + return true + } + resume := d.workerResume + d.workerStateMu.Unlock() + <-resume + } +} + +// endWorkerExecution releases pause waiters after an admitted handler completes. +func (d *localQueue) endWorkerExecution() { + d.workerStateMu.Lock() + d.workerActive-- + if d.workerActive == 0 && d.workerDrained != nil { + close(d.workerDrained) + d.workerDrained = nil + } + d.workerStateMu.Unlock() +} + +// resumeWorkersLocked releases workers reserved behind the lifecycle gate. +func (d *localQueue) resumeWorkersLocked() { + if !d.workerPaused { + return + } + d.workerPaused = false + close(d.workerResume) + d.workerResume = nil +} + // validateEnqueue rejects work before uniqueness is claimed or an acceptance fact is committed. func (d *localQueue) validateEnqueue(job Job, queueName string) error { if d.isPaused(queueName) { diff --git a/queue_local_test.go b/queue_local_test.go index ef7f3cb..39a99c4 100644 --- a/queue_local_test.go +++ b/queue_local_test.go @@ -17,6 +17,93 @@ func TestLocalQueue_Driver(t *testing.T) { } } +func TestLocalQueueWorkerLifecyclePausesIntakeAndDrainsActiveHandlers(t *testing.T) { + d := newLocalQueueWithConfig(DriverWorkerpool, WorkerpoolConfig{Workers: 1, QueueCapacity: 4}) + started := make(chan struct{}) + release := make(chan struct{}) + finished := make(chan int, 2) + var calls atomic.Int64 + d.Register("job:lifecycle", func(context.Context, Job) error { + call := int(calls.Add(1)) + if call == 1 { + close(started) + <-release + } + finished <- call + return nil + }) + if err := d.StartWorkers(context.Background()); err != nil { + t.Fatalf("start workers: %v", err) + } + if err := d.Dispatch(context.Background(), NewJob("job:lifecycle")); err != nil { + t.Fatalf("dispatch active job: %v", err) + } + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("active handler did not start") + } + pauseDone := make(chan error, 1) + go func() { pauseDone <- d.PauseWorkers(context.Background()) }() + select { + case err := <-pauseDone: + t.Fatalf("pause returned before active handler drained: %v", err) + case <-time.After(20 * time.Millisecond): + } + close(release) + if err := <-pauseDone; err != nil { + t.Fatalf("pause workers: %v", err) + } + if err := d.Dispatch(context.Background(), NewJob("job:lifecycle")); err != nil { + t.Fatalf("dispatch while paused: %v", err) + } + if got := <-finished; got != 1 { + t.Fatalf("first finished job = %d, want 1", got) + } + select { + case got := <-finished: + t.Fatalf("paused worker executed job %d", got) + case <-time.After(30 * time.Millisecond): + } + if err := d.ResumeWorkers(context.Background()); err != nil { + t.Fatalf("resume workers: %v", err) + } + select { + case got := <-finished: + if got != 2 { + t.Fatalf("resumed worker executed job %d, want 2", got) + } + case <-time.After(time.Second): + t.Fatal("queued job did not execute after resume") + } + if err := d.Shutdown(context.Background()); err != nil { + t.Fatalf("shutdown: %v", err) + } +} + +func TestLocalQueuePausedShutdownDoesNotStartPendingHandlers(t *testing.T) { + d := newLocalQueueWithConfig(DriverWorkerpool, WorkerpoolConfig{Workers: 1, QueueCapacity: 2}) + handled := make(chan struct{}, 1) + d.Register("job:pending", func(context.Context, Job) error { + handled <- struct{}{} + return nil + }) + if err := d.PauseWorkers(context.Background()); err != nil { + t.Fatalf("pause workers: %v", err) + } + if err := d.Dispatch(context.Background(), NewJob("job:pending")); err != nil { + t.Fatalf("dispatch while paused: %v", err) + } + if err := d.Shutdown(context.Background()); err != nil { + t.Fatalf("shutdown paused workers: %v", err) + } + select { + case <-handled: + t.Fatal("paused shutdown started a pending handler") + default: + } +} + func TestLocalQueue_DispatchRunsRegisteredHandler(t *testing.T) { d := newLocalQueue(DriverSync) var calls atomic.Int64 diff --git a/queue_runtime_unit_test.go b/queue_runtime_unit_test.go index 59e87e1..e724c1e 100644 --- a/queue_runtime_unit_test.go +++ b/queue_runtime_unit_test.go @@ -1958,3 +1958,51 @@ func TestExternalQueueRuntimeStartWorkersErrorBranches(t *testing.T) { } }) } + +func TestExternalQueueRuntimeWorkerLifecyclePreservesProducer(t *testing.T) { + inner := &queueBackendRecorder{} + workers := []*runtimeBackendStub{{}, {}} + nextWorker := 0 + q := &externalQueueRuntime{ + common: &queueCommon{inner: inner, cfg: Config{DefaultQueue: "default"}, driver: DriverNATS}, + newWorker: func(int) (driverWorkerBackend, error) { + worker := workers[nextWorker] + nextWorker++ + return worker, nil + }, + externalQueueRuntimeState: &externalQueueRuntimeState{ + registered: map[string]Handler{"job:external": func(context.Context, Job) error { return nil }}, + }, + } + + if err := q.StartWorkers(context.Background()); err != nil { + t.Fatalf("start workers: %v", err) + } + if err := q.PauseWorkers(context.Background()); err != nil { + t.Fatalf("pause workers: %v", err) + } + if workers[0].stopCalls != 1 { + t.Fatalf("first worker shutdowns = %d, want 1", workers[0].stopCalls) + } + if inner.shutdowns != 0 { + t.Fatalf("producer shutdowns while paused = %d, want 0", inner.shutdowns) + } + if err := q.Dispatch(NewJob("job:external")); err != nil { + t.Fatalf("dispatch while paused: %v", err) + } + if err := q.ResumeWorkers(context.Background()); err != nil { + t.Fatalf("resume workers: %v", err) + } + if workers[1].startCalls != 1 { + t.Fatalf("replacement worker starts = %d, want 1", workers[1].startCalls) + } + if _, ok := workers[1].registered["job:external"]; !ok { + t.Fatal("replacement worker did not restore registrations") + } + if err := q.Shutdown(context.Background()); err != nil { + t.Fatalf("shutdown: %v", err) + } + if inner.shutdowns != 1 { + t.Fatalf("producer shutdowns = %d, want 1", inner.shutdowns) + } +} diff --git a/runtime.go b/runtime.go index 9136b83..02329e2 100644 --- a/runtime.go +++ b/runtime.go @@ -547,6 +547,18 @@ func (r *Queue) StartWorkers(ctx context.Context) error { return r.b.StartWorkers(ctx) } +// PauseWorkers stops new worker intake after active handlers finish while keeping dispatch available. +// @group Queue +func (r *Queue) PauseWorkers(ctx context.Context) error { + return r.q.PauseWorkers(ctx) +} + +// ResumeWorkers restarts worker intake after PauseWorkers. +// @group Queue +func (r *Queue) ResumeWorkers(ctx context.Context) error { + return r.q.ResumeWorkers(ctx) +} + // Run starts worker processing, blocks until ctx is canceled, then gracefully shuts down. // @group Queue // From 58701d5f5fbc2a8f088555108c78f0eb448d2569 Mon Sep 17 00:00:00 2001 From: Chris Miles Date: Sun, 16 Aug 2026 23:48:32 +0000 Subject: [PATCH 2/4] fix: register paused native workers --- queue.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/queue.go b/queue.go index e6ed5ee..214e636 100644 --- a/queue.go +++ b/queue.go @@ -907,6 +907,13 @@ func (q *nativeQueueRuntime) PauseWorkers(ctx context.Context) error { attempt := &runtimePauseAttempt{done: make(chan struct{})} q.pause = attempt started := q.started + if !started { + for jobType, handler := range q.registered { + var slot *runtimeHandlerSlot + q.handlerSlots, slot = updateRuntimeHandlerSlot(q.handlerSlots, jobType, handler) + q.runtimeRegistrations = installRuntimeHandler(q.runtime, q.common, q.runtimeRegistrations, jobType, handler, slot) + } + } q.mu.Unlock() var err error From e6366af9b25ffeddafd44a129d2ee52ca6c2c198 Mon Sep 17 00:00:00 2001 From: Chris Miles Date: Sun, 16 Aug 2026 23:50:10 +0000 Subject: [PATCH 3/4] fix: gate native intake before startup --- queue.go | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/queue.go b/queue.go index 214e636..79d50f7 100644 --- a/queue.go +++ b/queue.go @@ -917,10 +917,8 @@ func (q *nativeQueueRuntime) PauseWorkers(ctx context.Context) error { q.mu.Unlock() var err error - if started { - if lifecycle, ok := q.runtime.(workerLifecycleBackend); ok { - err = lifecycle.PauseWorkers(ctx) - } + if lifecycle, ok := q.runtime.(workerLifecycleBackend); ok { + err = lifecycle.PauseWorkers(ctx) } q.mu.Lock() if err == nil { From 3b99da924c152c037089ad8087ffd5622e6ebc72 Mon Sep 17 00:00:00 2001 From: Chris Miles Date: Mon, 17 Aug 2026 00:07:47 +0000 Subject: [PATCH 4/4] fix: harden worker lifecycle transitions --- driver/natsqueue/worker_nats_impl.go | 58 +++++++++-- driver/natsqueue/worker_nats_impl_test.go | 10 +- driver/sqlitequeue/worker_lifecycle_test.go | 105 ++++++++++++++++++++ driver/sqlqueuecore/queue_database_impl.go | 54 +++++++--- driver/sqlqueuecore/wrapper.go | 12 ++- driver_runtime.go | 2 +- internal/driverbridge/bridge.go | 18 ++++ queue.go | 52 ++++++++-- queue_local.go | 59 +++++------ queue_local_test.go | 61 +++++++++++- queue_runtime_unit_test.go | 62 ++++++++++++ 11 files changed, 417 insertions(+), 76 deletions(-) create mode 100644 driver/sqlitequeue/worker_lifecycle_test.go diff --git a/driver/natsqueue/worker_nats_impl.go b/driver/natsqueue/worker_nats_impl.go index fa10b26..208e2ce 100644 --- a/driver/natsqueue/worker_nats_impl.go +++ b/driver/natsqueue/worker_nats_impl.go @@ -31,9 +31,11 @@ type natsWorker struct { sub natsWorkerSubscription sem chan struct{} - running sync.WaitGroup - delayed sync.WaitGroup - observer queue.Observer + running sync.WaitGroup + delayed sync.WaitGroup + delayedMu sync.Mutex + delayedJobs map[*time.Timer]natsMessage + observer queue.Observer } type natsWorkerSubscription interface { @@ -84,6 +86,7 @@ func newNATSWorkerWithConfig(cfg natsWorkerConfig) *natsWorker { defaultQueue: cfg.DefaultQueue, workers: cfg.Workers, handlers: make(map[string]queue.Handler), + delayedJobs: make(map[*time.Timer]natsMessage), observer: cfg.Observer, } } @@ -177,7 +180,7 @@ func (w *natsWorker) Shutdown(ctx context.Context) error { stopErr = sub.Drain() } w.running.Wait() - w.delayed.Wait() + w.flushDelayedJobs() if conn != nil { if drainErr := conn.Drain(); stopErr == nil { stopErr = drainErr @@ -216,13 +219,7 @@ func (w *natsWorker) processMessage(message *nats.Msg) { if incoming.AvailableAtMS > 0 { remaining := time.Until(time.UnixMilli(incoming.AvailableAtMS)) if remaining > 0 { - w.delayed.Add(1) - time.AfterFunc(remaining, func() { - defer w.delayed.Done() - if err := w.republish(incoming); err != nil { - w.observeRepublishFailure(context.Background(), incoming, err) - } - }) + w.scheduleDelayedJob(incoming, remaining) return } } @@ -268,6 +265,45 @@ func (w *natsWorker) processMessage(message *nats.Msg) { } } +// scheduleDelayedJob keeps Core NATS delayed work recoverable because the broker cannot retain a claimed callback until its due time. +func (w *natsWorker) scheduleDelayedJob(message natsMessage, remaining time.Duration) { + w.delayedMu.Lock() + w.delayed.Add(1) + var timer *time.Timer + timer = time.AfterFunc(remaining, func() { + w.delayedMu.Lock() + delete(w.delayedJobs, timer) + w.delayedMu.Unlock() + defer w.delayed.Done() + if err := w.republish(message); err != nil { + w.observeRepublishFailure(context.Background(), message, err) + } + }) + w.delayedJobs[timer] = message + w.delayedMu.Unlock() +} + +// flushDelayedJobs returns timer-owned work to Core NATS before the consumer connection closes so maintenance does not wait for long delays. +func (w *natsWorker) flushDelayedJobs() { + w.delayedMu.Lock() + messages := make([]natsMessage, 0, len(w.delayedJobs)) + for timer, message := range w.delayedJobs { + if !timer.Stop() { + continue + } + delete(w.delayedJobs, timer) + messages = append(messages, message) + w.delayed.Done() + } + w.delayedMu.Unlock() + for _, message := range messages { + if err := w.republish(message); err != nil { + w.observeRepublishFailure(context.Background(), message, err) + } + } + w.delayed.Wait() +} + // connectNATSWorker creates the Core NATS subscription owned by one worker lifecycle. func connectNATSWorker(url, subject string, callback nats.MsgHandler) (natsConnection, natsWorkerSubscription, error) { nc, err := nats.Connect(url) diff --git a/driver/natsqueue/worker_nats_impl_test.go b/driver/natsqueue/worker_nats_impl_test.go index 5720bf3..c795e40 100644 --- a/driver/natsqueue/worker_nats_impl_test.go +++ b/driver/natsqueue/worker_nats_impl_test.go @@ -473,8 +473,8 @@ func TestNATSWorkerShutdownWaitsForInFlightRepublish(t *testing.T) { } } -// TestNATSWorkerShutdownTracksDelayedRepublish verifies timer-backed accepted work finishes before connection drain. -func TestNATSWorkerShutdownTracksDelayedRepublish(t *testing.T) { +// TestNATSWorkerShutdownFlushesDelayedRepublish verifies maintenance can return long-delay work to Core NATS without waiting for its due time. +func TestNATSWorkerShutdownFlushesDelayedRepublish(t *testing.T) { w := newNATSWorker("nats://example:4222") connection, subscription := newNATSWorkerLifecycleStubs() w.conn = connection @@ -483,13 +483,15 @@ func TestNATSWorkerShutdownTracksDelayedRepublish(t *testing.T) { payload, err := json.Marshal(natsMessage{ Type: "job:delayed-shutdown", Queue: "default", - AvailableAtMS: time.Now().Add(25 * time.Millisecond).UnixMilli(), + AvailableAtMS: time.Now().Add(time.Hour).UnixMilli(), }) if err != nil { t.Fatalf("marshal: %v", err) } w.processMessage(&nats.Msg{Data: payload}) - if err := w.Shutdown(context.Background()); err != nil { + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + if err := w.Shutdown(ctx); err != nil { t.Fatalf("shutdown: %v", err) } select { diff --git a/driver/sqlitequeue/worker_lifecycle_test.go b/driver/sqlitequeue/worker_lifecycle_test.go new file mode 100644 index 0000000..2e47ae6 --- /dev/null +++ b/driver/sqlitequeue/worker_lifecycle_test.go @@ -0,0 +1,105 @@ +package sqlitequeue + +import ( + "context" + "path/filepath" + "sync/atomic" + "testing" + "time" + + "github.com/goforj/queue" +) + +// TestSQLiteWorkerLifecycleLeavesPausedJobsPending verifies maintenance-style pauses happen before durable claims. +func TestSQLiteWorkerLifecycleLeavesPausedJobsPending(t *testing.T) { + dsn := filepath.Join(t.TempDir(), "queue.db") + q, err := New(dsn, queue.WithWorkers(1)) + if err != nil { + t.Fatalf("new SQLite queue: %v", err) + } + t.Cleanup(func() { _ = q.Shutdown(context.Background()) }) + + handled := make(chan struct{}, 1) + var calls atomic.Int64 + q.Register("reports:build", func(context.Context, queue.Message) error { + calls.Add(1) + handled <- struct{}{} + return nil + }) + if err := q.PauseWorkers(context.Background()); err != nil { + t.Fatalf("pause workers before startup: %v", err) + } + if _, err := q.Dispatch(queue.NewJob("reports:build").OnQueue("default")); err != nil { + t.Fatalf("dispatch while paused: %v", err) + } + select { + case <-handled: + t.Fatal("paused SQLite worker claimed the job") + case <-time.After(150 * time.Millisecond): + } + stats, err := q.Stats(context.Background()) + if err != nil { + t.Fatalf("read paused queue stats: %v", err) + } + counters := stats.ByQueue["default"] + if counters.Pending != 1 || counters.Active != 0 || counters.Failed != 0 || counters.Archived != 0 { + t.Fatalf("paused queue counters = %+v, want one untouched pending job", counters) + } + + if err := q.ResumeWorkers(context.Background()); err != nil { + t.Fatalf("resume workers: %v", err) + } + select { + case <-handled: + case <-time.After(2 * time.Second): + t.Fatal("resumed SQLite worker did not execute the pending job") + } + counters = waitForSQLiteQueueSettlement(t, q) + if calls.Load() != 1 || counters.Failed != 0 || counters.Archived != 0 { + t.Fatalf("resumed queue counters = %+v, want one successful execution", counters) + } + + if err := q.PauseWorkers(context.Background()); err != nil { + t.Fatalf("pause live workers: %v", err) + } + if _, err := q.Dispatch(queue.NewJob("reports:build").OnQueue("default")); err != nil { + t.Fatalf("dispatch during live pause: %v", err) + } + select { + case <-handled: + t.Fatal("live-paused SQLite worker claimed the job") + case <-time.After(150 * time.Millisecond): + } + if err := q.ResumeWorkers(context.Background()); err != nil { + t.Fatalf("resume live workers: %v", err) + } + select { + case <-handled: + case <-time.After(2 * time.Second): + t.Fatal("live-resumed SQLite worker did not execute the pending job") + } + counters = waitForSQLiteQueueSettlement(t, q) + if calls.Load() != 2 || counters.Failed != 0 || counters.Archived != 0 { + t.Fatalf("live-resumed queue counters = %+v, want two successful executions", counters) + } +} + +// waitForSQLiteQueueSettlement waits for the durable success update that follows handler return. +func waitForSQLiteQueueSettlement(t *testing.T, q *queue.Queue) queue.QueueCounters { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for { + stats, err := q.Stats(context.Background()) + if err != nil { + t.Fatalf("read resumed queue stats: %v", err) + } + counters := stats.ByQueue["default"] + if counters.Pending == 0 && counters.Active == 0 { + return counters + } + if time.Now().After(deadline) { + t.Fatalf("queue did not settle after handler completion; counters = %+v", counters) + } + time.Sleep(10 * time.Millisecond) + } +} diff --git a/driver/sqlqueuecore/queue_database_impl.go b/driver/sqlqueuecore/queue_database_impl.go index 67ba2fe..dd4b8d5 100644 --- a/driver/sqlqueuecore/queue_database_impl.go +++ b/driver/sqlqueuecore/queue_database_impl.go @@ -109,12 +109,14 @@ type databaseQueue struct { cfg localDatabaseConfig db *sql.DB - ownsDB bool + ownsDB bool + prepareSchemaOnUse bool mu sync.RWMutex handlers map[string]queue.Handler startMu sync.Mutex + schemaMu sync.Mutex shutdownOnce sync.Once shutdownDone chan struct{} closeOnce sync.Once @@ -124,6 +126,7 @@ type databaseQueue struct { shutdownCh chan struct{} started atomic.Bool + schemaReady atomic.Bool shuttingDown atomic.Bool uniqueClaims atomic.Uint64 continuation *busruntime.ContinuationScope @@ -204,13 +207,14 @@ func New(cfg queue.DatabaseConfig) (*databaseQueue, error) { } d := &databaseQueue{ - cfg: local, - db: cfg.DB, - handlers: make(map[string]queue.Handler), - shutdownCh: make(chan struct{}), - ownsDB: ownsDB, - continuation: busruntime.NewContinuationScope(), - observer: cfg.Observer, + cfg: local, + db: cfg.DB, + handlers: make(map[string]queue.Handler), + shutdownCh: make(chan struct{}), + ownsDB: ownsDB, + prepareSchemaOnUse: true, + continuation: busruntime.NewContinuationScope(), + observer: cfg.Observer, } if cfg.DriverName == "sqlite" { d.db.SetMaxOpenConns(1) @@ -267,11 +271,7 @@ func (d *databaseQueue) StartWorkers(ctx context.Context) error { if d.started.Load() { return nil } - if d.cfg.AutoMigrate && !d.cfg.DisableAutoMigrate { - if err := d.ensureSchema(ctx); err != nil { - return err - } - } else if err := d.requireManagedQueueSchema(ctx); err != nil { + if err := d.prepareSchema(ctx); err != nil { return err } for i := 0; i < d.cfg.Workers; i++ { @@ -282,6 +282,29 @@ func (d *databaseQueue) StartWorkers(ctx context.Context) error { return nil } +// prepareSchema keeps producer-only and worker database handles on the same lazy, retryable schema boundary. +func (d *databaseQueue) prepareSchema(ctx context.Context) error { + if d.schemaReady.Load() { + return nil + } + d.schemaMu.Lock() + defer d.schemaMu.Unlock() + if d.schemaReady.Load() { + return nil + } + var err error + if d.cfg.AutoMigrate && !d.cfg.DisableAutoMigrate { + err = d.ensureSchema(ctx) + } else { + err = d.requireManagedQueueSchema(ctx) + } + if err != nil { + return err + } + d.schemaReady.Store(true) + return nil +} + // Shutdown drains workers before closing only database handles opened by this queue. func (d *databaseQueue) Shutdown(ctx context.Context) error { if err := d.DrainWorkers(ctx); err != nil { @@ -374,6 +397,11 @@ func (d *databaseQueue) Dispatch(ctx context.Context, job queue.Job) error { return err } } + if d.prepareSchemaOnUse { + if err := d.prepareSchema(ctx); err != nil { + return err + } + } parsed := queuecore.DriverOptions(job) payloadBytes := job.PayloadBytes() if payloadBytes == nil { diff --git a/driver/sqlqueuecore/wrapper.go b/driver/sqlqueuecore/wrapper.go index 630981d..d3f6a25 100644 --- a/driver/sqlqueuecore/wrapper.go +++ b/driver/sqlqueuecore/wrapper.go @@ -24,7 +24,7 @@ type ModuleConfig struct { // for a specific SQL driver name. func NewQueue(driverName string, cfg ModuleConfig, opts ...queue.Option) (*queue.Queue, error) { observer := driverbridge.NewObserverSink(cfg.Observer) - backend, err := New(queue.DatabaseConfig{ + databaseConfig := queue.DatabaseConfig{ DB: cfg.DB, DriverName: driverName, DSN: cfg.DSN, @@ -34,7 +34,8 @@ func NewQueue(driverName string, cfg ModuleConfig, opts ...queue.Option) (*queue ProcessingLeaseNoTimeout: cfg.ProcessingLeaseNoTimeout, Observer: observer, Logger: cfg.Logger, - }) + } + backend, err := New(databaseConfig) if err != nil { return nil, err } @@ -43,5 +44,10 @@ func NewQueue(driverName string, cfg ModuleConfig, opts ...queue.Option) (*queue DefaultQueue: cfg.DefaultQueue, Logger: cfg.Logger, } - return driverbridge.NewQueueFromDriver(rootCfg, observer, backend, nil, opts...) + workerFactory := func(workers int) (any, error) { + workerConfig := databaseConfig + workerConfig.Workers = workers + return New(workerConfig) + } + return driverbridge.NewQueueFromDriver(rootCfg, observer, backend, workerFactory, opts...) } diff --git a/driver_runtime.go b/driver_runtime.go index afb143f..0b2ede9 100644 --- a/driver_runtime.go +++ b/driver_runtime.go @@ -37,7 +37,7 @@ func newQueueFromDriver(cfg Config, observer Observer, backend driverQueueBacken var q queueBackend var runtime runtimeQueueBackend - if native, ok := backend.(driverRuntimeQueueBackend); ok { + if native, ok := backend.(driverRuntimeQueueBackend); ok && workerFactory == nil { runtime = driverRuntimeQueueBackendAdapter{native} q = runtime } else { diff --git a/internal/driverbridge/bridge.go b/internal/driverbridge/bridge.go index 53c0daa..4dc3d92 100644 --- a/internal/driverbridge/bridge.go +++ b/internal/driverbridge/bridge.go @@ -139,6 +139,24 @@ func (a runtimeQueueBackendAdapter) StartWorkers(ctx context.Context) error { return a.inner.StartWorkers(ctx) } +// PauseWorkers preserves optional native lifecycle control without claiming unsupported backends have paused. +func (a runtimeQueueBackendAdapter) PauseWorkers(ctx context.Context) error { + controller, ok := a.inner.(interface{ PauseWorkers(context.Context) error }) + if !ok { + return queue.ErrPauseUnsupported + } + return controller.PauseWorkers(ctx) +} + +// ResumeWorkers preserves optional native lifecycle control without claiming unsupported backends have resumed. +func (a runtimeQueueBackendAdapter) ResumeWorkers(ctx context.Context) error { + controller, ok := a.inner.(interface{ ResumeWorkers(context.Context) error }) + if !ok { + return queue.ErrPauseUnsupported + } + return controller.ResumeWorkers(ctx) +} + // DrainWorkers forwards the native backend's worker-drain lifecycle phase. func (a runtimeQueueBackendAdapter) DrainWorkers(ctx context.Context) error { return a.inner.DrainWorkers(ctx) diff --git a/queue.go b/queue.go index 79d50f7..2fab362 100644 --- a/queue.go +++ b/queue.go @@ -232,6 +232,7 @@ type nativeQueueRuntimeState struct { runtimeRegistrations map[string]struct{} started bool paused bool + pauseIncomplete bool draining bool closed bool start *runtimeStartAttempt @@ -887,7 +888,7 @@ func (q *nativeQueueRuntime) PauseWorkers(ctx context.Context) error { q.mu.Unlock() return ErrQueuerShuttingDown } - if q.paused { + if q.paused && !q.pauseIncomplete { q.mu.Unlock() return nil } @@ -906,6 +907,8 @@ func (q *nativeQueueRuntime) PauseWorkers(ctx context.Context) error { } attempt := &runtimePauseAttempt{done: make(chan struct{})} q.pause = attempt + q.paused = true + q.pauseIncomplete = true started := q.started if !started { for jobType, handler := range q.registered { @@ -919,12 +922,16 @@ func (q *nativeQueueRuntime) PauseWorkers(ctx context.Context) error { var err error if lifecycle, ok := q.runtime.(workerLifecycleBackend); ok { err = lifecycle.PauseWorkers(ctx) + } else { + err = ErrPauseUnsupported } q.mu.Lock() - if err == nil { - q.paused = true + if err != nil { + // The desired state stays paused after a caller deadline so ResumeWorkers can reconcile a backend gate that may already be active. + attempt.err = err + } else { + q.pauseIncomplete = false } - attempt.err = err q.pause = nil close(attempt.done) q.mu.Unlock() @@ -944,7 +951,7 @@ func (q *nativeQueueRuntime) ResumeWorkers(ctx context.Context) error { if q.pause != nil { pause := q.pause q.mu.Unlock() - if err := waitForRuntimePause(ctx, pause); err != nil { + if err := waitForRuntimePauseCompletion(ctx, pause); err != nil { return err } return q.ResumeWorkers(ctx) @@ -965,14 +972,18 @@ func (q *nativeQueueRuntime) ResumeWorkers(ctx context.Context) error { if err := lifecycle.ResumeWorkers(ctx); err != nil { return err } + } else { + return ErrPauseUnsupported } q.mu.Lock() q.paused = false + q.pauseIncomplete = false q.mu.Unlock() return nil } q.mu.Lock() q.paused = false + q.pauseIncomplete = false q.mu.Unlock() if err := q.StartWorkers(ctx); err != nil { q.mu.Lock() @@ -993,7 +1004,7 @@ func (q *externalQueueRuntime) PauseWorkers(ctx context.Context) error { q.mu.Unlock() return ErrQueuerShuttingDown } - if q.paused { + if q.paused && q.worker == nil { q.mu.Unlock() return nil } @@ -1012,6 +1023,7 @@ func (q *externalQueueRuntime) PauseWorkers(ctx context.Context) error { } attempt := &runtimePauseAttempt{done: make(chan struct{})} q.pause = attempt + q.paused = true w := q.worker q.mu.Unlock() @@ -1024,7 +1036,6 @@ func (q *externalQueueRuntime) PauseWorkers(ctx context.Context) error { q.worker = nil q.workerRegistrations = nil q.started = false - q.paused = true } attempt.err = err q.pause = nil @@ -1046,7 +1057,7 @@ func (q *externalQueueRuntime) ResumeWorkers(ctx context.Context) error { if q.pause != nil { pause := q.pause q.mu.Unlock() - if err := waitForRuntimePause(ctx, pause); err != nil { + if err := waitForRuntimePauseCompletion(ctx, pause); err != nil { return err } return q.ResumeWorkers(ctx) @@ -1059,6 +1070,21 @@ func (q *externalQueueRuntime) ResumeWorkers(ctx context.Context) error { } return q.StartWorkers(ctx) } + w := q.worker + q.mu.Unlock() + if w != nil { + if err := w.Shutdown(ctx); err != nil { + return err + } + q.mu.Lock() + if q.worker == w { + q.worker = nil + q.workerRegistrations = nil + q.started = false + } + q.mu.Unlock() + } + q.mu.Lock() q.paused = false q.mu.Unlock() if err := q.StartWorkers(ctx); err != nil { @@ -1264,6 +1290,16 @@ func waitForRuntimePause(ctx context.Context, attempt *runtimePauseAttempt) erro } } +// waitForRuntimePauseCompletion waits for transition ownership without inheriting another caller's diagnostic. +func waitForRuntimePauseCompletion(ctx context.Context, attempt *runtimePauseAttempt) error { + select { + case <-attempt.done: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + // waitForRuntimeStart lets concurrent callers share one startup attempt while honoring their own deadline. func waitForRuntimeStart(ctx context.Context, attempt *runtimeStartAttempt) error { select { diff --git a/queue_local.go b/queue_local.go index c792787..c7ef002 100644 --- a/queue_local.go +++ b/queue_local.go @@ -13,28 +13,27 @@ import ( // localQueue is an in-memory queue implementation supporting sync and workerpool drivers. type localQueue struct { - driver Driver - cfg WorkerpoolConfig - mu sync.RWMutex - metricsMu sync.RWMutex - queueMu sync.RWMutex - handlers map[string]Handler - unique uniqueness.MemoryStore - metrics map[string]*localQueueMetrics - pausedQueues map[string]bool - workQueue chan queuedJob - workPending int - workIdle chan struct{} - continuation *busruntime.ContinuationScope - resizeBuffer bool - shutdownOnce sync.Once - workerWG sync.WaitGroup - workerStateMu sync.Mutex - workerPaused bool - workerStopping bool - workerActive int - workerResume chan struct{} - workerDrained chan struct{} + driver Driver + cfg WorkerpoolConfig + mu sync.RWMutex + metricsMu sync.RWMutex + queueMu sync.RWMutex + handlers map[string]Handler + unique uniqueness.MemoryStore + metrics map[string]*localQueueMetrics + pausedQueues map[string]bool + workQueue chan queuedJob + workPending int + workIdle chan struct{} + continuation *busruntime.ContinuationScope + resizeBuffer bool + shutdownOnce sync.Once + workerWG sync.WaitGroup + workerStateMu sync.Mutex + workerPaused bool + workerActive int + workerResume chan struct{} + workerDrained chan struct{} syncWorkMu sync.Mutex syncWorkPending int @@ -210,7 +209,6 @@ func (d *localQueue) DrainWorkers(ctx context.Context) error { d.shutdownOnce.Do(func() { d.shuttingDown.Store(true) d.workerStateMu.Lock() - d.workerStopping = d.workerPaused d.resumeWorkersLocked() d.workerStateMu.Unlock() }) @@ -577,10 +575,7 @@ func (d *localQueue) worker(workQueue <-chan queuedJob) { defer d.workerWG.Done() jobTimeout := d.cfg.DefaultJobTimeout for job := range workQueue { - if !d.beginWorkerExecution() { - d.finishQueuedWork() - continue - } + d.beginWorkerExecution() func() { defer d.finishQueuedWork() defer d.endWorkerExecution() @@ -625,18 +620,14 @@ func (d *localQueue) worker(workQueue <-chan queuedJob) { } } -// beginWorkerExecution waits outside the backend claim path and rejects pending work during a paused shutdown. -func (d *localQueue) beginWorkerExecution() bool { +// beginWorkerExecution keeps accepted in-memory work behind the lifecycle gate until intake resumes or graceful shutdown drains it. +func (d *localQueue) beginWorkerExecution() { for { d.workerStateMu.Lock() - if d.workerStopping { - d.workerStateMu.Unlock() - return false - } if !d.workerPaused { d.workerActive++ d.workerStateMu.Unlock() - return true + return } resume := d.workerResume d.workerStateMu.Unlock() diff --git a/queue_local_test.go b/queue_local_test.go index 39a99c4..17bb44f 100644 --- a/queue_local_test.go +++ b/queue_local_test.go @@ -81,7 +81,56 @@ func TestLocalQueueWorkerLifecyclePausesIntakeAndDrainsActiveHandlers(t *testing } } -func TestLocalQueuePausedShutdownDoesNotStartPendingHandlers(t *testing.T) { +func TestQueueWorkerLifecycleResumesAfterPauseDeadline(t *testing.T) { + q, err := NewWorkerpool(WithWorkers(1)) + if err != nil { + t.Fatalf("new workerpool: %v", err) + } + started := make(chan struct{}) + release := make(chan struct{}) + handled := make(chan struct{}, 2) + q.Register("job:deadline", func(context.Context, Message) error { + select { + case <-started: + default: + close(started) + <-release + } + handled <- struct{}{} + return nil + }) + if err := q.StartWorkers(context.Background()); err != nil { + t.Fatalf("start workers: %v", err) + } + if _, err := q.Dispatch(NewJob("job:deadline")); err != nil { + t.Fatalf("dispatch active job: %v", err) + } + <-started + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + if err := q.PauseWorkers(ctx); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("pause error = %v, want deadline exceeded", err) + } + if _, err := q.Dispatch(NewJob("job:deadline")); err != nil { + t.Fatalf("dispatch behind timed-out pause: %v", err) + } + close(release) + if err := q.ResumeWorkers(context.Background()); err != nil { + t.Fatalf("resume after pause deadline: %v", err) + } + for range 2 { + select { + case <-handled: + case <-time.After(time.Second): + t.Fatal("resumed worker did not drain accepted jobs") + } + } + if err := q.Shutdown(context.Background()); err != nil { + t.Fatalf("shutdown: %v", err) + } +} + +func TestLocalQueuePausedShutdownDrainsAcceptedHandlers(t *testing.T) { d := newLocalQueueWithConfig(DriverWorkerpool, WorkerpoolConfig{Workers: 1, QueueCapacity: 2}) handled := make(chan struct{}, 1) d.Register("job:pending", func(context.Context, Job) error { @@ -99,8 +148,16 @@ func TestLocalQueuePausedShutdownDoesNotStartPendingHandlers(t *testing.T) { } select { case <-handled: - t.Fatal("paused shutdown started a pending handler") default: + t.Fatal("paused shutdown did not drain an accepted handler") + } + stats, err := d.Stats(context.Background()) + if err != nil { + t.Fatalf("read shutdown stats: %v", err) + } + counters := stats.ByQueue["default"] + if counters.Pending != 0 || counters.Active != 0 || counters.Processed != 1 { + t.Fatalf("shutdown queue counters = %+v, want one drained job", counters) } } diff --git a/queue_runtime_unit_test.go b/queue_runtime_unit_test.go index e724c1e..7799ce5 100644 --- a/queue_runtime_unit_test.go +++ b/queue_runtime_unit_test.go @@ -58,6 +58,32 @@ type phasedShutdownRuntimeBackendStub struct { drainOnce sync.Once } +type latchingWorkerShutdownStub struct { + runtimeBackendStub + stopStarted chan struct{} + stopRelease chan struct{} + stopDone chan struct{} + stopOnce sync.Once +} + +// Shutdown keeps a committed stop running after an individual caller's deadline. +func (s *latchingWorkerShutdownStub) Shutdown(ctx context.Context) error { + s.stopCalls++ + s.stopOnce.Do(func() { + close(s.stopStarted) + go func() { + <-s.stopRelease + close(s.stopDone) + }() + }) + select { + case <-s.stopDone: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + type strictRegistrationRuntimeBackendStub struct { runtimeBackendStub registrations map[string]int @@ -2006,3 +2032,39 @@ func TestExternalQueueRuntimeWorkerLifecyclePreservesProducer(t *testing.T) { t.Fatalf("producer shutdowns = %d, want 1", inner.shutdowns) } } + +func TestExternalQueueRuntimeResumeReconcilesTimedOutPause(t *testing.T) { + first := &latchingWorkerShutdownStub{ + stopStarted: make(chan struct{}), + stopRelease: make(chan struct{}), + stopDone: make(chan struct{}), + } + second := &runtimeBackendStub{} + workers := []driverWorkerBackend{first, second} + nextWorker := 0 + q := &externalQueueRuntime{ + common: &queueCommon{inner: &queueBackendRecorder{}, cfg: Config{DefaultQueue: "default"}, driver: DriverNATS}, + newWorker: func(int) (driverWorkerBackend, error) { + worker := workers[nextWorker] + nextWorker++ + return worker, nil + }, + externalQueueRuntimeState: &externalQueueRuntimeState{registered: map[string]Handler{}}, + } + if err := q.StartWorkers(context.Background()); err != nil { + t.Fatalf("start workers: %v", err) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + if err := q.PauseWorkers(ctx); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("pause error = %v, want deadline exceeded", err) + } + <-first.stopStarted + close(first.stopRelease) + if err := q.ResumeWorkers(context.Background()); err != nil { + t.Fatalf("resume after pause deadline: %v", err) + } + if second.startCalls != 1 { + t.Fatalf("replacement worker starts = %d, want 1", second.startCalls) + } +}