kvserver: add TestBaseQueueCallback

This commit adds TestBaseQueueCallbackOnEnqueueResult and
TestBaseQueueCallbackOnProcessResult to verify that callbacks are correctly
invoked with both enqueue and process results.

Author: wenyihu6

pkg/kv/kvserver/queue_helpers_testutil.go

@@ -21,6 +21,15 @@ func (bq *baseQueue) testingAdd(
 	return bq.addInternal(ctx, repl.Desc(), repl.ReplicaID(), priority, noopProcessCallback)
 }
 
+// testingAddWithCallback is the same as testingAdd, but allows the caller to
+// register a process callback that will be invoked when the replica is enqueued
+// or processed.
+func (bq *baseQueue) testingAddWithCallback(
+	ctx context.Context, repl replicaInQueue, priority float64, callback processCallback,
+) (bool, error) {
+	return bq.addInternal(ctx, repl.Desc(), repl.ReplicaID(), priority, callback)
+}
+
 func forceScanAndProcess(ctx context.Context, s *Store, q *baseQueue) error {
 	// Check that the system config is available. It is needed by many queues. If
 	// it's not available, some queues silently fail to process any replicas,

pkg/kv/kvserver/queue_test.go

@@ -1292,6 +1292,265 @@ func TestBaseQueueDisable(t *testing.T) {
 	}
 }
 
