(2.12) NRG: separate Processed and Applied for batch-semantics

Signed-off-by: Maurice van Veen <[email protected]>

Author: MauriceVanVeen

locksordering.txt

@@ -35,8 +35,10 @@ The "clsMu" lock protects the consumer list on a stream, used for signalling con
 
     stream -> clsMu
 
-The "clMu" and "ddMu" locks protect clustered and dedupe state respectively. The stream lock (`mset.mu`) is optional,
-but if holding "clMu" or "ddMu", locking the stream lock afterward would violate locking order.
+The "clMu", "ddMu" and "batchMu" locks protect clustered, dedupe and batch state respectively.
+The stream lock (`mset.mu`) is optional, but if holding "clMu", "ddMu" or "batchMu",
+locking the stream lock afterward would violate locking order.
 
     stream -> clMu
+    stream -> batchMu -> clMu
     stream -> ddMu

server/jetstream_batching_test.go

@@ -18,6 +18,7 @@ package server
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"strconv"
 	"testing"
@@ -634,3 +635,104 @@ func TestJetStreamAtomicBatchPublishProposeMultiplePartialBatches(t *testing.T)
 		})
 	}
 }
+
+// Test a continuous flow of batches spanning multiple append entries can still move applied up.
+// Also, test a server can become leader if the previous leader left it with a partial batch.
+func TestJetStreamAtomicBatchPublishContinuousBatchesStillMoveAppliedUp(t *testing.T) {
+	c := createJetStreamClusterExplicit(t, "R3S", 3)
+	defer c.shutdown()
+
+	nc, js := jsClientConnect(t, c.randomServer())
+	defer nc.Close()
+
+	_, err := jsStreamCreate(t, nc, &StreamConfig{
+		Name:               "TEST",
+		Subjects:           []string{"foo"},
+		Storage:            FileStorage,
+		Replicas:           3,
+		AllowAtomicPublish: true,
+	})
+	require_NoError(t, err)
+
+	sl := c.streamLeader(globalAccountName, "TEST")
+	mset, err := sl.globalAccount().lookupStream("TEST")
+	require_NoError(t, err)
+
+	pubAck, err := js.Publish("foo", nil)
+	require_NoError(t, err)
+	require_Equal(t, pubAck.Sequence, 1)
+
+	n := mset.raftNode().(*raft)
+	index, commit, applied := n.Progress()
+
+	// The first batch spans two append entries, but commits.
+	mset.clMu.Lock()
+	hdr := genHeader(nil, "Nats-Batch-Id", "ID_1")
+	hdr = setHeader("Nats-Batch-Sequence", "1", hdr)
+	esm := encodeStreamMsgAllowCompressAndBatch("foo", _EMPTY_, hdr, nil, mset.clseq, 0, false, "ID_1", 1, false)
+	mset.clseq++
+	mset.clMu.Unlock()
+	n.sendAppendEntry([]*Entry{newEntry(EntryNormal, esm)})
+
+	var entries []*Entry
+	mset.clMu.Lock()
+	hdr = genHeader(nil, "Nats-Batch-Id", "ID_1")
+	hdr = setHeader("Nats-Batch-Sequence", "2", hdr)
+	hdr = setHeader("Nats-Batch-Commit", "1", hdr)
+	esm = encodeStreamMsgAllowCompressAndBatch("foo", _EMPTY_, hdr, nil, mset.clseq, 0, false, "ID_1", 2, true)
+	mset.clseq++
+
+	// The second batch doesn't commit.
+	entries = append(entries, newEntry(EntryNormal, esm))
+	hdr = genHeader(nil, "Nats-Batch-Id", "ID_2")
+	hdr = setHeader("Nats-Batch-Sequence", "1", hdr)
+	esm = encodeStreamMsgAllowCompressAndBatch("foo", _EMPTY_, hdr, nil, mset.clseq, 0, false, "ID_2", 1, false)
+	mset.clseq++
+	entries = append(entries, newEntry(EntryNormal, esm))
+	mset.clMu.Unlock()
+	n.sendAppendEntry(entries)
+
+	checkFor(t, 2*time.Second, 200*time.Millisecond, func() error {
+		n.RLock()
+		nindex, ncommit, processed, napplied := n.pindex, n.commit, n.processed, n.applied
+		n.RUnlock()
+		if nindex == index {
+			return errors.New("index not updated")
+		} else if ncommit == commit {
+			return errors.New("commit not updated")
+		} else if napplied == applied {
+			return errors.New("applied not updated")
+		} else if napplied == ncommit {
+			return errors.New("applied should not be able to equal commit yet")
+		} else if processed != ncommit {
+			return errors.New("must have processed all commits")
+		}
+		return checkState(t, c, globalAccountName, "TEST")
+	})
+
+	// Followers are now stranded with a partial batch, one needs to become leader
+	// and have the batch be rejected since it was partially proposed.
+	sl.Shutdown()
+	c.waitOnStreamLeader(globalAccountName, "TEST")
+	checkFor(t, 2*time.Second, 200*time.Millisecond, func() error {
+		return checkState(t, c, globalAccountName, "TEST")
+	})
+
+	// Confirm the last batch gets rejected, and we are still able to publish with quorum.
+	pubAck, err = js.Publish("foo", nil)
+	require_NoError(t, err)
+	require_Equal(t, pubAck.Sequence, 4)
+
+	c.restartServer(sl)
+	checkFor(t, 2*time.Second, 200*time.Millisecond, func() error {
+		return checkState(t, c, globalAccountName, "TEST")
+	})
+
+	// Publish again, now with all servers online.
+	pubAck, err = js.Publish("foo", nil)
+	require_NoError(t, err)
+	require_Equal(t, pubAck.Sequence, 5)
+	checkFor(t, 2*time.Second, 200*time.Millisecond, func() error {
+		return checkState(t, c, globalAccountName, "TEST")
+	})
+}

