storage: introduce concurrent Raft proposal buffer

pkg/storage/client_raft_test.go

@@ -2171,9 +2171,6 @@ func TestQuotaPool(t *testing.T) {
 			if qLen := leaderRepl.QuotaReleaseQueueLen(); qLen != 1 {
 				return errors.Errorf("expected 1 queued quota release, found: %d", qLen)
 			}
-			if cLen := leaderRepl.CommandSizesLen(); cLen != 0 {
-				return errors.Errorf("expected zero-length command sizes map, found %d", cLen)
-			}
 			return nil
 		})
 
@@ -2196,9 +2193,6 @@ func TestQuotaPool(t *testing.T) {
 		if qLen := leaderRepl.QuotaReleaseQueueLen(); qLen != 0 {
 			return errors.Errorf("expected no queued quota releases, found: %d", qLen)
 		}
-		if cLen := leaderRepl.CommandSizesLen(); cLen != 0 {
-			return errors.Errorf("expected zero-length command sizes map, found %d", cLen)
-		}
 		return nil
 	})
 

pkg/storage/helpers_test.go

@@ -32,7 +32,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage/engine"
 	"github.com/cockroachdb/cockroach/pkg/storage/engine/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/storage/rditer"
-	"github.com/cockroachdb/cockroach/pkg/storage/storagebase"
 	"github.com/cockroachdb/cockroach/pkg/storage/storagepb"
 	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -256,7 +255,7 @@ func (r *Replica) GetLastIndex() (uint64, error) {
 func (r *Replica) LastAssignedLeaseIndex() uint64 {
 	r.mu.RLock()
 	defer r.mu.RUnlock()
-	return r.mu.lastAssignedLeaseIndex
+	return r.mu.proposalBuf.LastAssignedLeaseIndexRLocked()
 }
 
 // SetQuotaPool allows the caller to set a replica's quota pool initialized to
@@ -273,7 +272,6 @@ func (r *Replica) InitQuotaPool(quota int64) {
 	}
 	r.mu.proposalQuota = newQuotaPool(quota)
 	r.mu.quotaReleaseQueue = nil
-	r.mu.commandSizes = make(map[storagebase.CmdIDKey]int)
 }
 
 // QuotaAvailable returns the quota available in the replica's quota pool. Only
@@ -296,12 +294,6 @@ func (r *Replica) IsFollowerActive(ctx context.Context, followerID roachpb.Repli
 	return r.mu.lastUpdateTimes.isFollowerActive(ctx, followerID, timeutil.Now())
 }
 
