fix: Don't drop work queued by the last active table in shuffle-queue

scheduler/queue/scheduler.go

@@ -89,13 +89,12 @@ func (d *Scheduler) Sync(ctx context.Context, tableClients []WorkUnit, resolvedR
 	if len(tableClients) == 0 {
 		return
 	}
-	queue := NewConcurrentRandomQueue[WorkUnit](d.seed, len(tableClients))
+	queue := newWorkQueue(d.seed, len(tableClients))
 	for _, tc := range tableClients {
 		queue.Push(tc)
 	}
 
 	jobs := make(chan *WorkUnit)
-	activeWorkSignal := newActiveWorkSignal()
 
 	// Worker pool
 	workerPool, _ := errgroup.WithContext(ctx)
@@ -112,37 +111,16 @@ func (d *Scheduler) Sync(ctx context.Context, tableClients []WorkUnit, resolvedR
 				d.metrics,
 				msgChan,
 				d.errorClassifier,
-			).work(ctx, activeWorkSignal)
+			).work(ctx)
 			return nil
 		})
 	}
 
 	// Work distribution
 	go func() {
 		defer close(jobs)
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			default:
-				item := queue.Pop()
-
-				// There is work to do. Mark it active before handing it off so the
-				// idle check below cannot fire while the item is in flight to a worker.
-				if item != nil {
-					activeWorkSignal.Add()
-					jobs <- item
-					continue
-				}
-
-				// Queue is empty and no active work, done!
-				if activeWorkSignal.IsIdle() {
-					return
-				}
-
-				// Queue is empty and there is active work, wait for changes
-				activeWorkSignal.Wait()
-			}
+		for item := queue.Next(ctx); item != nil; item = queue.Next(ctx) {
+			jobs <- item
 		}
 	}()
 

scheduler/queue/work_queue.go

@@ -0,0 +1,76 @@
+package queue
+
+import (
+	"context"
+	"sync"
+)
+
+// workQueue holds the work units waiting for a worker and counts the ones in progress.
+//
+// The sync is done when no work unit is queued and none is in progress. Workers queue
+// relations while their table is in progress, so an empty queue alone does not mean done.
+//
+//   - Push queues a work unit.
+//   - Next takes a work unit off the queue and marks it in progress.
+//   - Done marks a work unit taken by Next as finished.
+type workQueue struct {
+	mu         sync.Mutex
+	queue      *ConcurrentRandomQueue[WorkUnit]
+	inProgress int
+	// Buffered with capacity 1 so a change that happens before Next starts to wait is kept rather than lost.
+	changed chan struct{}
+}
+
+func newWorkQueue(seed int64, capacityHint int) *workQueue {
+	return &workQueue{
+		queue:   NewConcurrentRandomQueue[WorkUnit](seed, capacityHint),
+		changed: make(chan struct{}, 1),
+	}
+}
+
+func (q *workQueue) Push(item WorkUnit) {
+	q.mu.Lock()
+	q.queue.Push(item)
+	q.mu.Unlock()
+	q.notify()
+}
+
+// Next blocks until a work unit is queued and returns it marked in progress.
+// It returns nil when no work unit is queued or in progress, or when ctx is done.
+func (q *workQueue) Next(ctx context.Context) *WorkUnit {
+	for ctx.Err() == nil {
+		q.mu.Lock()
+		item := q.queue.Pop()
+		if item != nil {
+			q.inProgress++
+		}
+		idle := item == nil && q.inProgress == 0
+		q.mu.Unlock()
+
+		if item != nil {
+			return item
+		}
+		if idle {
+			return nil
+		}
+		select {
+		case <-q.changed:
+		case <-ctx.Done():
+		}
+	}
+	return nil
+}
+
+func (q *workQueue) Done() {
+	q.mu.Lock()
+	q.inProgress--
+	q.mu.Unlock()
+	q.notify()
+}
+
+func (q *workQueue) notify() {
+	select {
+	case q.changed <- struct{}{}:
+	default:
+	}
+}

scheduler/queue/work_queue_test.go