server/jetstream_cluster.go

@@ -2463,9 +2463,13 @@ func (js *jetStream) monitorStream(mset *stream, sa *streamAssignment, sendSnaps
 				}
 
 				// Apply our entries.
-				if applied, err := js.applyStreamEntries(mset, ce, isRecovering); err == nil {
+				if maxApplied, err := js.applyStreamEntries(mset, ce, isRecovering); err == nil {
 					// Update our applied.
-					if applied {
+					if maxApplied > 0 {
+						// Indicate we've processed (but not applied) everything up to this point.
+						ne, nb = n.Processed(ce.Index, min(maxApplied, ce.Index))
+						// Don't return entry to the pool, this is handled by the in-progress batch.
+					} else {
 						ne, nb = n.Applied(ce.Index)
 						ce.ReturnToPool()
 					}
@@ -2928,8 +2932,8 @@ func isControlHdr(hdr []byte) bool {
 }
 
 // Apply our stream entries.
-// Return whether the ce is applied and error.
-func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isRecovering bool) (bool, error) {
+// Return maximum allowed applied value, if currently inside a batch, zero otherwise.
+func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isRecovering bool) (uint64, error) {
 	mset.batchMu.Lock()
 	batchActiveId := mset.batchId
 	mset.batchMu.Unlock()
@@ -2968,6 +2972,7 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 				// At that point we'll commit the whole batch.
 				if batchSeq == 1 {
 					mset.batchEntryStart = i
+					mset.batchMaxApplied = ce.Index - 1
 				}
 				mset.batchId, batchActiveId = batchId, batchId
 				mset.batchMu.Unlock()
@@ -3011,7 +3016,9 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 						}
 						if err = js.applyStreamMsgOp(mset, op, buf, isRecovering); err != nil {
 							mset.batchMu.Unlock()
-							return false, err
+							// Make sure to return previous entries to the pool on error.
+							bce.ReturnToPool()
+							return 0, err
 						}
 					}
 					// Return the entry to the pool now.
@@ -3033,7 +3040,7 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 					}
 					if err = js.applyStreamMsgOp(mset, op, buf, isRecovering); err != nil {
 						mset.batchMu.Unlock()
-						return false, err
+						return 0, err
 					}
 				}
 				// Clear state, batch was successful.
