WIP: finalize

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

Author: MauriceVanVeen

server/jetstream_batching.go

@@ -34,8 +34,9 @@ var (
 )
 
 type batching struct {
-	mu    sync.Mutex
-	group map[string]*batchGroup
+	mu     sync.Mutex
+	atomic map[string]*batchGroup
+	fast   map[string]*batchGroup
 }
 
 type batchGroup struct {
@@ -124,7 +125,7 @@ func (b *batchGroup) readyForCommit() bool {
 // Also returns whether a flow control message should be sent.
 // Lock should be held.
 func (batches *batching) fastBatchRegisterSequences(batchId string, batchSeq, streamSeq uint64) (*batchGroup, bool, string) {
-	if b, ok := batches.group[batchId]; ok {
+	if b, ok := batches.fast[batchId]; ok {
 		b.sseq = streamSeq
 		b.pseq = batchSeq
 		// If the PubAck needs to be sent now as a result of a commit.
@@ -209,8 +210,10 @@ func (b *batchGroup) cleanupLocked(batchId string, batches *batching) {
 	b.timer.Stop()
 	if b.store != nil {
 		b.store.Delete(true)
+		delete(batches.atomic, batchId)
+	} else {
+		delete(batches.fast, batchId)
 	}
-	delete(batches.group, batchId)
 }
 
 // Lock should be held.

server/jetstream_batching_test.go

@@ -132,6 +132,7 @@ func TestJetStreamAtomicBatchPublish(t *testing.T) {
 
 			pubAck = JSPubAckResponse{}
 			require_NoError(t, json.Unmarshal(rmsg.Data, &pubAck))
+			require_True(t, pubAck.Error == nil)
 			require_Equal(t, pubAck.Sequence, 6)
 			require_Equal(t, pubAck.BatchId, "uuid")
 			require_Equal(t, pubAck.BatchSize, 5)
@@ -393,7 +394,7 @@ func TestJetStreamAtomicBatchPublishLimits(t *testing.T) {
 				return errors.New("batches not found")
 			}
 			batches.mu.Lock()
-			groups := len(batches.group)
+			groups := len(batches.atomic)
 			batches.mu.Unlock()
 			if groups != 0 {
 				return fmt.Errorf("expected 0 groups, got %d", groups)
@@ -684,8 +685,8 @@ func TestJetStreamAtomicBatchPublishCleanup(t *testing.T) {
 		require_NotNil(t, batches)
 		require_NotNil(t, batch)
 		batches.mu.Lock()
-		groups := len(batches.group)
-		b := batches.group["uuid"]
+		groups := len(batches.atomic)
+		b := batches.atomic["uuid"]
 		batches.mu.Unlock()
 		require_Len(t, groups, 1)
 		require_NotNil(t, b)
@@ -720,7 +721,7 @@ func TestJetStreamAtomicBatchPublishCleanup(t *testing.T) {
 					return fmt.Errorf("expected no batches")
 				}
 				batches.mu.Lock()
-				groups = len(batches.group)
+				groups = len(batches.atomic)
 				batches.mu.Unlock()
 				if groups > 0 {
 					return fmt.Errorf("expected 0 groups, got %d", groups)
@@ -2893,7 +2894,7 @@ func TestJetStreamAtomicBatchPublishCommitUnsupported(t *testing.T) {
 		batches := mset.batches
 		mset.mu.RUnlock()
 		batches.mu.Lock()
-		groups := len(batches.group)
+		groups := len(batches.atomic)
 		batches.mu.Unlock()
 		require_Len(t, groups, 0)
 	}

server/stream.go

@@ -4548,14 +4548,18 @@ func (mset *stream) unsubscribeToStream(stopping, shuttingDown bool) error {
 func (mset *stream) deleteInflightBatches(shuttingDown bool) {
 	if mset.batches != nil {
 		mset.batches.mu.Lock()
-		for batchId, b := range mset.batches.group {
-			// If shutting down, do fixup during startup. In-memory batches don't require manual cleanup.
-			if shuttingDown {
-				b.stopLocked()
-			} else {
-				b.cleanupLocked(batchId, mset.batches)
+		clearBatches := func(group map[string]*batchGroup) {
+			for batchId, b := range group {
+				// If shutting down, do fixup during startup. In-memory batches don't require manual cleanup.
+				if shuttingDown {
+					b.stopLocked()
+				} else {
+					b.cleanupLocked(batchId, mset.batches)
+				}
 			}
 		}
+		clearBatches(mset.batches.atomic)
+		clearBatches(mset.batches.fast)
 		mset.batches.mu.Unlock()
 		mset.batches = nil
 	}
@@ -6396,9 +6400,9 @@ func (mset *stream) processJetStreamMsg(subject, reply string, hdr, msg []byte,
 	return nil
 }
 
-// processJetStreamBatchMsg processes a JetStream message that's part of an atomic or fast batch publish.
+// processJetStreamAtomicBatchMsg processes a JetStream message that's part of an atomic batch publish.
 // Handles constraints around the batch, storing messages, doing consistency checks, and performing the commit.
-func (mset *stream) processJetStreamBatchMsg(batchId, subject, reply string, hdr, msg []byte, mt *msgTrace) (retErr error) {
+func (mset *stream) processJetStreamAtomicBatchMsg(batchId, subject, reply string, hdr, msg []byte, mt *msgTrace) (retErr error) {
 	mset.mu.RLock()
 	canRespond := !mset.cfg.NoAck && len(reply) > 0
 	name, stype := mset.cfg.Name, mset.cfg.Storage
@@ -6482,9 +6486,7 @@ func (mset *stream) processJetStreamBatchMsg(batchId, subject, reply string, hdr
 
 	mset.mu.Lock()
 	if mset.batches == nil {
-		mset.batches = &batching{
-			group: make(map[string]*batchGroup, 1),
-		}
+		mset.batches = &batching{}
 	}
 	batches := mset.batches
 	mset.mu.Unlock()
@@ -6495,7 +6497,7 @@ func (mset *stream) processJetStreamBatchMsg(batchId, subject, reply string, hdr
 
 	// Get batch.
 	batches.mu.Lock()
-	b, ok := batches.group[batchId]
+	b, ok := batches.atomic[batchId]
 	if !ok {
 		if batchSeq != 1 {
 			batches.mu.Unlock()
@@ -6523,7 +6525,7 @@ func (mset *stream) processJetStreamBatchMsg(batchId, subject, reply string, hdr
 		}
 
 		// Confirm we can facilitate an additional batch.
-		if len(batches.group)+1 > maxInflightPerStream {
+		if len(batches.atomic)+1 > maxInflightPerStream {
 			batches.mu.Unlock()
 			return respondIncompleteBatch()
 		}
@@ -6542,7 +6544,10 @@ func (mset *stream) processJetStreamBatchMsg(batchId, subject, reply string, hdr
 			batches.mu.Unlock()
 			return respondIncompleteBatch()
 		}
-		batches.group[batchId] = b
+		if batches.atomic == nil {
+			batches.atomic = make(map[string]*batchGroup, 1)
+		}
+		batches.atomic[batchId] = b
 	}
 
 	var commit, commitEob bool
@@ -6770,6 +6775,8 @@ func (mset *stream) processJetStreamBatchMsg(batchId, subject, reply string, hdr
 	return nil
 }
 
+// processJetStreamFastBatchMsg processes a JetStream message that's part of an atomic batch publish.
+// Handles constraints around the batch, storing messages, doing consistency checks, and performing the commit.
 func (mset *stream) processJetStreamFastBatchMsg(subject, reply string, hdr, msg []byte, mt *msgTrace) (retErr error) {
 	mset.mu.RLock()
 	canRespond := !mset.cfg.NoAck && len(reply) > 0
@@ -6854,16 +6861,14 @@ func (mset *stream) processJetStreamFastBatchMsg(subject, reply string, hdr, msg
 
 	mset.mu.Lock()
 	if mset.batches == nil {
-		mset.batches = &batching{
-			group: make(map[string]*batchGroup, 1),
-		}
+		mset.batches = &batching{}
 	}
 	batches := mset.batches
 	mset.mu.Unlock()
 
 	// Get batch.
 	batches.mu.Lock()
-	b, ok := batches.group[batch.id]
+	b, ok := batches.fast[batch.id]
 	if !ok {
 		if batch.seq != 1 {
 			batches.mu.Unlock()
@@ -6872,7 +6877,10 @@ func (mset *stream) processJetStreamFastBatchMsg(subject, reply string, hdr, msg
 		// We'll need a copy as we'll use it as a key and later for cleanup.
 		batchId := copyString(batch.id)
 		b = batches.newFastBatchGroup(mset, batchId, batch.gapOk, batch.flow)
-		batches.group[batchId] = b
+		if batches.fast == nil {
+			batches.fast = make(map[string]*batchGroup, 1)
+		}
+		batches.fast[batchId] = b
 	}
 
 	// The required API level can have the batch be rejected. But the header is always removed.
@@ -7264,7 +7272,7 @@ func (mset *stream) internalLoop() {
 				if len(im.rply) > 0 && strings.HasSuffix(im.rply, FastBatchSuffix) {
 					mset.processJetStreamFastBatchMsg(im.subj, im.rply, im.hdr, im.msg, im.mt)
 				} else if batchId := getBatchId(im.hdr); batchId != _EMPTY_ {
-					mset.processJetStreamBatchMsg(batchId, im.subj, im.rply, im.hdr, im.msg, im.mt)
+					mset.processJetStreamAtomicBatchMsg(batchId, im.subj, im.rply, im.hdr, im.msg, im.mt)
 				} else if isClustered {
 					mset.processClusteredInboundMsg(im.subj, im.rply, im.hdr, im.msg, im.mt, false)
 				} else {