@@ -0,0 +1,108 @@
+package queue
+
+import (
+	"context"
+	"testing"
+	"testing/synctest"
+
+	"github.com/cloudquery/plugin-sdk/v4/schema"
+	"github.com/stretchr/testify/require"
+)
+
+func workUnit(name string) WorkUnit {
+	return WorkUnit{Table: &schema.Table{Name: name}}
+}
+
+// startNext calls Next in a goroutine and returns after Next blocks or returns.
+func startNext(ctx context.Context, q *workQueue) <-chan *WorkUnit {
+	next := make(chan *WorkUnit, 1)
+	go func() { next <- q.Next(ctx) }()
+	synctest.Wait()
+	return next
+}
+
+func requireNextReturned(t *testing.T, next <-chan *WorkUnit) *WorkUnit {
+	t.Helper()
+	synctest.Wait()
+	select {
+	case item := <-next:
+		return item
+	default:
+		require.FailNow(t, "Next is still blocked")
+		return nil
+	}
+}
+
+func TestWorkQueue(t *testing.T) {
+	t.Run("returns nil when nothing is queued or in progress", func(t *testing.T) {
+		synctest.Test(t, func(t *testing.T) {
+			require.Nil(t, newWorkQueue(0, 0).Next(t.Context()))
+		})
+	})
+
+	t.Run("wakes up when a work unit in progress queues a relation", func(t *testing.T) {
+		synctest.Test(t, func(t *testing.T) {
+			q := newWorkQueue(0, 1)
+			q.Push(workUnit("parent"))
+			require.Equal(t, "parent", q.Next(t.Context()).Table.Name)
+
+			next := startNext(t.Context(), q)
+			q.Push(workUnit("relation"))
+
+			require.Equal(t, "relation", requireNextReturned(t, next).Table.Name)
+		})
+	})
+
+	t.Run("returns a relation that the last work unit in progress queued before it finished", func(t *testing.T) {
+		synctest.Test(t, func(t *testing.T) {
+			q := newWorkQueue(0, 1)
+			q.Push(workUnit("parent"))
+			require.Equal(t, "parent", q.Next(t.Context()).Table.Name)
+
+			next := startNext(t.Context(), q)
+			q.Push(workUnit("relation"))
+			q.Done()
+
+			require.Equal(t, "relation", requireNextReturned(t, next).Table.Name)
+		})
+	})
+
+	t.Run("returns nil when the last work unit in progress finishes", func(t *testing.T) {
+		synctest.Test(t, func(t *testing.T) {
+			q := newWorkQueue(0, 1)
+			q.Push(workUnit("parent"))
+			require.Equal(t, "parent", q.Next(t.Context()).Table.Name)
+
+			next := startNext(t.Context(), q)
+			q.Done()
+
+			require.Nil(t, requireNextReturned(t, next))
+		})
+	})
+
+	t.Run("returns nil when the context is done and work is still queued", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(t.Context())
+		q := newWorkQueue(0, 2)
+		q.Push(workUnit("first"))
+		q.Push(workUnit("second"))
+		require.NotNil(t, q.Next(ctx))
+
+		cancel()
+
+		require.Nil(t, q.Next(ctx))
+	})
+
+	t.Run("returns nil when the context is done", func(t *testing.T) {
+		synctest.Test(t, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(t.Context())
+			q := newWorkQueue(0, 1)
+			q.Push(workUnit("parent"))
+			require.Equal(t, "parent", q.Next(ctx).Table.Name)
+
+			next := startNext(ctx, q)
+			cancel()
+
+			require.Nil(t, requireNextReturned(t, next))
+		})
+	})
+}

scheduler/queue/worker.go

@@ -23,7 +23,7 @@ import (
 
 type worker struct {
 	jobs              <-chan *WorkUnit
-	queue             *ConcurrentRandomQueue[WorkUnit]
+	queue             *workQueue
 	resolvedResources chan<- *schema.Resource
 
 	logger            zerolog.Logger
@@ -36,20 +36,16 @@ type worker struct {
 	errorClassifier schema.ErrorClassifier
 }
 
-func (w *worker) work(ctx context.Context, activeWorkSignal *activeWorkSignal) {
+func (w *worker) work(ctx context.Context) {
 	for j := range w.jobs {
-		// the work unit was already marked active by the dispatcher before it was
-		// handed off, so the dispatcher can never observe an idle state while a
-		// job is in flight between the queue and a worker
 		w.resolveTable(ctx, j.Table, j.Client, j.Parent)
-
-		activeWorkSignal.Done()
+		w.queue.Done()
 	}
 }
 
 func newWorker(
 	jobs <-chan *WorkUnit,
-	queue *ConcurrentRandomQueue[WorkUnit],
+	queue *workQueue,
 	resolvedResources chan<- *schema.Resource,
 
 	logger zerolog.Logger,