Merge #143705

143705: raft: replace MsgStorageAppend protocol r=tbg a=pav-kv

This PR replaces the `MsgStorageAppend` protocol with a more ergonomic / strongly-typed API.

Benefits of this change:

- 3 fields are removed from `raftpb.Message`, and its in-memory size is decreased from 184 to 144 bytes.
- Storage interaction is completely removed from `raftpb.Message`, and there is a cleaner separation between messages and local storage interaction.
- Improved readability and safety of the `Ready` API.
- Log storage code upstream of `pkg/raft` can use the new `StorageAppend/Ack` types to carry their semantics.

---

The old flow:

1. If there are pending storage writes, `RawNode.Ready()` constructs a `MsgStorageAppend` and puts it in `Ready.Messages`.
2. The application filters it out from other messages, and sends it to local storage.
3. When the write is done, the `RawNode` must be notified, and some messages (such as `MsgAppResp`) must be sent to other peers. The `MsgStorageAppend` contains a list of remote and local messages in `Responses`, which includes a `MsgStorageAppendResp`. The application filters out local messages and steps them into the `RawNode`. It sends all other messages via raft transport.

---

The new flow:

1. If there are pending storage writes, `RawNode.Ready()` constructs a `StorageAppend` struct and embeds it into `Ready`.
2. All messages in `Ready.Messages` are not subject to durability, so can be sent immediately without filtering.
3. The application sends `Ready.StorageWrite` to local storage.
4. When the write is done/durable, the application routes the `StorageAppend.Ack()` to `RawNode.AckAppend(ack)` method, to notify the local `RawNode`, and sends all remote messages in `ack.Responses` via the raft transport.

----

Part of #143652, #124440 

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

Author: craig[bot]

pkg/kv/kvserver/client_manual_proposal_test.go

@@ -205,14 +205,13 @@ LIMIT
 
 		stats := &logstore.AppendStats{}
 
-		msgApp := raftpb.Message{
-			Type:      raftpb.MsgStorageAppend,
-			To:        raft.LocalAppendThread,
-			Term:      lastTerm,
-			LogTerm:   lastTerm,
-			Index:     uint64(lastIndex),
+		app := raft.StorageAppend{
+			HardState: raftpb.HardState{
+				Term:   lastTerm,
+				Commit: uint64(lastIndex) + uint64(len(ents)),
+			},
 			Entries:   ents,
-			Commit:    uint64(lastIndex) + uint64(len(ents)),
+			LeadTerm:  lastTerm,
 			Responses: []raftpb.Message{{}}, // need >0 responses so StoreEntries will sync
 		}
 
@@ -246,7 +245,7 @@ LIMIT
 		_, err = ls.StoreEntries(ctx, logstore.RaftState{
 			LastIndex: lastIndex,
 			LastTerm:  kvpb.RaftTerm(lastTerm),
-		}, logstore.MakeMsgStorageAppend(msgApp), (*wgSyncCallback)(wg), stats)
+		}, app, (*wgSyncCallback)(wg), stats)
 		require.NoError(t, err)
 		wg.Wait()
 
@@ -259,7 +258,7 @@ LIMIT
 type wgSyncCallback sync.WaitGroup
 
 func (w *wgSyncCallback) OnLogSync(
-	context.Context, logstore.MsgStorageAppendDone, storage.BatchCommitStats,
+	context.Context, raft.StorageAppendAck, storage.BatchCommitStats,
 ) {
 	(*sync.WaitGroup)(w).Done()
 }

pkg/kv/kvserver/client_raft_log_queue_test.go

@@ -205,7 +205,8 @@ func TestRaftTracing(t *testing.T) {
 					// the ordering may change between 1->2 and 1->3. It should be
 					// sufficient to just check one of them for tracing.
 					`replica_raft.* 1->2 MsgApp`,
-					`replica_raft.* AppendThread->1 MsgStorageAppendResp`,
+					`replica_raft.* appended entries`,
+					`replica_raft.* synced log storage write at mark`,
 					`replica_raft.* applying entries`,
 					`ack-ing replication success to the client`,
 				}

pkg/kv/kvserver/kvflowcontrol/rac2/range_controller.go

