Batch apply state cleanup

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

Author: MauriceVanVeen

server/jetstream_batching_test.go

@@ -336,17 +336,21 @@ func TestJetStreamAtomicBatchPublishCleanup(t *testing.T) {
 		require_NoError(t, err)
 		mset.mu.RLock()
 		batches := mset.batches
+		batch := mset.batchApply
 		mset.mu.RUnlock()
 		require_True(t, batches == nil)
+		require_True(t, batch == nil)
 
 		// Enabling doesn't need to populate the batching state.
 		cfg.AllowAtomicPublish = true
 		_, err = jsStreamUpdate(t, nc, cfg)
 		require_NoError(t, err)
 		mset.mu.RLock()
 		batches = mset.batches
+		batch = mset.batchApply
 		mset.mu.RUnlock()
 		require_True(t, batches == nil)
+		require_True(t, batch == nil)
 
 		// Publish a partial batch that needs to be cleaned up.
 		m := nats.NewMsg("foo")
@@ -362,8 +366,10 @@ func TestJetStreamAtomicBatchPublishCleanup(t *testing.T) {
 
 		mset.mu.RLock()
 		batches = mset.batches
+		batch = mset.batchApply
 		mset.mu.RUnlock()
 		require_NotNil(t, batches)
+		require_NotNil(t, batch)
 		batches.mu.Lock()
 		groups := len(batches.group)
 		b := batches.group["uuid"]
@@ -372,6 +378,7 @@ func TestJetStreamAtomicBatchPublishCleanup(t *testing.T) {
 		require_NotNil(t, b)
 		store := b.store
 		require_Equal(t, store.State().Msgs, 1)
+		clfs := mset.getCLFS()
 
 		// Should fully clean up the in-progress batch.
 		switch mode {
@@ -396,6 +403,15 @@ func TestJetStreamAtomicBatchPublishCleanup(t *testing.T) {
 			}
 			return nil
 		})
+		// Should clean up the batch apply state.
+		if mode == Disable || mode == Delete {
+			mset.mu.RLock()
+			batch = mset.batchApply
+			mset.mu.RUnlock()
+			nclfs := mset.getCLFS()
+			require_True(t, batch == nil)
+			require_Equal(t, clfs, nclfs)
+		}
 	}
 
 	t.Run("Disable", func(t *testing.T) { test(t, Disable) })

server/stream.go

@@ -2228,6 +2228,7 @@ func (mset *stream) updateWithAdvisory(config *StreamConfig, sendAdvisory bool,
 	// If atomic publish is disabled, delete any in-progress batches.
 	if !cfg.AllowAtomicPublish {
 		mset.deleteInflightBatches()
+		mset.deleteBatchApplyState()
 	}
 
 	// Now update config and store's version of our config.
@@ -4161,8 +4162,10 @@ func (mset *stream) unsubscribeToStream(stopping bool) error {
 	// Clear batching state.
 	mset.deleteInflightBatches()
 
-	// In case we had a direct get subscriptions.
 	if stopping {
+		mset.deleteBatchApplyState()
+
+		// In case we had a direct get subscriptions.
 		mset.unsubscribeToDirect()
 	}
 
@@ -4182,6 +4185,17 @@ func (mset *stream) deleteInflightBatches() {
 	}
 }
 
+// Lock should be held.
+func (mset *stream) deleteBatchApplyState() {
+	if batch := mset.batchApply; batch != nil {
+		// Need to return entries (if any) to the pool.
+		for _, bce := range batch.entries {
+			bce.ReturnToPool()
+		}
+		mset.batchApply = nil
+	}
+}
+
 // Lock does NOT need to be held, we set the client on setup and never change it at this point.
 func (mset *stream) subscribeInternal(subject string, cb msgHandler) (*subscription, error) {
 	if mset.closed.Load() {