+// TestBaseQueueCallbackOnEnqueueResult tests the callback onEnqueueResult for
+// 1. successful case: the replica is successfully enqueued.
+// 2. priority update: updates the priority of the replica and not enqueuing
+// again.
+// 3. disabled: queue is disabled and the replica is not enqueued.
+// 4. stopped: queue is stopped and the replica is not enqueued.
+// 5. already queued: the replica is already in the queue and not enqueued
+// again.
+// 6. purgatory: the replica is in purgatory and not enqueued again.
+// 7. processing: the replica is already being processed and not enqueued again.
+// 8. full queue: the queue is full and the replica is not enqueued again.
+func TestBaseQueueCallbackOnEnqueueResult(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	tc := testContext{}
+	stopper := stop.NewStopper()
+	ctx := context.Background()
+	defer stopper.Stop(ctx)
+	tc.Start(ctx, t, stopper)
+
+	t.Run("successfuladd", func(t *testing.T) {
+		testQueue := &testQueueImpl{}
+		bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: 1})
+		r, err := tc.store.GetReplica(1)
+		require.NoError(t, err)
+		queued, _ := bq.testingAddWithCallback(ctx, r, 1.0, processCallback{
+			onEnqueueResult: func(indexOnHeap int, err error) {
+				require.Equal(t, 0, indexOnHeap)
+				require.NoError(t, err)
+			},
+			onProcessResult: func(err error) {
+				t.Fatal("unexpected call to onProcessResult")
+			},
+		})
+		require.True(t, queued)
+	})
+
+	t.Run("priority", func(t *testing.T) {
+		testQueue := &testQueueImpl{}
+		bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: 5})
+		r, err := tc.store.GetReplica(1)
+		require.NoError(t, err)
+		priorities := []float64{5.0, 4.0, 8.0, 1.0, 3.0}
+		expectedIndices := []int{0, 1, 0, 3, 4}
+		// When inserting 5, [5], index 0.
+		// When inserting 4, [5, 4], index 1.
+		// When inserting 8, [8, 4, 5], index 0.
+		// When inserting 1, [8, 4, 5, 1], index 3.
+		// When inserting 3, [8, 4, 5, 1, 3], index 4.
+		for i, priority := range priorities {
+			r.Desc().RangeID = roachpb.RangeID(i + 1)
+			queued, _ := bq.testingAddWithCallback(ctx, r, priority, processCallback{
+				onEnqueueResult: func(indexOnHeap int, err error) {
+					require.Equal(t, expectedIndices[i], indexOnHeap)
+					require.NoError(t, err)
+				},
+				onProcessResult: func(err error) {
+					t.Fatal("unexpected call to onProcessResult")
+				},
+			})
+			require.True(t, queued)
+		}
+		// Set range id back to 1.
+		r.Desc().RangeID = 1
+	})
+	t.Run("disabled", func(t *testing.T) {
+		testQueue := &testQueueImpl{}
+		bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: 2})
+		bq.SetDisabled(true)
+		r, err := tc.store.GetReplica(1)
+		require.NoError(t, err)
+		queued, _ := bq.testingAddWithCallback(ctx, r, 1.0, processCallback{
+			onEnqueueResult: func(indexOnHeap int, err error) {
+				require.Equal(t, -1, indexOnHeap)
+				require.ErrorIs(t, err, errQueueDisabled)
+			},
+			onProcessResult: func(err error) {
+				t.Fatal("unexpected call to onProcessResult")
+			},
+		})
+		require.False(t, queued)
+	})
+	t.Run("stopped", func(t *testing.T) {
+		testQueue := &testQueueImpl{}
+		bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: 2})
+		bq.mu.stopped = true
+		r, err := tc.store.GetReplica(1)
+		require.NoError(t, err)
+		queued, _ := bq.testingAddWithCallback(ctx, r, 1.0, processCallback{
+			onEnqueueResult: func(indexOnHeap int, err error) {
+				require.Equal(t, -1, indexOnHeap)
+				require.ErrorIs(t, err, errQueueStopped)
+			},
+			onProcessResult: func(err error) {
+				t.Fatal("unexpected call to onProcessResult")
+			},
+		})
+		require.False(t, queued)
+	})
+
+	t.Run("alreadyqueued", func(t *testing.T) {
+		testQueue := &testQueueImpl{}
+		bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: 2})
+		r, err := tc.store.GetReplica(1)
+		require.NoError(t, err)
+		queued, _ := bq.testingAddWithCallback(ctx, r, 1.0, processCallback{
+			onEnqueueResult: func(indexOnHeap int, err error) {
+				require.Equal(t, 0, indexOnHeap)
+				require.NoError(t, err)
+			},
+			onProcessResult: func(err error) {
+				t.Fatal("unexpected call to onProcessResult")
+			},
+		})
+		require.True(t, queued)
+
+		// Inserting again on the same range id should fail.
+		queued, _ = bq.testingAddWithCallback(ctx, r, 1.0, processCallback{
+			onEnqueueResult: func(indexOnHeap int, err error) {
+				require.Equal(t, -1, indexOnHeap)
+				require.ErrorIs(t, err, errReplicaAlreadyInQueue)
+			},
+			onProcessResult: func(err error) {
+				t.Fatal("unexpected call to onProcessResult")
+			},
+		})
+		require.False(t, queued)
+	})
+
+	t.Run("purgatory", func(t *testing.T) {
+		testQueue := &testQueueImpl{
+			pChan: make(chan time.Time, 1),
+		}
+		bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: 2})
+		r, err := tc.store.GetReplica(1)
+		require.NoError(t, err)
+		bq.mu.Lock()
+		bq.addToPurgatoryLocked(ctx, stopper, r, &testPurgatoryError{}, 1.0, nil)
+		bq.mu.Unlock()
+		// Inserting a range in purgatory should not enqueue again.
+		queued, _ := bq.testingAddWithCallback(ctx, r, 1.0, processCallback{
+			onEnqueueResult: func(indexOnHeap int, err error) {
+				require.Equal(t, -1, indexOnHeap)
+				require.ErrorIs(t, err, errReplicaAlreadyInPurgatory)
+			},
+			onProcessResult: func(err error) {
+				t.Fatal("unexpected call to onProcessResult")
+			},
+		})
+		require.False(t, queued)
+	})
+
+	t.Run("processing", func(t *testing.T) {
+		testQueue := &testQueueImpl{}
+		bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: 2})
+		r, err := tc.store.GetReplica(1)
+		require.NoError(t, err)
+		item := &replicaItem{rangeID: r.Desc().RangeID, replicaID: r.ReplicaID(), index: -1}
+		item.setProcessing()
+		bq.addLocked(item)
+		// Inserting a range that is already being processed should not enqueue again.
+		requeued, _ := bq.testingAddWithCallback(ctx, r, 1.0, processCallback{
+			onEnqueueResult: func(indexOnHeap int, err error) {
+				require.Equal(t, -1, indexOnHeap)
+				require.ErrorIs(t, err, errReplicaAlreadyProcessing)
+			},
+			onProcessResult: func(err error) {
+				t.Fatal("unexpected call to onProcessResult")
+			},
+		})
+		require.True(t, requeued)
+	})
+	t.Run("fullqueue", func(t *testing.T) {
+		testQueue := &testQueueImpl{}
+		bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: 0})
+		r, err := tc.store.GetReplica(1)
+		require.NoError(t, err)
+		// Max size is 0, so the replica should not be enqueued.
+		queued, _ := bq.testingAddWithCallback(ctx, r, 1.0, processCallback{
+			onEnqueueResult: func(indexOnHeap int, err error) {
+				// It may be called with err = nil.
+				if err != nil {
+					require.ErrorIs(t, err, errDroppedDueToFullQueueSize)
+				}
+			},
+			onProcessResult: func(err error) {
+				t.Fatal("unexpected call to onProcessResult")
+			},
+		})
+		require.True(t, queued)
+	})
+}
+
+// TestBaseQueueCallbackOnProcessResult tests that the processCallback is
+// invoked when the replica is processed and will be invoked again if the
+// replica ends up in the purgatory queue and being processed again.
+func TestBaseQueueCallbackOnProcessResult(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	tc := testContext{}
+	stopper := stop.NewStopper()
+	ctx := context.Background()
+	defer stopper.Stop(ctx)
+	tsc := TestStoreConfig(nil)
+	tc.StartWithStoreConfig(ctx, t, stopper, tsc)
+
+	testQueue := &testQueueImpl{
+		duration: time.Nanosecond,
+		pChan:    make(chan time.Time, 1),
+		err:      &testPurgatoryError{},
+	}
+
+	const replicaCount = 10
+	repls := createReplicas(t, &tc, replicaCount)
+
+	bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: replicaCount})
+	bq.Start(stopper)
+
+	var totalProcessedCalledWithErr atomic.Int32
+	for _, r := range repls {
+		queued, _ := bq.testingAddWithCallback(context.Background(), r, 1.0, processCallback{
+			onEnqueueResult: func(indexOnHeap int, err error) {
+				require.NoError(t, err)
+			},
+			onProcessResult: func(err error) {
+				if err != nil {
+					totalProcessedCalledWithErr.Add(1)
+				}
+			},
+		})
+		require.True(t, queued)
+	}
+
+	testutils.SucceedsSoon(t, func() error {
+		if pc := testQueue.getProcessed(); pc != replicaCount {
+			return errors.Errorf("expected %d processed replicas; got %d", replicaCount, pc)
+		}
+
+		if totalProcessedCalledWithErr.Load() != int32(replicaCount) {
+			return errors.Errorf("expected %d processed replicas with err; got %d", replicaCount, totalProcessedCalledWithErr.Load())
+		}
+		return nil
+	})
+
+	// Now, signal that purgatoried replicas should retry.
+	testQueue.pChan <- timeutil.Now()
+
+	testutils.SucceedsSoon(t, func() error {
+		if pc := testQueue.getProcessed(); pc != replicaCount*2 {
+			return errors.Errorf("expected %d processed replicas; got %d", replicaCount, pc)
+		}
+
+		if totalProcessedCalledWithErr.Load() != int32(replicaCount*2) {
+			return errors.Errorf("expected %d processed replicas with err; got %d", replicaCount, totalProcessedCalledWithErr.Load())
+		}
+		return nil
+	})
+}
+
 // TestQueueDisable verifies that setting the set of queue.enabled cluster
 // settings actually disables the base queue. This test works alongside
 // TestBaseQueueDisable to verify the entire disable workflow.