@@ -432,18 +432,19 @@ type MsgAppSender interface {
 func RaftEventFromMsgStorageAppendAndMsgApps(
 	mode RaftMsgAppMode,
 	replicaID roachpb.ReplicaID,
-	appendMsg raftpb.Message,
+	appendMsg raft.StorageAppend,
 	outboundMsgs []raftpb.Message,
 	logSnapshot raft.LogSnapshot,
 	msgAppScratch map[roachpb.ReplicaID][]raftpb.Message,
 	replicaStateInfoMap map[roachpb.ReplicaID]ReplicaStateInfo,
 ) RaftEvent {
 	event := RaftEvent{
-		MsgAppMode: mode, LogSnapshot: logSnapshot, ReplicasStateInfo: replicaStateInfoMap}
-	if appendMsg.Type == raftpb.MsgStorageAppend {
-		event.Term = appendMsg.LogTerm
-		event.Snap = appendMsg.Snapshot
-		event.Entries = appendMsg.Entries
+		MsgAppMode:        mode,
+		Term:              appendMsg.LeadTerm,
+		Snap:              appendMsg.Snapshot,
+		Entries:           appendMsg.Entries,
+		LogSnapshot:       logSnapshot,
+		ReplicasStateInfo: replicaStateInfoMap,
 	}
 	if len(outboundMsgs) == 0 || mode == MsgAppPull {
 		// MsgAppPull mode can have MsgApps with entries under some cases: (a)

pkg/kv/kvserver/kvflowcontrol/rac2/range_controller_test.go

@@ -1766,14 +1766,11 @@ func testingFirst(args ...interface{}) interface{} {
 func TestRaftEventFromMsgStorageAppendAndMsgAppsBasic(t *testing.T) {
 	// raftpb.Entry and raftpb.Message are only partially populated below, which
 	// could be improved in the future.
-	appendMsg := raftpb.Message{
-		Type:     raftpb.MsgStorageAppend,
-		LogTerm:  10,
+	appendMsg := raft.StorageAppend{
+		LeadTerm: 10,
 		Snapshot: &raftpb.Snapshot{},
 		Entries: []raftpb.Entry{
-			{
-				Term: 9,
-			},
+			{Term: 9},
 		},
 	}
 	outboundMsgs := []raftpb.Message{
@@ -1822,7 +1819,7 @@ func TestRaftEventFromMsgStorageAppendAndMsgAppsBasic(t *testing.T) {
 	checkSnapAndMap(event)
 	// Only LogSnapshot and ReplicasStateInfo set.
 	event = RaftEventFromMsgStorageAppendAndMsgApps(
-		MsgAppPush, 20, raftpb.Message{}, nil, logSnap, msgAppScratch, infoMap)
+		MsgAppPush, 20, raft.StorageAppend{}, nil, logSnap, msgAppScratch, infoMap)
 	checkSnapAndMap(event)
 	event.LogSnapshot = raft.LogSnapshot{}
 	event.ReplicasStateInfo = nil

pkg/kv/kvserver/logstore/BUILD.bazel

@@ -61,6 +61,7 @@ go_test(
         "//pkg/kv/kvserver/kvserverpb",
         "//pkg/kv/kvserver/raftentry",
         "//pkg/kv/kvserver/raftlog",
+        "//pkg/raft",
         "//pkg/raft/raftpb",
         "//pkg/roachpb",
         "//pkg/settings/cluster",

pkg/kv/kvserver/logstore/logstore.go

@@ -8,7 +8,6 @@ package logstore
 
 import (
 	"context"
-	"fmt"
 	"math"
 	"math/rand"
 	"slices"
@@ -78,70 +77,6 @@ var enableNonBlockingRaftLogSync = settings.RegisterBoolSetting(
 var raftLogTruncationClearRangeThreshold = kvpb.RaftIndex(metamorphic.ConstantWithTestRange(
 	"raft-log-truncation-clearrange-threshold", 100000 /* default */, 1 /* min */, 1e6 /* max */))
 
-// MsgStorageAppend is a raftpb.Message with type MsgStorageAppend.
-type MsgStorageAppend raftpb.Message
-
-// MakeMsgStorageAppend constructs a MsgStorageAppend from a raftpb.Message.
-func MakeMsgStorageAppend(m raftpb.Message) MsgStorageAppend {
-	if m.Type != raftpb.MsgStorageAppend {
-		panic(fmt.Sprintf("unexpected message type %s", m.Type))
-	}
-	return MsgStorageAppend(m)
-}
-
-// HardState returns the hard state assembled from the message.
-func (m *MsgStorageAppend) HardState() raftpb.HardState {
-	return raftpb.HardState{
-		Term:      m.Term,
-		Vote:      m.Vote,
-		Commit:    m.Commit,
-		Lead:      m.Lead,
-		LeadEpoch: m.LeadEpoch,
-	}
-}
-
-// MustSync returns true if this storage write must be synced.
-func (m *MsgStorageAppend) MustSync() bool {
-	return len(m.Responses) != 0
-}
-
-// OnDone returns the storage write post-processing information.
-func (m *MsgStorageAppend) OnDone() MsgStorageAppendDone { return m.Responses }
-
-// MsgStorageAppendDone encapsulates the actions to do after MsgStorageAppend is
-// done, such as sending messages back to raft node and its peers.
-type MsgStorageAppendDone []raftpb.Message
-
-// Responses returns the messages to send after the write/sync is completed.
-func (m MsgStorageAppendDone) Responses() []raftpb.Message { return m }
-
-// Mark returns the LogMark of the raft log in storage after the write/sync is
-// completed. Returns zero value if the write does not update the log mark.
-func (m MsgStorageAppendDone) Mark() raft.LogMark {
-	if len(m) == 0 {
-		return raft.LogMark{}
-	}
-	// Optimization: the MsgStorageAppendResp message, if any, is always the last
-	// one in the list.
-	// TODO(pav-kv): this is an undocumented API quirk. Refactor the raft write
-	// API to be more digestible outside the package.
-	if buildutil.CrdbTestBuild {
-		for _, msg := range m[:len(m)-1] {
-			if msg.Type == raftpb.MsgStorageAppendResp {
-				panic("unexpected MsgStorageAppendResp not in last position")
-			}
-		}
-	}
-	if msg := m[len(m)-1]; msg.Type != raftpb.MsgStorageAppendResp {
-		return raft.LogMark{}
-	} else if msg.Index != 0 {
-		return raft.LogMark{Term: msg.LogTerm, Index: msg.Index}
-	} else if msg.Snapshot != nil {
-		return raft.LogMark{Term: msg.LogTerm, Index: msg.Snapshot.Metadata.Index}
-	}
-	return raft.LogMark{}
-}
-
 // RaftState stores information about the last entry and the size of the log.
 type RaftState struct {
 	LastIndex kvpb.RaftIndex
@@ -198,7 +133,7 @@ type LogStore struct {
 //
 // commitStats is populated iff this was a non-blocking sync.
 type SyncCallback interface {
-	OnLogSync(context.Context, MsgStorageAppendDone, storage.BatchCommitStats)
+	OnLogSync(context.Context, raft.StorageAppendAck, storage.BatchCommitStats)
 }
 
 func newStoreEntriesBatch(eng storage.Engine) storage.Batch {
@@ -218,18 +153,18 @@ func newStoreEntriesBatch(eng storage.Engine) storage.Batch {
 // Accepts the state of the log before the operation, returns the state after.
 // Persists HardState atomically with, or strictly after Entries.
 func (s *LogStore) StoreEntries(
-	ctx context.Context, state RaftState, m MsgStorageAppend, cb SyncCallback, stats *AppendStats,
+	ctx context.Context, state RaftState, app raft.StorageAppend, cb SyncCallback, stats *AppendStats,
 ) (RaftState, error) {
 	batch := newStoreEntriesBatch(s.Engine)
-	return s.storeEntriesAndCommitBatch(ctx, state, m, cb, stats, batch)
+	return s.storeEntriesAndCommitBatch(ctx, state, app, cb, stats, batch)
 }
 
 // storeEntriesAndCommitBatch is like StoreEntries, but it accepts a
 // storage.Batch, which it takes responsibility for committing and closing.
 func (s *LogStore) storeEntriesAndCommitBatch(
 	ctx context.Context,
 	state RaftState,
-	m MsgStorageAppend,
+	m raft.StorageAppend,
 	cb SyncCallback,
 	stats *AppendStats,
 	batch storage.Batch,
@@ -270,7 +205,7 @@ func (s *LogStore) storeEntriesAndCommitBatch(
 		stats.End = crtime.NowMono()
 	}
 
-	if hs := m.HardState(); !raft.IsEmptyHardState(hs) {
+	if hs := m.HardState; !raft.IsEmptyHardState(hs) {
 		// NB: Note that without additional safeguards, it's incorrect to write
 		// the HardState before appending m.Entries. When catching up, a follower
 		// will receive Entries that are immediately Committed in the same
@@ -336,7 +271,7 @@ func (s *LogStore) storeEntriesAndCommitBatch(
 		*waiterCallback = nonBlockingSyncWaiterCallback{
 			ctx:            ctx,
 			cb:             cb,
-			onDone:         m.OnDone(),
+			onDone:         m.Ack(),
 			batch:          batch,
 			metrics:        s.Metrics,
 			logCommitBegin: stats.PebbleBegin,
@@ -354,7 +289,7 @@ func (s *LogStore) storeEntriesAndCommitBatch(
 		if wantsSync {
 			logCommitEnd := stats.PebbleEnd
 			s.Metrics.RaftLogCommitLatency.RecordValue(logCommitEnd.Sub(stats.PebbleBegin).Nanoseconds())
-			cb.OnLogSync(ctx, m.OnDone(), storage.BatchCommitStats{})
+			cb.OnLogSync(ctx, m.Ack(), storage.BatchCommitStats{})
 		}
 	}
 	stats.Sync = wantsSync
@@ -406,7 +341,7 @@ type nonBlockingSyncWaiterCallback struct {
 	// Used to run SyncCallback.
 	ctx    context.Context
 	cb     SyncCallback
-	onDone MsgStorageAppendDone
+	onDone raft.StorageAppendAck
 	// Used to extract stats. This is the batch that has been synced.
 	batch storage.WriteBatch
 	// Used to record Metrics.

pkg/kv/kvserver/logstore/logstore_bench_test.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/raftentry"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/raftlog"
+	"github.com/cockroachdb/cockroach/pkg/raft"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage"
@@ -33,7 +34,10 @@ func (b *discardBatch) Commit(bool) error {
 
 type noopSyncCallback struct{}
 
-func (noopSyncCallback) OnLogSync(context.Context, MsgStorageAppendDone, storage.BatchCommitStats) {}
+func (noopSyncCallback) OnLogSync(
+	context.Context, raft.StorageAppendAck, storage.BatchCommitStats,
+) {
+}
 
 func BenchmarkLogStore_StoreEntries(b *testing.B) {
 	defer log.Scope(b).Close(b)
@@ -96,7 +100,7 @@ func runBenchmarkLogStore_StoreEntries(b *testing.B, bytes int64) {
 	batch := &discardBatch{}
 	for i := 0; i < b.N; i++ {
 		batch.Batch = newStoreEntriesBatch(eng)
-		m := MsgStorageAppend{Entries: ents}
+		m := raft.StorageAppend{Entries: ents}
 		cb := noopSyncCallback{}
 		var err error
 		rs, err = s.storeEntriesAndCommitBatch(ctx, rs, m, cb, stats, batch)

pkg/kv/kvserver/raft.go

@@ -167,8 +167,8 @@ func logRaftReady(ctx context.Context, ready raft.Ready) {
 			i, raft.DescribeEntry(e, raftEntryFormatter))
 	}
 	fmt.Fprintf(&buf, "  Committed: %v\n", ready.Committed)
-	if !raft.IsEmptySnap(ready.Snapshot) {
-		snap := ready.Snapshot
+	if ready.Snapshot != nil {
+		snap := *ready.Snapshot
 		snap.Data = nil
 		fmt.Fprintf(&buf, "  Snapshot updated: %v\n", snap)
 	}

pkg/kv/kvserver/rafttrace/BUILD.bazel

@@ -28,6 +28,7 @@ go_test(
     deps = [
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver/kvserverpb",
+        "//pkg/raft",
         "//pkg/raft/raftpb",
         "//pkg/settings/cluster",
         "//pkg/testutils",