diff --git a/backend/pkg/queue/queue.go b/backend/pkg/queue/queue.go index 55d30a64..75ca14fe 100644 --- a/backend/pkg/queue/queue.go +++ b/backend/pkg/queue/queue.go @@ -34,8 +34,14 @@ type queue[I any, O any] struct { mx *sync.Mutex wg *sync.WaitGroup - ctx context.Context - cancel context.CancelFunc + // ctx tracks "running"; it is cancelled by Stop() AND by the reader on normal + // input-close. stopCtx is a HARD stop, cancelled only by Stop() — workers and + // the reader bail on stopCtx (not ctx) so a normal input-close still drains and + // delivers every result instead of dropping in-flight ones. + ctx context.Context + cancel context.CancelFunc + stopCtx context.Context + stopCancel context.CancelFunc instance uuid.UUID workers int @@ -49,6 +55,8 @@ func NewQueue[I, O any](input <-chan I, output chan O, workers int, process func mx, wg := &sync.Mutex{}, &sync.WaitGroup{} ctx, cancel := context.WithCancel(context.Background()) cancel() + stopCtx, stopCancel := context.WithCancel(context.Background()) + stopCancel() if workers <= 0 { workers = defaultWorkersAmount @@ -58,8 +66,10 @@ func NewQueue[I, O any](input <-chan I, output chan O, workers int, process func mx: mx, wg: wg, - ctx: ctx, - cancel: cancel, + ctx: ctx, + cancel: cancel, + stopCtx: stopCtx, + stopCancel: stopCancel, instance: uuid.New(), workers: workers, @@ -89,6 +99,7 @@ func (q *queue[I, O]) Start() error { } q.ctx, q.cancel = context.WithCancel(context.Background()) + q.stopCtx, q.stopCancel = context.WithCancel(context.Background()) q.queue = make(chan *message[I], q.workers*2) q.wg.Add(q.workers) @@ -112,11 +123,14 @@ func (q *queue[I, O]) Start() error { func (q *queue[I, O]) Stop() error { q.mx.Lock() - if q.ctx.Err() != nil { + // Guard on stopCtx, not ctx: a normal input-close already cancelled ctx, but + // the queue is not hard-stopped until Stop() runs. + if q.stopCtx.Err() != nil { q.mx.Unlock() return ErrAlreadyStopped } + q.stopCancel() q.cancel() q.mx.Unlock() q.wg.Wait() @@ -150,22 +164,20 @@ func (q *queue[I, O]) worker(wid int) { } else if result, err := q.process(msg.value); err != nil { logger.WithError(err).Error("failed to process message") } else { - // Wait for the previous message to be sent (preserves output order), - // then send this one. Both waits also select on q.ctx so a stopped - // queue whose consumer quit reading output can't block the worker - // forever and deadlock Stop() -> wg.Wait(). On stop we return WITHOUT - // msg.cancel(): leaving the next message's doneCtx uncancelled makes - // every later worker bail via q.ctx too, so output ends at a contiguous - // prefix instead of developing gaps. + // Bail only on a hard Stop() (q.stopCtx), never q.ctx — the reader + // cancels q.ctx on a NORMAL input-close, and bailing there would drop + // still-undelivered results. On a hard stop, return without msg.cancel() + // so later workers stay blocked on their doneCtx and bail too, keeping + // output a contiguous prefix. select { case <-msg.doneCtx.Done(): - case <-q.ctx.Done(): + case <-q.stopCtx.Done(): return } select { case q.output <- result: - case <-q.ctx.Done(): + case <-q.stopCtx.Done(): return } } @@ -195,7 +207,7 @@ func (q *queue[I, O]) reader() { for { select { - case <-q.ctx.Done(): + case <-q.stopCtx.Done(): return case value, ok := <-q.input: // check if the input channel is closed and exit if so @@ -209,15 +221,15 @@ func (q *queue[I, O]) reader() { // create a new context for each message newCtx, cancel := context.WithCancel(context.Background()) - // select on q.ctx so a stopped queue whose workers have exited - // cannot block the reader forever on a full q.queue. + // Bail only on a hard Stop() (q.stopCtx) so a stopped queue whose + // workers have exited can't block the reader on a full q.queue. select { case q.queue <- &message[I]{ value: value, doneCtx: lastDoneCtx, cancel: cancel, }: - case <-q.ctx.Done(): + case <-q.stopCtx.Done(): cancel() return } diff --git a/backend/pkg/queue/queue_test.go b/backend/pkg/queue/queue_test.go index 6713431b..d654480d 100644 --- a/backend/pkg/queue/queue_test.go +++ b/backend/pkg/queue/queue_test.go @@ -129,12 +129,12 @@ func TestQueue_ProcessOrdering(t *testing.T) { } } -// A consumer that stops reading output (e.g. ListContainerDir bailing on the -// first stat error) leaves workers blocked on the unbuffered send; Stop() must -// still return instead of hanging on wg.Wait(). +// Stop() must return when the consumer stops reading output (e.g. +// ListContainerDir bailing on a stat error), leaving workers blocked on the +// send; otherwise wg.Wait() hangs. func TestQueue_StopDoesNotDeadlockWithUnreadOutput(t *testing.T) { input := make(chan int, 20) - output := make(chan int) // unbuffered; deliberately left unread below + output := make(chan int) workers := 4 q := queue.NewQueue(input, output, workers, func(i int) (int, error) { @@ -149,16 +149,13 @@ func TestQueue_StopDoesNotDeadlockWithUnreadOutput(t *testing.T) { } close(input) - <-output // take one result, then abandon the channel with items still in flight + <-output // read one, then abandon output with items still in flight done := make(chan error, 1) go func() { done <- q.Stop() }() select { - case err := <-done: - if err != nil { - t.Fatalf("Stop returned an error: %v", err) - } + case <-done: case <-time.After(3 * time.Second): t.Fatal("Stop() deadlocked with unread output") } @@ -168,6 +165,73 @@ func TestQueue_StopDoesNotDeadlockWithUnreadOutput(t *testing.T) { } } +// The ListContainerDir happy path: the consumer reads every result after input +// closes. A normal input-close must not make workers drop in-flight results, or +// the consumer waits forever for the last one. +func TestQueue_DeliversEveryResultAfterInputClose(t *testing.T) { + const n = 200 + input := make(chan int, n) + output := make(chan int) // unbuffered, like ListContainerDir's outputStats + for i := 0; i < n; i++ { + input <- i + } + close(input) + + q := queue.NewQueue(input, output, 20, func(i int) (int, error) { return i, nil }) + if err := q.Start(); err != nil { + t.Fatalf("failed to start queue: %v", err) + } + + done := make(chan struct{}) + go func() { + for i := 0; i < n; i++ { + <-output + } + close(done) + }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("consumer hung waiting for a result after input close") + } + + _ = q.Stop() +} + +// Small input closes before the consumer bails, so the reader reaches +// input-close (which cancels ctx) first. Stop() must still hard-stop the +// blocked workers instead of short-circuiting as "already stopped". +func TestQueue_StopHardStopsAfterInputClose(t *testing.T) { + input := make(chan int, 4) + output := make(chan int) + for i := 0; i < 4; i++ { + input <- i + } + close(input) + + q := queue.NewQueue(input, output, 4, func(i int) (int, error) { return i, nil }) + if err := q.Start(); err != nil { + t.Fatalf("failed to start queue: %v", err) + } + + <-output // read one, abandon the rest + for q.Running() { + time.Sleep(time.Millisecond) // wait until the reader processes input-close + } + + done := make(chan error, 1) + go func() { done <- q.Stop() }() + select { + case err := <-done: + if err != nil { + t.Fatalf("Stop() short-circuited (%v) instead of hard-stopping blocked workers", err) + } + case <-time.After(3 * time.Second): + t.Fatal("Stop() hung after input close") + } +} + func BenchmarkQueue_DefaultWorkers(b *testing.B) { simpleBenchmark(b, 0) }