-func (r *Replica) CommandSizesLen() int {
-	r.mu.Lock()
-	defer r.mu.Unlock()
-	return len(r.mu.commandSizes)
-}
-
 // GetTSCacheHighWater returns the high water mark of the replica's timestamp
 // cache.
 func (r *Replica) GetTSCacheHighWater() hlc.Timestamp {

pkg/storage/replica.go

@@ -222,8 +222,6 @@ type Replica struct {
 		mergeComplete chan struct{}
 		// The state of the Raft state machine.
 		state storagepb.ReplicaState
-		// Counter used for assigning lease indexes for proposals.
-		lastAssignedLeaseIndex uint64
 		// Last index/term persisted to the raft log (not necessarily
 		// committed). Note that lastTerm may be 0 (and thus invalid) even when
 		// lastIndex is known, in which case the term will have to be retrieved
@@ -282,6 +280,12 @@ type Replica struct {
 		minLeaseProposedTS hlc.Timestamp
 		// A pointer to the zone config for this replica.
 		zone *config.ZoneConfig
+		// proposalBuf buffers Raft commands as they are passed to the Raft
+		// replication subsystem. The buffer is populated by requests after
+		// evaluation and is consumed by the Raft processing thread. Once
+		// consumed, commands are proposed through Raft and moved to the
+		// proposals map.
+		proposalBuf propBuf
 		// proposals stores the Raft in-flight commands which originated at
 		// this Replica, i.e. all commands for which propose has been called,
 		// but which have not yet applied.
@@ -381,8 +385,6 @@ type Replica struct {
 		// newly recreated replica will have a complete range descriptor.
 		lastToReplica, lastFromReplica roachpb.ReplicaDescriptor
 
-		// submitProposalFn can be set to mock out the propose operation.
-		submitProposalFn func(*ProposalData) error
 		// Computed checksum at a snapshot UUID.
 		checksums map[uuid.UUID]ReplicaChecksum
 
@@ -396,16 +398,11 @@ type Replica struct {
 
 		proposalQuotaBaseIndex uint64
 
-		// For command size based allocations we keep track of the sizes of all
-		// in-flight commands.
-		commandSizes map[storagebase.CmdIDKey]int
-
-		// Once the leader observes a proposal come 'out of Raft', we consult
-		// the 'commandSizes' map to determine the size of the associated
-		// command and add it to a queue of quotas we have yet to release back
-		// to the quota pool. We only do so when all replicas have persisted
-		// the corresponding entry into their logs.
-		quotaReleaseQueue []int
+		// Once the leader observes a proposal come 'out of Raft', we add the
+		// size of the associated command to a queue of quotas we have yet to
+		// release back to the quota pool. We only do so when all replicas have
+		// persisted the corresponding entry into their logs.
+		quotaReleaseQueue []int64
 
 		// Counts calls to Replica.tick()
 		ticks int
@@ -586,9 +583,8 @@ func (r *Replica) cleanupFailedProposalLocked(p *ProposalData) {
 	// NB: We may be double free-ing here in cases where proposals are
 	// duplicated. To counter this our quota pool is capped at the initial
 	// quota size.
-	if cmdSize, ok := r.mu.commandSizes[p.idKey]; ok {
-		r.mu.proposalQuota.add(int64(cmdSize))
-		delete(r.mu.commandSizes, p.idKey)
+	if r.mu.proposalQuota != nil {
+		r.mu.proposalQuota.add(p.quotaSize)
 	}
 }
 

pkg/storage/replica_closedts.go

@@ -20,13 +20,13 @@ import (
 // closed timestamp tracker. This is called to emit an update about this
 // replica in the absence of write activity.
 func (r *Replica) EmitMLAI() {
-	r.mu.Lock()
-	lai := r.mu.lastAssignedLeaseIndex
+	r.mu.RLock()
+	lai := r.mu.proposalBuf.LastAssignedLeaseIndexRLocked()
 	if r.mu.state.LeaseAppliedIndex > lai {
 		lai = r.mu.state.LeaseAppliedIndex
 	}
 	epoch := r.mu.state.Lease.Epoch
-	r.mu.Unlock()
+	r.mu.RUnlock()
 
 	ctx := r.AnnotateCtx(context.Background())
 	_, untrack := r.store.cfg.ClosedTimestamp.Tracker.Track(ctx)

pkg/storage/replica_destroy.go

@@ -153,6 +153,7 @@ func (r *Replica) destroyRaftMuLocked(ctx context.Context, nextReplicaID roachpb
 
 func (r *Replica) cancelPendingCommandsLocked() {
 	r.mu.AssertHeld()
+	r.mu.proposalBuf.FlushLockedWithoutProposing()
 	for _, p := range r.mu.proposals {
 		r.cleanupFailedProposalLocked(p)
 		// NB: each proposal needs its own version of the error (i.e. don't try to

pkg/storage/replica_init.go

@@ -104,6 +104,7 @@ func (r *Replica) initRaftMuLockedReplicaMuLocked(
 	// reloading the raft state below, it isn't safe to use the existing raft
 	// group.
 	r.mu.internalRaftGroup = nil
+	r.mu.proposalBuf.Init((*replicaProposer)(r))
 
 	var err error
 	if r.mu.state, err = r.mu.stateLoader.Load(ctx, r.store.Engine(), desc); err != nil {

pkg/storage/replica_proposal.go

@@ -71,6 +71,18 @@ type ProposalData struct {
 	// reproposals its MaxLeaseIndex field is mutated.
 	command *storagepb.RaftCommand
 
+	// encodedCommand is the encoded Raft command, with an optional prefix
+	// containing the command ID.
+	encodedCommand []byte
+
+	// quotaSize is the encoded size of command that was used to acquire
+	// proposal quota. command.Size can change slightly as the object is
+	// mutated, so it's safer to record the exact value used here.
+	quotaSize int64
+
+	// tmpFooter is used to avoid an allocation.
+	tmpFooter storagepb.RaftCommandFooter
+
 	// endCmds.finish is called after command execution to update the
 	// timestamp cache & release latches.
 	endCmds *endCmds