Merge #143615

143615: raft: replace MsgStorageApply protocol r=tbg a=pav-kv

Raft interaction with storage mainly consists of two protocols: `MsgStorageAppend` and `MsgStorageApply`. This PR focuses on the latter. The protocol is replaced with a more ergonomic / strongly-typed API, which also provides the user with a better / direct control of the apply flow.

---

The old flow:

1. If there are unapplied committed entries, `RawNode.Ready()` constructs a `MsgStorageApply` including some/all of these entries.
2. The `MsgStorageApply` is included in `Ready.Messages`.
3. The application filters it out from other messages, and acts on it.
4. When the job is done, the `RawNode` must be notified. The `MsgStorageApply` contains a `MsgStorageApplyResp` in its `Responses` field, which must be stepped back to `RawNode`.

Downsides of this approach:

- The decision of fetching entries from storage and constructing `MsgStorageApply` is made by raft. This necessitated some form of built-in flow control, but we never ended up using it: #143576. It would be better to have ability to control this flow from outside raft.
- `MsgStorageApply` construction and the associated IO happens under `Replica.{raftMu+mu}` during the `Ready()` call. We would like to avoid IO when holding `Replica.mu`: #140235.
- The raft messaging system is slightly abused with this approach. All raft messages contain the extra `Responses` [field](https://github.com/cockroachdb/cockroach/blob/acb74317523b0a0849d827a968167a9243e2bb88/pkg/raft/raftpb/raft.proto#L109-L112) that only the two storage APIs use. Most raft messages are external, and don't require order/delivery, while the `MsgStorageAppend/Apply` protocols require strict order.
- Storage interaction being encoded in a `raftpb.Message` with many unused fields is simply confusing and not a great API. The preference is to have dedicated types with stronger/narrower semantics.

---

The downsides are addressed by the new flow:

1. `RawNode.Ready()` includes a `Ready.Committed LogSpan`, to signify the committed but not applied span of the log.
2. The caller obtains `RawNode.LogSnapshot()` and has the freedom of fetching/applying the entire committed span or a prefix. It also can do so after releasing `Replica.mu`.
3. When applied, the caller notifies using the `RawNode.AckApplied()` method.

---

Part of #143652
Related to #124440
Epic: CRDB-46488
Release note: none

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

Author: craig[bot]

pkg/kv/kvpb/BUILD.bazel

@@ -31,6 +31,7 @@ go_library(
         "//pkg/col/coldata",
         "//pkg/kv/kvnemesis/kvnemesisutil",
         "//pkg/kv/kvserver/concurrency/lock",
+        "//pkg/raft/raftpb",
         "//pkg/roachpb",
         "//pkg/storage/enginepb",
         "//pkg/util/admission/admissionpb",

pkg/kv/kvpb/data.go

@@ -7,6 +7,7 @@
 package kvpb
 
 import (
+	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -183,36 +184,6 @@ type LeaseAppliedIndex uint64
 // SafeValue implements the redact.SafeValue interface.
 func (s LeaseAppliedIndex) SafeValue() {}
 
-// RaftTerm represents the term of a raft message. This corresponds to Term in
-// HardState.Term in the Raft library. That type is a uint64, so it is necessary
-// to cast to/from that type when dealing with the Raft library, however
-// internally RaftTerm is used for all fields in CRDB.
-type RaftTerm uint64
-
-// SafeValue implements the redact.SafeValue interface.
-func (s RaftTerm) SafeValue() {}
-
-// RaftIndex represents the term of a raft message. This corresponds to Index in
-// HardState.Index in the Raft library. That type is a uint64, so it is
-// necessary to cast to/from that type when dealing with the Raft library,
-// however internally RaftIndex is used for all fields in CRDB.
-type RaftIndex uint64
-
-// SafeValue implements the redact.SafeValue interface.
-func (s RaftIndex) SafeValue() {}
-
-// RaftSpan represents a (begin, end] span of indices in a raft log. The choice
-// of excluding the left bound and including the right bound is deliberate and
-// principled. When working with raft logs, it almost always helps to avoid
-// off-by-one errors and risk of integer underflow.
-type RaftSpan struct {
-	// After is the left bound of the log indices span. Exclusive.
-	After RaftIndex
-	// Last is the right bound of the log indices span. Inclusive.
-	Last RaftIndex
-}
-
-// Contains returns true iff the given index is within the span.
-func (s RaftSpan) Contains(index RaftIndex) bool {
-	return index > s.After && index <= s.Last
-}
+type RaftTerm = raftpb.Term
+type RaftIndex = raftpb.Index
+type RaftSpan = raftpb.LogSpan

pkg/kv/kvserver/raft.go

@@ -166,10 +166,7 @@ func logRaftReady(ctx context.Context, ready raft.Ready) {
 		fmt.Fprintf(&buf, "  New Entry[%d]: %.200s\n",
 			i, raft.DescribeEntry(e, raftEntryFormatter))
 	}
-	for i, e := range ready.CommittedEntries {
-		fmt.Fprintf(&buf, "  Committed Entry[%d]: %.200s\n",
-			i, raft.DescribeEntry(e, raftEntryFormatter))
-	}
+	fmt.Fprintf(&buf, "  Committed: %v\n", ready.Committed)
 	if !raft.IsEmptySnap(ready.Snapshot) {
 		snap := ready.Snapshot
 		snap.Data = nil

pkg/kv/kvserver/replica_raft.go

@@ -947,7 +947,8 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 
 	var hasReady bool
 	var outboundMsgs []raftpb.Message
-	var msgStorageAppend, msgStorageApply raftpb.Message
+	var msgStorageAppend raftpb.Message
+	var toApply []raftpb.Entry
 	rac2ModeToUse := r.replicationAdmissionControlModeToUse(ctx)
 	// Replication AC v2 state that is initialized while holding Replica.mu.
 	replicaStateInfoMap := r.raftMu.replicaStateScratchForFlowControl
@@ -977,31 +978,37 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 			}
 			r.mu.currentRACv2Mode = rac2ModeToUse
 		}
+		logSnapshot = raftGroup.LogSnapshot()
 		if hasReady = raftGroup.HasReady(); hasReady {
-			// Since we are holding raftMu, only this Ready() call will use
+			// Since we are holding raftMu, only the Slice() call below will use
 			// raftMu.bytesAccount. It tracks memory usage that this Ready incurs.
 			r.attachRaftEntriesMonitorRaftMuLocked()
-			// TODO(pav-kv): currently, Ready() only accounts for entry bytes loaded
+			// TODO(pav-kv): currently, Slice() only accounts for entry bytes loaded
 			// from log storage, and ignores the in-memory unstable entries. Pass a
 			// flow control struct down the stack, and do a more complete accounting
 			// in raft. This will also eliminate the "side channel" plumbing hack with
 			// this bytesAccount.
-			syncRd := raftGroup.Ready()
+			rd := raftGroup.Ready()
+			if !rd.Committed.Empty() {
+				// TODO(pav-kv): do this loading when Replica.mu is released. We don't
+				// want IO under Replica.mu.
+				if toApply, err = logSnapshot.Slice(
+					rd.Committed, r.store.cfg.RaftMaxCommittedSizePerReady,
+				); err != nil {
+					return false, err
+				}
+			}
 			// We apply committed entries during this handleRaftReady, so it is ok to
 			// release the corresponding memory tokens at the end of this func. Next
 			// time we enter this function, the account will be empty again.
 			defer r.detachRaftEntriesMonitorRaftMuLocked()
 
-			logRaftReady(ctx, syncRd)
-			asyncRd := makeAsyncReady(syncRd)
-			outboundMsgs, msgStorageAppend, msgStorageApply = splitLocalStorageMsgs(asyncRd.Messages)
+			logRaftReady(ctx, rd)
+			outboundMsgs, msgStorageAppend = splitLocalStorageMsgs(rd.Messages)
 		}
 		if switchToPullModeAfterReady {
 			raftGroup.SetLazyReplication(true)
 		}
-		if rac2ModeForReady == rac2.MsgAppPull {
-			logSnapshot = raftGroup.LogSnapshot()
-		}
 		raftNodeBasicState = replica_rac2.MakeRaftNodeBasicStateLocked(
 			raftGroup, r.shMu.state.Lease.Replica.ReplicaID)
 		replica_rac2.MakeReplicaStateInfos(raftGroup, replicaStateInfoMap)
@@ -1020,7 +1027,7 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 		unquiesceAndWakeLeader := hasReady || numFlushed > 0 || len(r.mu.proposals) > 0
 		return unquiesceAndWakeLeader, nil
 	})