@@ -3053,7 +3060,7 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 			case streamMsgOp, compressedStreamMsgOp:
 				mbuf := buf[1:]
 				if err := js.applyStreamMsgOp(mset, op, mbuf, isRecovering); err != nil {
-					return false, err
+					return 0, err
 				}
 
 			case deleteMsgOp:
@@ -3077,7 +3084,7 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 
 				// Cluster reset error.
 				if err == ErrStoreEOF {
-					return false, err
+					return 0, err
 				}
 
 				if err != nil && !isRecovering {
@@ -3170,13 +3177,13 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 				ss, err = DecodeStreamState(e.Data)
 				if err != nil {
 					onBadState(err)
-					return false, err
+					return 0, err
 				}
 			} else {
 				var snap streamSnapshot
 				if err := json.Unmarshal(e.Data, &snap); err != nil {
 					onBadState(err)
-					return false, err
+					return 0, err
 				}
 				// Convert over to StreamReplicatedState
 				ss = &StreamReplicatedState{
@@ -3193,7 +3200,7 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 
 			if isRecovering || !mset.IsLeader() {
 				if err := mset.processSnapshot(ss, ce.Index); err != nil {
-					return false, err
+					return 0, err
 				}
 			}
 		} else if e.Type == EntryRemovePeer {
@@ -3228,9 +3235,11 @@ func (js *jetStream) applyStreamEntries(mset *stream, ce *CommittedEntry, isReco
 		} else {
 			mset.batchEntries = append(mset.batchEntries, ce)
 		}
+		maxApplied := mset.batchMaxApplied
 		mset.batchMu.Unlock()
+		return maxApplied, nil
 	}
-	return batchActiveId == _EMPTY_, nil
+	return 0, nil
 }
 
 // clearBatchStateLocked clears in-memory apply-batch-related state.
@@ -3240,6 +3249,7 @@ func (mset *stream) clearBatchStateLocked() {
 	mset.batchCount = 0
 	mset.batchEntries = nil
 	mset.batchEntryStart = 0
+	mset.batchMaxApplied = 0
 }
 
 // rejectBatchStateLocked rejects the batch and clears in-memory apply-batch-related state.
@@ -3463,7 +3473,7 @@ func (js *jetStream) processStreamLeaderChange(mset *stream, isLeader bool) {
 		}
 
 		// Clear clseq. If we become leader again, it will be fixed up
