Harden backup and restore operations

internal/queue/queue.go

@@ -25,10 +25,12 @@ type Queue struct {
 	stopping bool
 	handler  Handler
 	onChange OnChange
+	done     chan struct{}
+	doneOnce sync.Once
 }
 
 func New(handler Handler, onChange OnChange) *Queue {
-	q := &Queue{handler: handler, onChange: onChange}
+	q := &Queue{handler: handler, onChange: onChange, done: make(chan struct{})}
 	q.cond = sync.NewCond(&q.mu)
 	return q
 }
@@ -71,6 +73,7 @@ func (q *Queue) Enqueue(run model.Run) error {
 }
 
 func (q *Queue) Run(ctx context.Context) {
+	defer q.doneOnce.Do(func() { close(q.done) })
 	for {
 		q.mu.Lock()
 		for len(q.pending) == 0 && !q.stopping {
@@ -104,6 +107,15 @@ func (q *Queue) Run(ctx context.Context) {
 	}
 }
 
+func (q *Queue) Wait(ctx context.Context) error {
+	select {
+	case <-q.done:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 func (q *Queue) Cancel(id string) bool {
 	q.mu.Lock()
 	defer q.mu.Unlock()

internal/queue/queue_test.go

@@ -90,3 +90,33 @@ func TestPauseAndResumeActiveRun(t *testing.T) {
 	}
 	close(release)
 }
+
+func TestWaitDoesNotReturnBeforeCancelledHandlerCleanup(t *testing.T) {
+	started := make(chan struct{})
+	cleanup := make(chan struct{})
+	q := New(func(ctx context.Context, run model.Run) model.Run {
+		close(started)
+		<-ctx.Done()
+		close(cleanup)
+		run.Status = "cancelled"
+		return run
+	}, nil)
+	ctx, cancel := context.WithCancel(context.Background())
+	go q.Run(ctx)
+	if err := q.Enqueue(model.Run{ID: "run", JobID: "job"}); err != nil {
+		t.Fatal(err)
+	}
+	<-started
+	q.Stop()
+	waitCtx, waitCancel := context.WithTimeout(context.Background(), time.Second)
+	defer waitCancel()
+	if err := q.Wait(waitCtx); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-cleanup:
+	default:
+		t.Fatal("queue wait returned before handler cleanup")
+	}
+	cancel()
+}