-	r.mu.applyingEntries = hasMsg(msgStorageApply)
+	r.mu.applyingEntries = len(toApply) != 0
 	pausedFollowers := r.mu.pausedFollowers
 	r.mu.Unlock()
 	if errors.Is(err, errRemoved) {
@@ -1077,9 +1084,9 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 	// entries and acknowledge as many as we can trivially prove will not be
 	// rejected beneath raft.
 	//
-	// Note that the Entries slice in the MsgStorageApply cannot refer to entries
-	// that are also in the Entries slice in the MsgStorageAppend. Raft will not
-	// allow unstable entries to be applied.
+	// Note that the Ready.Committed span cannot refer to entries that are also in
+	// the Entries slice in the MsgStorageAppend. Raft will not allow unstable
+	// entries to be applied.
 	// TODO(pav-kv): Reconsider if this can be relaxed.
 	//
 	// If we disable async storage writes in the future, this property will no
@@ -1097,12 +1104,12 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 	sm := r.getStateMachine()
 	dec := r.getDecoder()
 	var appTask apply.Task
-	if hasMsg(msgStorageApply) {
-		r.mu.raftTracer.MaybeTraceApplying(msgStorageApply.Entries)
+	if len(toApply) != 0 {
+		r.mu.raftTracer.MaybeTraceApplying(toApply)
 		appTask = apply.MakeTask(sm, dec)
 		appTask.SetMaxBatchSize(r.store.TestingKnobs().MaxApplicationBatchSize)
 		defer appTask.Close()
-		if err := appTask.Decode(ctx, msgStorageApply.Entries); err != nil {
+		if err := appTask.Decode(ctx, toApply); err != nil {
 			return stats, err
 		}
 		if knobs := r.store.TestingKnobs(); knobs == nil || !knobs.DisableCanAckBeforeApplication {
@@ -1265,8 +1272,8 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 	}
 
 	stats.tApplicationBegin = crtime.NowMono()
-	if hasMsg(msgStorageApply) {
-		r.traceEntries(msgStorageApply.Entries, "committed, before applying any entries")
+	if len(toApply) != 0 {
+		r.traceEntries(toApply, "committed, before applying any entries")
 
 		err := appTask.ApplyCommittedEntries(ctx)
 		stats.apply = sm.moveStats()
@@ -1302,15 +1309,12 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 				refreshReason = reasonNewLeaderOrConfigChange
 			}
 		}
-
-		r.mu.raftTracer.MaybeTraceApplied(msgStorageApply.Entries)
-		// Send MsgStorageApply's responses.
-		r.sendRaftMessages(ctx, msgStorageApply.Responses, nil /* blocked */, true /* willDeliverLocal */)
+		r.mu.raftTracer.MaybeTraceApplied(toApply)
 	}
 	stats.tApplicationEnd = crtime.NowMono()
 	applicationElapsed := stats.tApplicationEnd.Sub(stats.tApplicationBegin).Nanoseconds()
 	r.store.metrics.RaftApplyCommittedLatency.RecordValue(applicationElapsed)
-	r.store.metrics.RaftCommandsApplied.Inc(int64(len(msgStorageApply.Entries)))
+	r.store.metrics.RaftCommandsApplied.Inc(int64(len(toApply)))
 	if r.store.TestingKnobs().EnableUnconditionalRefreshesInRaftReady {
 		refreshReason = reasonNewLeaderOrConfigChange
 	}
@@ -1328,6 +1332,7 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 	r.mu.Lock()
 	err = r.withRaftGroupLocked(func(raftGroup *raft.RawNode) (bool, error) {
 		r.deliverLocalRaftMsgsRaftMuLockedReplicaMuLocked(ctx, raftGroup)
+		raftGroup.AckApplied(toApply)
 
 		if stats.apply.numConfChangeEntries > 0 {
 			// If the raft leader got removed, campaign on the leaseholder. Uses
@@ -1381,26 +1386,6 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 	return stats, nil
 }
 
-// asyncReady encapsulates the messages that are ready to be sent to other peers
-// or to be sent to local storage routines when async storage writes are enabled.
-// All fields in asyncReady are read-only.
-// TODO(nvanbenschoten): move this into go.etcd.io/raft.
-type asyncReady struct {
-	// Messages specifies outbound messages to other peers and to local storage
-	// threads. These messages can be sent in any order.
-	//
-	// If it contains a MsgSnap message, the application MUST report back to raft
-	// when the snapshot has been received or has failed by calling ReportSnapshot.
-	Messages []raftpb.Message
-}
-
-// makeAsyncReady constructs an asyncReady from the provided Ready.
-func makeAsyncReady(rd raft.Ready) asyncReady {
-	return asyncReady{
-		Messages: rd.Messages,
-	}
-}
-
 // hasMsg returns whether the provided raftpb.Message is present.
 // It serves as a poor man's Optional[raftpb.Message].
 func hasMsg(m raftpb.Message) bool { return m.Type != 0 }
@@ -1409,30 +1394,25 @@ func hasMsg(m raftpb.Message) bool { return m.Type != 0 }
 // message slice and returns them separately.
 func splitLocalStorageMsgs(
 	msgs []raftpb.Message,
-) (otherMsgs []raftpb.Message, msgStorageAppend, msgStorageApply raftpb.Message) {
+) (otherMsgs []raftpb.Message, msgStorageAppend raftpb.Message) {
 	for i := len(msgs) - 1; i >= 0; i-- {
 		switch msgs[i].Type {
 		case raftpb.MsgStorageAppend:
 			if hasMsg(msgStorageAppend) {
 				panic("two MsgStorageAppend")
 			}
 			msgStorageAppend = msgs[i]
-		case raftpb.MsgStorageApply:
-			if hasMsg(msgStorageApply) {
-				panic("two MsgStorageApply")
-			}
-			msgStorageApply = msgs[i]
 		default:
 			// Local storage messages will always be at the end of the messages slice,
 			// so we can terminate iteration as soon as we reach any other message
-			// type. This is leaking an implementation detail from etcd/raft which may
+			// type. This is leaking an implementation detail from pkg/raft which may
 			// not always hold, but while it does, we use it for convenience and
 			// assert against it changing in sendRaftMessages.
-			return msgs[:i+1], msgStorageAppend, msgStorageApply
+			return msgs[:i+1], msgStorageAppend
 		}
 	}
 	// Only local storage messages.
-	return nil, msgStorageAppend, msgStorageApply
+	return nil, msgStorageAppend
 }
 
 // maybeFatalOnRaftReadyErr will fatal if err is neither nil nor
@@ -1899,12 +1879,6 @@ func (r *Replica) sendRaftMessages(
 			// Instead, we handle messages to LocalAppendThread inline on the raft
 			// scheduler goroutine, so this code path is unused.
 			panic("unsupported, currently processed inline on raft scheduler goroutine")
-		case raft.LocalApplyThread:
-			// To local apply thread.
-			// NOTE: we don't currently split apply work off into an async goroutine.
-			// Instead, we handle messages to LocalAppendThread inline on the raft
-			// scheduler goroutine, so this code path is unused.
-			panic("unsupported, currently processed inline on raft scheduler goroutine")
 		case raftpb.PeerID(r.ReplicaID()):
 			// To local raft state machine, from local storage append and apply work.
 			// NOTE: For async Raft log appends, these messages come from calls to

pkg/raft/log.go

@@ -29,12 +29,6 @@ import (
 //
 // To access it safely, the user must not mutate the underlying raft log storage
 // between when the snapshot is obtained and the reads are done.
-//
-// TODO(pav-kv): this should be part of the Ready API. Instead of pre-fetching
-// entries (e.g. the committed entries subject to state machine application),
-// allow the application to read them from LogSnapshot in the Ready handler.
-// This gives the application direct control on resource allocation, and
-// flexibility to do raft log IO without blocking RawNode operation.
 type LogSnapshot struct {
 	// compacted is the compacted log index.
 	compacted uint64
@@ -291,23 +285,27 @@ func (l *raftLog) hasNextUnstableEnts() bool {
 // appended them to the local raft log yet. If allowUnstable is true, committed
 // entries from the unstable log may be returned; otherwise, only entries known
 // to reside locally on stable storage will be returned.
+//
+// TODO(pav-kv): only used in tests. Downgrade to a test helper or remove.
 func (l *raftLog) nextCommittedEnts(allowUnstable bool) (ents []pb.Entry) {
-	lo, hi := l.applying, l.maxAppliableIndex(allowUnstable) // (lo, hi]
-	if lo >= hi {
-		// Nothing to apply.
+	span := l.nextCommittedSpan(allowUnstable)
+	if span.Empty() {
 		return nil
 	}
-	ents, err := l.slice(lo, hi, l.maxApplyingEntsSize)
+	ents, err := l.slice(uint64(span.After), uint64(span.Last), l.maxApplyingEntsSize)
 	if err != nil {
 		l.logger.Panicf("unexpected error when getting unapplied entries (%v)", err)
 	}
 	return ents
 }
 
-// hasNextCommittedEnts returns if there is any available entries for execution.
+// nextCommittedSpan returns the span of committed entries that can be applied.
 // This is a fast check without heavy raftLog.slice() in nextCommittedEnts().
-func (l *raftLog) hasNextCommittedEnts(allowUnstable bool) bool {
-	return l.applying < l.maxAppliableIndex(allowUnstable)
+func (l *raftLog) nextCommittedSpan(allowUnstable bool) pb.LogSpan {
+	return pb.LogSpan{
+		After: pb.Index(l.applying),
+		Last:  pb.Index(l.maxAppliableIndex(allowUnstable)),
+	}
 }
 
 // maxAppliableIndex returns the maximum committed index that can be applied.
@@ -577,7 +575,15 @@ func (l LogSnapshot) LeadSlice(lo, hi uint64, maxSize uint64) (LeadSlice, error)
 	}, nil
 }
 
-// TODO(pav-kv): return LogSlice.
+// Slice returns log entries forming a prefix of the given log span, with the
+// total entries size not exceeding maxSize.
+//
+// Returns at least one entry if the interval contains any. The maxSize can only
+// be exceeded if the first entry (span.After+1) is larger.
+func (l LogSnapshot) Slice(span pb.LogSpan, maxSize uint64) ([]pb.Entry, error) {
+	return l.slice(uint64(span.After), uint64(span.Last), entryEncodingSize(maxSize))
+}
+
 func (l LogSnapshot) slice(lo, hi uint64, maxSize entryEncodingSize) ([]pb.Entry, error) {
 	if err := l.mustCheckOutOfBounds(lo, hi); err != nil {
 		return nil, err