-		// automatically on the next processClusteredInboundMsg call.
+		// automatically on the next mset.setLeader call.
 		mset.clMu.Lock()
 		if mset.clseq > 0 {
 			mset.clseq = 0
@@ -8373,14 +8383,7 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 
 		// Proceed with proposing the batch.
 
-		// We only use mset.clseq for clustering and in case we run ahead of actual commits.
-		// Check if we need to set initial value here
 		mset.clMu.Lock()
-		if mset.clseq == 0 || mset.clseq < lseq+mset.clfs {
-			// Re-capture
-			lseq = mset.lastSeq()
-			mset.clseq = lseq + mset.clfs
-		}
 
 		// TODO(mvv): support message tracing
 		ts := time.Now().UnixNano()
@@ -8478,11 +8481,6 @@ func (mset *stream) processClusteredInboundMsg(subject, reply string, hdr, msg [
 	// We only use mset.clseq for clustering and in case we run ahead of actual commits.
 	// Check if we need to set initial value here
 	mset.clMu.Lock()
-	if mset.clseq == 0 || mset.clseq < lseq+mset.clfs {
-		// Re-capture
-		lseq = mset.lastSeq()
-		mset.clseq = lseq + mset.clfs
-	}
 
 	var (
 		dseq   uint64

server/raft.go

@@ -44,6 +44,7 @@ type RaftNode interface {
 	SendSnapshot(snap []byte) error
 	NeedSnapshot() bool
 	Applied(index uint64) (entries uint64, bytes uint64)
+	Processed(index uint64, applied uint64) (entries uint64, bytes uint64)
 	State() RaftState
 	Size() (entries, bytes uint64)
 	Progress() (index, commit, applied uint64)
@@ -164,11 +165,12 @@ type raft struct {
 	llqrt  time.Time   // Last quorum lost time
 	lsut   time.Time   // Last scale-up time
 
-	term    uint64 // The current vote term
-	pterm   uint64 // Previous term from the last snapshot
-	pindex  uint64 // Previous index from the last snapshot
-	commit  uint64 // Index of the most recent commit
-	applied uint64 // Index of the most recently applied commit
+	term      uint64 // The current vote term
+	pterm     uint64 // Previous term from the last snapshot
+	pindex    uint64 // Previous index from the last snapshot
+	commit    uint64 // Index of the most recent commit
+	processed uint64 // Index of the most recently processed commit
+	applied   uint64 // Index of the most recently applied commit
 
 	aflr uint64 // Index when to signal initial messages have been applied after becoming leader. 0 means signaling is disabled.
 
@@ -1086,6 +1088,16 @@ func (n *raft) ResumeApply() {
 // apply queue. It will return the number of entries and an estimation of the
 // byte size that could be removed with a snapshot/compact.
 func (n *raft) Applied(index uint64) (entries uint64, bytes uint64) {
+	return n.Processed(index, index)
+}
+
+// Processed is a callback that must be called by the upper layer when it
+// has processed the committed entries that it received from the apply queue,
+// but it (maybe) hasn't applied all the processed entries yet.
+// Used to indicate a commit was processed, even if it wasn't applied yet and
+// can't be compacted away by a snapshot just yet. Which allows us to try to
+// become leader if we've processed all commits, even if they're not all applied.
+func (n *raft) Processed(index uint64, applied uint64) (entries uint64, bytes uint64) {
 	n.Lock()
 	defer n.Unlock()
 
@@ -1094,13 +1106,22 @@ func (n *raft) Applied(index uint64) (entries uint64, bytes uint64) {
 		return 0, 0
 	}
 
+	// Ignore if already processed.
+	if index > n.processed {
+		n.processed = index
+	}
+
 	// Ignore if already applied.
-	if index > n.applied {
-		n.applied = index
+	if applied > index {
+		applied = index
+	}
+	if applied > n.applied {
+		n.applied = applied
 	}
 
-	// If it was set, and we reached the minimum applied index, reset and send signal to upper layer.
-	if n.aflr > 0 && index >= n.aflr {
+	// If it was set, and we reached the minimum processed index, reset and send signal to upper layer.
+	// We're not waiting for processed AND applied, because applying could take longer.
+	if n.aflr > 0 && n.processed >= n.aflr {
 		n.aflr = 0
 		// Quick sanity-check to confirm we're still leader.
 		// In which case we must signal, since switchToLeader would not have done so already.
@@ -3333,8 +3354,11 @@ func (n *raft) truncateWAL(term, index uint64) {
 		if n.commit > n.pindex {
 			n.commit = n.pindex
 		}
-		if n.applied > n.commit {
-			n.applied = n.commit
+		if n.processed > n.commit {
+			n.processed = n.commit
+		}
+		if n.applied > n.processed {
+			n.applied = n.processed
 		}
 	}()
 
@@ -4447,7 +4471,7 @@ func (n *raft) switchToCandidate() {
 
 	// If we are catching up or are in observer mode we can not switch.
 	// Avoid petitioning to become leader if we're behind on applies.
-	if n.observer || n.paused || n.applied < n.commit {
+	if n.observer || n.paused || n.processed < n.commit {
 		n.resetElect(minElectionTimeout / 4)
 		return
 	}