Merge #155953 #159067

155953: kvserver: simplify rewriteRaftState r=arulajmani a=pav-kv

Keys in the RangeID-local unreplicated space have both state machine and raft state. With separated engines, these will reside in different engines, and are also interleaved in an unfortunate way. So we can't use one `ClearRawRange` to cover them all, or even one per engine.

This PR clears all unreplicated keys in `rewriteRaftState` manually, so that we can control which goes to which engine/`Writer`. Additionally, this uncovers that half of these keys don't need to be cleared at all because they don't exist, or don't need to be changed.

The end effect of the PR is that `rewriteRaftState` now only mutates the raft engine state, which is very convenient. Incidentally, this function finally lives up to its name.

Part of #152845
Epic: CRDB-55220

159067: util/ioctx: add RandomAccessReader for cloud storage r=sravotto a=sravotto

Add a general-purpose wrapper that provides ReadAt/Seek support for any cloud storage reader using an opener factory function pattern.

This enables random access (required by formats like Parquet) without requiring the cloud storage implementation itself to natively support ReadAt/Seek. Instead, it opens new readers at specific offsets for each ReadAt call.

Key features:
- Thread-safe ReadAt: Each call opens a new connection, allowing concurrent reads from different offsets
- Seek support: Tracks position logically without seeking underlying readers
- Generic: Works with any ReadCloserCtx via OpenerAtFunc factory

Release note: None

Epic: CRDB-23802

Co-authored-by: Pavel Kalinnikov <[email protected]>
Co-authored-by: Silvano Ravotto <[email protected]>

Author: 3 people

pkg/kv/kvserver/client_raft_test.go

@@ -3259,15 +3259,18 @@ func TestRaftRemoveRace(t *testing.T) {
 	}
 	tc.AddVotersOrFatal(t, key, targets...)
 
+	sl := kvstorage.MakeStateLoader(desc.RangeID)
+	s2 := tc.GetFirstStoreFromServer(t, 2)
 	for i := 0; i < 10; i++ {
 		tc.RemoveVotersOrFatal(t, key, tc.Target(2))
 		tc.AddVotersOrFatal(t, key, tc.Target(2))
 
-		// Verify the tombstone key does not exist. See #12130.
-		ts, err := kvstorage.MakeStateLoader(desc.RangeID).LoadRangeTombstone(
-			ctx, tc.GetFirstStoreFromServer(t, 2).StateEngine())
+		replID, err := sl.LoadRaftReplicaID(ctx, s2.StateEngine())
 		require.NoError(t, err)
-		require.Equal(t, kvserverpb.RangeTombstone{}, ts)
+		ts, err := sl.LoadRangeTombstone(ctx, s2.StateEngine())
+		require.NoError(t, err)
+		// ReplicaID leads the RangeTombstone, which means a replica exists.
+		require.GreaterOrEqual(t, replID.ReplicaID, ts.NextReplicaID)
 	}
 }
 

pkg/kv/kvserver/kvstorage/destroy.go

@@ -12,6 +12,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/rditer"
+	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
@@ -239,6 +240,37 @@ func RemoveStaleRHSFromSplit(
 	return nil
 }
 
+// RewriteRaftState rewrites the raft state of the given replica with the
+// provided state. Specifically, it rewrites HardState and RaftTruncatedState,
+// and clears the raft log. All writes are generated in the engine keys order.
+//
+// TODO(pav-kv): get rid of the returned cleared spans.
+func RewriteRaftState(
+	ctx context.Context,
+	raftWO RaftWO,
+	sl StateLoader,
+	hs raftpb.HardState,
+	ts kvserverpb.RaftTruncatedState,
+) (cleared roachpb.Span, _ error) {
+	// Update HardState.
+	if err := sl.SetHardState(ctx, raftWO, hs); err != nil {
+		return roachpb.Span{}, errors.Wrapf(err, "unable to write HardState")
+	}
+	// Clear the raft log. Note that there are no Pebble range keys in this span.
+	logPrefix := sl.RaftLogPrefix().Clone()
+	raftLog := roachpb.Span{Key: logPrefix, EndKey: logPrefix.PrefixEnd()}
+	if err := raftWO.ClearRawRange(
+		raftLog.Key, raftLog.EndKey, true /* pointKeys */, false, /* rangeKeys */
+	); err != nil {
+		return roachpb.Span{}, errors.Wrapf(err, "unable to clear the raft log")
+	}
+	// Update the log truncation state.
+	if err := sl.SetRaftTruncatedState(ctx, raftWO, &ts); err != nil {
+		return roachpb.Span{}, errors.Wrapf(err, "unable to write RaftTruncatedState")
+	}
+	return raftLog, nil
+}
+
 // TestingForceClearRange changes the value of ClearRangeThresholdPointKeys to
 // 1, which effectively forces ClearRawRange in the replica destruction path,
 // instead of point deletions. This can be used for making the storage

