allocatorimpl: add invariants on priority to base queue tests

This commit adds additional invariants to verify the correctness of priority
plumbing for range items in base queue tests.

Author: wenyihu6

pkg/kv/kvserver/allocator/allocatorimpl/BUILD.bazel

@@ -28,6 +28,7 @@ go_library(
         "//pkg/settings",
         "//pkg/settings/cluster",
         "//pkg/util/admission/admissionpb",
+        "//pkg/util/buildutil",
         "//pkg/util/log",
         "//pkg/util/metric",
         "//pkg/util/stop",

pkg/kv/kvserver/queue.go

@@ -1092,19 +1092,24 @@ func IsPurgatoryError(err error) (PurgatoryError, bool) {
 }
 
 // assertInvariants codifies the guarantees upheld by the data structures in the
-// base queue. In summary, a replica is one of:
+// base queue.
+// 1. In summary, a replica is one of:
 // - "queued" and in mu.replicas and mu.priorityQ
 // - "processing" and only in mu.replicas
 // - "purgatory" and in mu.replicas and mu.purgatory
+// 2. For every item in bq.mu.priorityQ.sl, bq.mu.purgatory, and bq.mu.replicas,
+// assertOnReplicaItem callback is called with the item. Note that we expect
+// items in priorityQ and purgatory to be in replicas.
 //
 // Note that in particular, nothing is ever in both mu.priorityQ and
 // mu.purgatory.
-func (bq *baseQueue) assertInvariants() {
+func (bq *baseQueue) assertInvariants(assertOnReplicaItem func(item *replicaItem)) {
 	bq.mu.Lock()
 	defer bq.mu.Unlock()
 
 	ctx := bq.AnnotateCtx(context.Background())
 	for _, item := range bq.mu.priorityQ.sl {
+		assertOnReplicaItem(item)
 		if item.processing {
 			log.Fatalf(ctx, "processing item found in prioQ: %v", item)
 		}
@@ -1117,6 +1122,7 @@ func (bq *baseQueue) assertInvariants() {
 	}
 	for rangeID := range bq.mu.purgatory {
 		item, inReplicas := bq.mu.replicas[rangeID]
+		assertOnReplicaItem(item)
 		if !inReplicas {
 			log.Fatalf(ctx, "item found in purg but not in mu.replicas: %v", item)
 		}
@@ -1131,6 +1137,7 @@ func (bq *baseQueue) assertInvariants() {
 	// that there aren't any non-processing replicas *only* in bq.mu.replicas.
 	var nNotProcessing int
 	for _, item := range bq.mu.replicas {
+		assertOnReplicaItem(item)
 		if !item.processing {
 			nNotProcessing++
 		}

pkg/kv/kvserver/queue_concurrency_test.go

@@ -112,7 +112,7 @@ func TestBaseQueueConcurrent(t *testing.T) {
 			})
 		}
 		g.Go(func() error {
-			bq.assertInvariants()
+			bq.assertInvariants(func(item *replicaItem) {})
 			return nil
 		})
 	}

pkg/kv/kvserver/queue_test.go

@@ -222,6 +222,11 @@ func TestBaseQueueAddUpdateAndRemove(t *testing.T) {
 	} else {
 		bq.finishProcessingReplica(ctx, stopper, r2, nil)
 	}
+	bq.assertInvariants(func(item *replicaItem) {
+		replica, err := bq.getReplica(item.rangeID)
+		require.NoError(t, err)
+		require.Equal(t, priorityMap[replica.(*Replica)], item.priority)
+	})
 	if v := bq.pending.Value(); v != 1 {
 		t.Errorf("expected 1 pending replicas; got %d", v)
 	}
@@ -294,6 +299,11 @@ func TestBaseQueueAddUpdateAndRemove(t *testing.T) {
 	if r, _ := bq.pop(); r != nil {
 		t.Errorf("expected empty queue; got %v", r)
 	}
+	bq.assertInvariants(func(item *replicaItem) {
+		replica, err := bq.getReplica(item.rangeID)
+		require.NoError(t, err)
+		require.Equal(t, priorityMap[replica.(*Replica)], item.priority)
+	})
 
 	// Try removing a replica.
 	bq.maybeAdd(ctx, r1, hlc.ClockTimestamp{})
@@ -313,6 +323,11 @@ func TestBaseQueueAddUpdateAndRemove(t *testing.T) {
 	if v := bq.pending.Value(); v != 0 {
 		t.Errorf("expected 0 pending replicas; got %d", v)
 	}
+	bq.assertInvariants(func(item *replicaItem) {
+		replica, err := bq.getReplica(item.rangeID)
+		require.NoError(t, err)
+		require.Equal(t, priorityMap[replica.(*Replica)], item.priority)
+	})
 }
 
 // TestBaseQueueSamePriorityFIFO verifies that if multiple items are queued at
@@ -538,19 +553,27 @@ func TestBaseQueueAddRemove(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	const testPriority = 1.0
 	testQueue := &testQueueImpl{
 		blocker: make(chan struct{}, 1),
 		shouldQueueFn: func(now hlc.ClockTimestamp, r *Replica) (shouldQueue bool, priority float64) {
 			shouldQueue = true
-			priority = 1.0
+			priority = testPriority
 			return
 		},
 	}
 	bq := makeTestBaseQueue("test", testQueue, tc.store, queueConfig{maxSize: 2})
 	bq.Start(stopper)
 
 	bq.maybeAdd(ctx, r, hlc.ClockTimestamp{})
+	bq.assertInvariants(func(item *replicaItem) {
+		require.Equal(t, testPriority, item.priority)
+	})
+
 	bq.MaybeRemove(r.RangeID)
+	bq.assertInvariants(func(item *replicaItem) {
+		require.Equal(t, testPriority, item.priority)
+	})
 
 	// Wake the queue
 	close(testQueue.blocker)
@@ -837,10 +860,19 @@ func TestBaseQueuePurgatory(t *testing.T) {
 		bq.maybeAdd(context.Background(), r, hlc.ClockTimestamp{})
 	}
 
+	// Make sure priority is preserved during processing.
+	bq.assertInvariants(func(item *replicaItem) {
+		require.Equal(t, float64(item.rangeID), item.priority)
+	})
+
 	testutils.SucceedsSoon(t, func() error {
 		if pc := testQueue.getProcessed(); pc != replicaCount {
 			return errors.Errorf("expected %d processed replicas; got %d", replicaCount, pc)
 		}
+		// Make sure priorities are preserved with the purgatory queue.
+		bq.assertInvariants(func(item *replicaItem) {
+			require.Equal(t, float64(item.rangeID), item.priority)
+		})
 		// We have to loop checking the following conditions because the increment
 		// of testQueue.processed does not happen atomically with the replica being
 		// placed in purgatory.
@@ -852,6 +884,9 @@ func TestBaseQueuePurgatory(t *testing.T) {
 		if l := bq.Length(); l != 0 {
 			return errors.Errorf("expected empty priorityQ; got %d", l)
 		}
+		bq.assertInvariants(func(item *replicaItem) {
+			require.Equal(t, float64(item.rangeID), item.priority)
+		})
 		// Check metrics.
 		if v := bq.successes.Count(); v != 0 {
 			return errors.Errorf("expected 0 processed replicas; got %d", v)
@@ -886,6 +921,9 @@ func TestBaseQueuePurgatory(t *testing.T) {
 		if l := bq.Length(); l != 0 {
 			return errors.Errorf("expected empty priorityQ; got %d", l)
 		}
+		bq.assertInvariants(func(item *replicaItem) {
+			require.Equal(t, float64(item.rangeID), item.priority)
+		})
 		// Check metrics.
 		if v := bq.successes.Count(); v != 0 {
 			return errors.Errorf("expected 0 processed replicas; got %d", v)
@@ -918,6 +956,9 @@ func TestBaseQueuePurgatory(t *testing.T) {
 		if pc := testQueue.getProcessed(); pc != replicaCount*3-rmReplCount {
 			return errors.Errorf("expected %d processed replicas; got %d", replicaCount*3-rmReplCount, pc)
 		}
+		bq.assertInvariants(func(item *replicaItem) {
+			require.Equal(t, float64(item.rangeID), item.priority)
+		})
 		// Check metrics.
 		if v := bq.successes.Count(); v != int64(replicaCount)-rmReplCount {
 			return errors.Errorf("expected %d processed replicas; got %d", replicaCount-rmReplCount, v)
@@ -961,6 +1002,9 @@ func TestBaseQueuePurgatory(t *testing.T) {
 		if pc := testQueue.getProcessed(); pc != beforeProcessCount+1 {
 			return errors.Errorf("expected %d processed replicas; got %d", beforeProcessCount+1, pc)
 		}
+		bq.assertInvariants(func(item *replicaItem) {
+			require.Equal(t, float64(item.rangeID), item.priority)
+		})
 		if v := bq.successes.Count(); v != beforeSuccessCount+1 {
 			return errors.Errorf("expected %d processed replicas; got %d", beforeSuccessCount+1, v)
 		}
@@ -1365,11 +1409,12 @@ func TestBaseQueueProcessConcurrently(t *testing.T) {
 	repls := createReplicas(t, &tc, 3)
 	r1, r2, r3 := repls[0], repls[1], repls[2]
 
+	const testPriority = 1
 	pQueue := &parallelQueueImpl{
 		testQueueImpl: testQueueImpl{
 			blocker: make(chan struct{}, 1),
 			shouldQueueFn: func(now hlc.ClockTimestamp, r *Replica) (shouldQueue bool, priority float64) {
-				return true, 1
+				return true, testPriority
 			},
 		},
 		processBlocker: make(chan struct{}, 1),
@@ -1414,6 +1459,9 @@ func TestBaseQueueProcessConcurrently(t *testing.T) {
 
 	pQueue.processBlocker <- struct{}{}
 	assertProcessedAndProcessing(3, 0)
+	bq.assertInvariants(func(item *replicaItem) {
+		require.Equal(t, float64(testPriority), item.priority)
+	})
 }
 
 // TestBaseQueueReplicaChange ensures that if a replica is added to the queue