Merge pull request #15825 from tamird/storage-context

storage: plumb contexts

pkg/storage/gc_queue_test.go

@@ -552,7 +552,7 @@ func TestGCQueueTransactionTable(t *testing.T) {
 
 	batch := tc.engine.NewSnapshot()
 	defer batch.Close()
-	tc.repl.assertState(batch) // check that in-mem and on-disk state were updated
+	tc.repl.assertState(context.TODO(), batch) // check that in-mem and on-disk state were updated
 
 	tc.repl.mu.Lock()
 	txnSpanThreshold := tc.repl.mu.state.TxnSpanGCThreshold

pkg/storage/replica.go

@@ -629,7 +629,7 @@ func (r *Replica) initLocked(
 	if err := r.setReplicaIDLocked(replicaID); err != nil {
 		return err
 	}
-	r.assertStateRLocked(r.store.Engine())
+	r.assertStateRLocked(ctx, r.store.Engine())
 	return nil
 }
 
@@ -1333,19 +1333,18 @@ func (r *Replica) State() storagebase.RangeInfo {
 	return ri
 }
 
-func (r *Replica) assertState(reader engine.Reader) {
+func (r *Replica) assertState(ctx context.Context, reader engine.Reader) {
 	r.mu.RLock()
 	defer r.mu.RUnlock()
-	r.assertStateRLocked(reader)
+	r.assertStateRLocked(ctx, reader)
 }
 
 // assertStateRLocked can be called from the Raft goroutine to check that the
 // in-memory and on-disk states of the Replica are congruent. See also
 // assertState if the replica mutex is not currently held.
 //
 // TODO(tschottdorf): Consider future removal (for example, when #7224 is resolved).
-func (r *Replica) assertStateRLocked(reader engine.Reader) {
-	ctx := r.AnnotateCtx(context.TODO())
+func (r *Replica) assertStateRLocked(ctx context.Context, reader engine.Reader) {
 	diskState, err := r.stateLoader.load(ctx, reader, r.mu.state.Desc)
 	if err != nil {
 		log.Fatal(ctx, err)

pkg/storage/replica_proposal.go

@@ -799,6 +799,6 @@ func (r *Replica) handleEvalResult(
 	if shouldAssert {
 		// Assert that the on-disk state doesn't diverge from the in-memory
 		// state as a result of the side effects.
-		r.assertState(r.store.Engine())
+		r.assertState(ctx, r.store.Engine())
 	}
 }

pkg/storage/replica_raftstorage.go

@@ -684,7 +684,7 @@ func (r *Replica) applySnapshot(
 	r.store.metrics.subtractMVCCStats(r.mu.state.Stats)
 	r.store.metrics.addMVCCStats(s.Stats)
 	r.mu.state = s
-	r.assertStateRLocked(r.store.Engine())
+	r.assertStateRLocked(ctx, r.store.Engine())
 	r.mu.Unlock()
 
 	// As the last deferred action after committing the batch, update other

pkg/storage/replica_test.go

@@ -304,7 +304,7 @@ func (tc *testContext) addBogusReplicaToRangeDesc(
 	}
 
 	tc.repl.setDescWithoutProcessUpdate(&newDesc)
-	tc.repl.assertState(tc.engine)
+	tc.repl.assertState(ctx, tc.engine)
 	return secondReplica, nil
 }