kvserver: pass context into write functor

Epic: none
Release note: none

Author: pav-kv

pkg/kv/kvserver/kvstorage/snaprecv/sst_snapshot_storage.go

@@ -144,7 +144,7 @@ func (s *SSTSnapshotStorageScratch) NewFile(
 // WriteSST creates an SST populated with the given write function, and writes
 // it to a file. Does nothing if no data is written.
 func (s *SSTSnapshotStorageScratch) WriteSST(
-	ctx context.Context, write func(storage.Writer) error,
+	ctx context.Context, write func(context.Context, storage.Writer) error,
 ) error {
 	if s.closed {
 		return errors.AssertionFailedf("SSTSnapshotStorageScratch closed")
@@ -154,7 +154,7 @@ func (s *SSTSnapshotStorageScratch) WriteSST(
 	sstFile := &storage.MemObject{}
 	w := storage.MakeIngestionSSTWriter(ctx, s.st, sstFile)
 	defer w.Close()
-	if err := write(&w); err != nil {
+	if err := write(ctx, &w); err != nil {
 		return err
 	}
 	if err := w.Finish(); err != nil {

pkg/kv/kvserver/snapshot_apply_prepare.go

@@ -27,7 +27,7 @@ type snapWriteBuilder struct {
 
 	todoEng  storage.Engine
 	sl       stateloader.StateLoader
-	writeSST func(context.Context, func(storage.Writer) error) error
+	writeSST func(context.Context, func(context.Context, storage.Writer) error) error
 
 	truncState    kvserverpb.RaftTruncatedState
 	hardState     raftpb.HardState
@@ -42,14 +42,9 @@ type snapWriteBuilder struct {
 // prepareSnapApply writes the unreplicated SST for the snapshot and clears disk data for subsumed replicas.
 func (s *snapWriteBuilder) prepareSnapApply(ctx context.Context) error {
 	_ = applySnapshotTODO // 3.1 + 1.1 + 2.5.
-	if err := s.writeSST(ctx, func(w storage.Writer) error {
-		// Clear the raft state/log, and initialize it again with the provided
-		// HardState and RaftTruncatedState.
-		return s.rewriteRaftState(ctx, w)
-	}); err != nil {
+	if err := s.writeSST(ctx, s.rewriteRaftState); err != nil {
 		return err
 	}
-
 	_ = applySnapshotTODO // 3.2 + 2.1 + 2.2 + 2.3
 	return s.clearSubsumedReplicaDiskData(ctx, s.todoEng)
 }
@@ -121,7 +116,7 @@ func (s *snapWriteBuilder) clearSubsumedReplicaDiskData(
 	totalKeySpans := append([]roachpb.Span(nil), keySpans...)
 	for _, subDesc := range s.subsumedDescs {
 		// We have to create an SST for the subsumed replica's range-id local keys.
-		if err := s.writeSST(ctx, func(w storage.Writer) error {
+		if err := s.writeSST(ctx, func(ctx context.Context, w storage.Writer) error {
 			// NOTE: We set mustClearRange to true because we are setting
 			// RangeTombstoneKey. Since Clears and Puts need to be done in increasing
 			// order of keys, it is not safe to use ClearRangeIter.
@@ -200,7 +195,7 @@ func (s *snapWriteBuilder) clearSubsumedReplicaDiskData(
 		//
 		// We need to additionally clear [b,sn).
 
-		if err := s.writeSST(ctx, func(w storage.Writer) error {
+		if err := s.writeSST(ctx, func(ctx context.Context, w storage.Writer) error {
 			return storage.ClearRangeWithHeuristic(
 				ctx, reader, w,
 				keySpans[i].EndKey, totalKeySpans[i].EndKey,

pkg/kv/kvserver/snapshot_apply_prepare_test.go

@@ -43,12 +43,12 @@ func TestPrepareSnapApply(t *testing.T) {
 
 	var sb redact.StringBuilder
 
-	writeSST := func(ctx context.Context, write func(storage.Writer) error) error {
+	writeSST := func(ctx context.Context, write func(context.Context, storage.Writer) error) error {
 		// Use WriteBatch so that we print the writes in exactly the order in which
 		// they are made. The real code creates an SST writer.
 		b := eng.NewWriteBatch()
 		defer b.Close()
-		if err := write(b); err != nil {
+		if err := write(ctx, b); err != nil {
 			return err
 		}
 		str, err := print.DecodeWriteBatch(b.Repr())