pkg/kv/kvserver/snapshot_apply_prepare.go

@@ -8,19 +8,19 @@ package kvserver
 import (
 	"context"
 
-	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/rditer"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
-	"github.com/cockroachdb/errors"
 )
 
 // snapWriteBuilder contains the data needed to prepare the on-disk state for a
 // snapshot.
+//
+// TODO(pav-kv): move this struct to kvstorage package.
 type snapWriteBuilder struct {
 	id roachpb.FullReplicaID
 
@@ -42,7 +42,11 @@ 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 {
+	// TODO(pav-kv): assert that our replica already exists in storage. Note that
+	// it can be either uninitialized or initialized.
 	_ = applySnapshotTODO // 1.1 + 1.3 + 2.4 + 3.1
+	// TODO(sep-raft-log): rewriteRaftState now only touches raft engine keys, so
+	// it will be convenient to redirect it to a raft engine batch.
 	if err := s.writeSST(ctx, s.rewriteRaftState); err != nil {
 		return err
 	}
@@ -55,42 +59,16 @@ func (s *snapWriteBuilder) prepareSnapApply(ctx context.Context) error {
 	return s.clearResidualDataOnNarrowSnapshot(ctx)
 }
 
-// rewriteRaftState clears and rewrites the unreplicated rangeID-local key space
-// of the given replica with the provided raft state. Note that it also clears
-// the raft log contents.
-//
-// The caller must make sure the log does not have entries newer than the
-// snapshot entry ID, and that clearing the log is applied atomically with the
-// snapshot write, or after the latter is synced.
+// rewriteRaftState rewrites the raft state of the given replica with the
+// provided state. Specifically, it rewrites HardState and RaftTruncatedState,
+// and clears the raft log. All writes are generated in the engine keys order.
 func (s *snapWriteBuilder) rewriteRaftState(ctx context.Context, w storage.Writer) error {
-	// Clearing the unreplicated state.
-	//
-	// NB: We do not expect to see range keys in the unreplicated state, so
-	// we don't drop a range tombstone across the range key space.
-	unreplicatedPrefixKey := keys.MakeRangeIDUnreplicatedPrefix(s.id.RangeID)
-	unreplicatedStart := unreplicatedPrefixKey
-	unreplicatedEnd := unreplicatedPrefixKey.PrefixEnd()
-	if err := w.ClearRawRange(
-		unreplicatedStart, unreplicatedEnd, true /* pointKeys */, false, /* rangeKeys */
-	); err != nil {
-		return errors.Wrapf(err, "error clearing the unreplicated space")
-	}
-
-	// Update HardState.
-	if err := s.sl.SetHardState(ctx, w, s.hardState); err != nil {
-		return errors.Wrapf(err, "unable to write HardState")
-	}
-	// We've cleared all the raft state above, so we are forced to write the
-	// RaftReplicaID again here.
-	if err := s.sl.SetRaftReplicaID(ctx, w, s.id.ReplicaID); err != nil {
-		return errors.Wrapf(err, "unable to write RaftReplicaID")
-	}
-	// Update the log truncation state.
-	if err := s.sl.SetRaftTruncatedState(ctx, w, &s.truncState); err != nil {
-		return errors.Wrapf(err, "unable to write RaftTruncatedState")
+	cleared, err := kvstorage.RewriteRaftState(
+		ctx, kvstorage.RaftWO(w), s.sl, s.hardState, s.truncState)
+	if err != nil {
+		return err
 	}
-
-	s.cleared = append(s.cleared, roachpb.Span{Key: unreplicatedStart, EndKey: unreplicatedEnd})
+	s.cleared = append(s.cleared, cleared)
 	return nil
 }
 

pkg/kv/kvserver/snapshot_apply_prepare_test.go

@@ -12,9 +12,11 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/print"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/rditer"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
@@ -130,4 +132,24 @@ func createRangeData(t *testing.T, eng storage.Engine, desc roachpb.RangeDescrip
 		}.ToEngineKey(nil)
 		require.NoError(t, eng.PutEngineKey(ek, nil))
 	}
+
+	// Add some raft state: HardState, TruncatedState and log entries.
+	const truncIndex, numEntries = 10, 3
+	ctx := context.Background()
+
+	sl := logstore.NewStateLoader(desc.RangeID)
+	require.NoError(t, sl.SetRaftTruncatedState(
+		ctx, eng, &kvserverpb.RaftTruncatedState{Index: truncIndex, Term: 5},
+	))
+	require.NoError(t, sl.SetHardState(
+		ctx, eng, raftpb.HardState{Term: 6, Commit: truncIndex + numEntries},
+	))
+	for i := truncIndex + 1; i <= truncIndex+numEntries; i++ {
+		require.NoError(t, storage.MVCCBlindPutProto(
+			ctx, eng,
+			sl.RaftLogKey(kvpb.RaftIndex(i)), hlc.Timestamp{},
+			&raftpb.Entry{Index: uint64(i), Term: 6},
+			storage.MVCCWriteOptions{},
+		))
+	}
 }

pkg/kv/kvserver/testdata/TestPrepareSnapApply.txt

@@ -1,16 +1,25 @@
 echo
 ----
 >> sst:
-Delete Range: [0,0 /Local/RangeID/123/u"" (0x0169f67b7500): , 0,0 /Local/RangeID/123/v"" (0x0169f67b7600): )
 Put: 0,0 /Local/RangeID/123/u/RaftHardState (0x0169f67b757266746800): term:20 vote:0 commit:100 lead:0 lead_epoch:0 
-Put: 0,0 /Local/RangeID/123/u/RaftReplicaID (0x0169f67b757266747200): replica_id:4 
+Delete Range: [0,0 /Local/RangeID/123/u/RaftLog (0x0169f67b757266746c00): , 0,0 /Local/RangeID/123/u"rftm" (0x0169f67b757266746d00): )
 Put: 0,0 /Local/RangeID/123/u/RaftTruncatedState (0x0169f67b757266747400): index:100 term:20 
 >> sst:
 Put: 0,0 /Local/RangeID/101/u/RangeTombstone (0x0169ed757266746200): next_replica_id:2147483647 
+Delete (Sized at 38): 0,0 /Local/RangeID/101/u/RaftHardState (0x0169ed757266746800): 
+Delete (Sized at 42): 0,0 /Local/RangeID/101/u/RaftLog/logIndex:11 (0x0169ed757266746c000000000000000b00): 
+Delete (Sized at 42): 0,0 /Local/RangeID/101/u/RaftLog/logIndex:12 (0x0169ed757266746c000000000000000c00): 
+Delete (Sized at 42): 0,0 /Local/RangeID/101/u/RaftLog/logIndex:13 (0x0169ed757266746c000000000000000d00): 
 Delete: 0,0 /Local/RangeID/101/u/RaftReplicaID (0x0169ed757266747200): 
+Delete (Sized at 32): 0,0 /Local/RangeID/101/u/RaftTruncatedState (0x0169ed757266747400): 
 >> sst:
 Put: 0,0 /Local/RangeID/102/u/RangeTombstone (0x0169ee757266746200): next_replica_id:2147483647 
+Delete (Sized at 38): 0,0 /Local/RangeID/102/u/RaftHardState (0x0169ee757266746800): 
+Delete (Sized at 42): 0,0 /Local/RangeID/102/u/RaftLog/logIndex:11 (0x0169ee757266746c000000000000000b00): 
+Delete (Sized at 42): 0,0 /Local/RangeID/102/u/RaftLog/logIndex:12 (0x0169ee757266746c000000000000000c00): 
+Delete (Sized at 42): 0,0 /Local/RangeID/102/u/RaftLog/logIndex:13 (0x0169ee757266746c000000000000000d00): 
 Delete: 0,0 /Local/RangeID/102/u/RaftReplicaID (0x0169ee757266747200): 
+Delete (Sized at 32): 0,0 /Local/RangeID/102/u/RaftTruncatedState (0x0169ee757266747400): 
 >> sst:
 Delete (Sized at 27): /Local/Lock"y\xff" 0300000000000000000000000000000000 (0x017a6b1279ff000100030000000000000000000000000000000012): 
 >> sst:
@@ -20,7 +29,7 @@ Delete (Sized at 12): 0.000000001,0 "y\xff" (0x79ff00000000000000000109):
 >> repl: /Local/Lock/Local/Range"{a"-k"}
 >> repl: /Local/Lock"{a"-k"}
 >> repl: {a-k}
->> cleared: /Local/RangeID/123/{u""-v""}
+>> cleared: /Local/RangeID/123/u{/RaftLog-"rftm"}
 >> cleared: /Local/RangeID/101/{r""-s""}
 >> cleared: /Local/RangeID/101/{u""-v""}
 >> cleared: /Local/RangeID/102/{r""-s""}

pkg/util/ioctx/BUILD.bazel

@@ -2,7 +2,10 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "ioctx",
-    srcs = ["reader.go"],
+    srcs = [
+        "random_access_reader.go",
+        "reader.go",
+    ],
     importpath = "github.com/cockroachdb/cockroach/pkg/util/ioctx",
     visibility = ["//visibility:public"],
     deps = ["@com_github_cockroachdb_errors//:errors"],
@@ -11,7 +14,10 @@ go_library(
 go_test(
     name = "ioctx_test",
     size = "small",
-    srcs = ["reader_test.go"],
+    srcs = [
+        "random_access_reader_test.go",
+        "reader_test.go",
+    ],
     embed = [":ioctx"],
     deps = [
         "@com_github_cockroachdb_errors//:errors",

pkg/util/ioctx/random_access_reader.go

@@ -0,0 +1,132 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package ioctx
+
+import (
+	"context"
+	"io"
+
+	"github.com/cockroachdb/errors"
+)
+
+// OpenerAtFunc is a factory function that opens a reader at a specific offset.
+// endHint can be used to optimize the read if the end position is known.
+type OpenerAtFunc func(ctx context.Context, offset int64, endHint int64) (ReadCloserCtx, error)
+
+// randomAccessReader wraps any cloud storage reader and adds ReadAt/Seek support
+// using a factory function that can open readers at specific offsets.
+//
+// This provides random access capabilities without requiring the cloud storage
+// implementation itself to support seeking. Each ReadAt call opens a new reader
+// at the requested offset.
+//
+// Thread Safety:
+//   - ReadAt: Safe for concurrent calls. Each call opens its own reader.
+//   - Read/Seek: NOT safe for concurrent calls. Should only be used from a single goroutine.
+type randomAccessReader struct {
+	ctx    context.Context
+	size   int64
+	openAt OpenerAtFunc
+
+	// Current sequential reader (for Read operations)
+	current ReadCloserCtx
+	pos     int64
+}
+
+var _ io.ReaderAt = &randomAccessReader{}
+var _ io.ReadSeekCloser = &randomAccessReader{}
+
+// NewRandomAccessReader creates a reader that supports ReadAt and Seek
+// by using the provided opener function to create readers at specific offsets.
+//
+// Parameters:
+//   - ctx: Context for operations
+//   - size: Total size of the file
+//   - openAt: Factory function that opens a reader at a given offset
+func NewRandomAccessReader(
+	ctx context.Context, size int64, openAt OpenerAtFunc,
+) *randomAccessReader {
+	return &randomAccessReader{
+		ctx:    ctx,
+		size:   size,
+		openAt: openAt,
+	}
+}
+
+// Read implements io.Reader using sequential access
+func (r *randomAccessReader) Read(p []byte) (int, error) {
+	if r.current == nil {
+		// Open reader at current position
+		opened, err := r.openAt(r.ctx, r.pos, r.size)
+		if err != nil {
+			return 0, err
+		}
+		r.current = opened
+	}
+
+	n, err := r.current.Read(r.ctx, p)
+	r.pos += int64(n)
+	return n, err
+}
+
+// ReadAt implements io.ReaderAt - safe for concurrent calls.
+// Each call opens a new reader at the requested offset, so multiple
+// goroutines can call ReadAt simultaneously.
+//
+// Note: This opens a new connection (HTTP request) for every ReadAt call.
+// For cloud storage, this means each random access read results in a new
+// range request to the storage service.
+func (r *randomAccessReader) ReadAt(p []byte, off int64) (int, error) {
+	if off >= r.size {
+		return 0, io.EOF
+	}
+
+	// Open a new reader at the requested offset
+	reader, err := r.openAt(r.ctx, off, off+int64(len(p)))
+	if err != nil {
+		return 0, err
+	}
+	defer reader.Close(r.ctx)
+
+	// Read exactly len(p) bytes or until EOF
+	return io.ReadFull(ReaderCtxAdapter(r.ctx, reader), p)
+}
+
+// Seek implements io.Seeker
+func (r *randomAccessReader) Seek(offset int64, whence int) (int64, error) {
+	var newPos int64
+	switch whence {
+	case io.SeekStart:
+		newPos = offset
+	case io.SeekCurrent:
+		newPos = r.pos + offset
+	case io.SeekEnd:
+		newPos = r.size + offset
+	default:
+		return 0, errors.Newf("invalid whence: %d", whence)
+	}
+
+	if newPos < 0 {
+		return 0, errors.New("negative position")
+	}
+
+	// Close current reader if position changed
+	if r.current != nil && newPos != r.pos {
+		r.current.Close(r.ctx)
+		r.current = nil
+	}
+
+	r.pos = newPos
+	return newPos, nil
+}
+
+// Close closes any open reader
+func (r *randomAccessReader) Close() error {
+	if r.current != nil {
+		return r.current.Close(r.ctx)
+	}
+	return nil
+}