kgo: do not trigger draining early if a partition moves sinks while lingering

twmb · twmb · commit 73c16c1d4584 · 2025-10-24T14:08:46.000-06:00
This extracts the 3x duplicated logic for triggering draining to a new
function. The buffering record logic LOOKED different, but was the same.

We now use this new function when a partition moves sinks: rather than
always triggering draining, we take lingering into consideration.
diff --git a/pkg/kgo/sink.go b/pkg/kgo/sink.go
@@ -134,8 +134,9 @@ func (s *sink) createReq(id int64, epoch int16) (*produceRequest, *kmsg.AddParti
 		recBuf.inflight++
 
 		recBuf.batchDrainIdx++
+		recBuf.lockedStopLinger()
 		recBuf.seq = incrementSequence(recBuf.seq, int32(len(batch.records)))
-		moreToDrain = recBuf.tryStopLingerForDraining() || moreToDrain
+		moreToDrain = recBuf.checkIfShouldDrainOrStartLinger() || moreToDrain
 		recBuf.mu.Unlock()
 
 		txnBuilder.add(recBuf)
@@ -1384,18 +1385,17 @@ func (recBuf *recBuf) bufferRecord(pr promisedRec, abortOnNewBatch bool) bool {
 	}
 
 	var (
-		newBatch       = true
-		onDrainBatch   = recBuf.batchDrainIdx == len(recBuf.batches)
+		mkNewBatch     = true
 		produceVersion = recBuf.sink.produceVersion.Load()
 	)
 
-	if !onDrainBatch {
+	if recBuf.batchDrainIdx != len(recBuf.batches) {
 		batch := recBuf.batches[len(recBuf.batches)-1]
 		appended, _ := batch.tryBuffer(pr, produceVersion, recBuf.maxRecordBatchBytes, false)
-		newBatch = !appended
+		mkNewBatch = !appended
 	}
 
-	if newBatch {
+	if mkNewBatch {
 		newBatch := recBuf.newRecordBatch()
 		appended, aborted := newBatch.tryBuffer(pr, produceVersion, recBuf.maxRecordBatchBytes, abortOnNewBatch)
 
@@ -1413,22 +1413,7 @@ func (recBuf *recBuf) bufferRecord(pr promisedRec, abortOnNewBatch bool) bool {
 		recBuf.batches = append(recBuf.batches, newBatch)
 	}
 
-	if recBuf.cl.cfg.linger == 0 {
-		if onDrainBatch {
-			recBuf.sink.maybeDrain()
-		}
-	} else {
-		// With linger, if this is a new batch but not the first, we
-		// stop lingering and begin draining. The drain loop will
-		// restart our linger once this buffer has one batch left.
-		if newBatch && !onDrainBatch ||
-			// If this is the first batch, try lingering; if
-			// we cannot, we are being flushed and must drain.
-			onDrainBatch && !recBuf.lockedMaybeLinger() {
-			recBuf.lockedStopLinger()
-			recBuf.sink.maybeDrain()
-		}
-	}
+	recBuf.maybeTriggerDrain()
 
 	recBuf.buffered.Add(1)
 	if recBuf.cl.producer.hooks != nil && len(recBuf.cl.producer.hooks.partitioned) > 0 {
@@ -1440,19 +1425,22 @@ func (recBuf *recBuf) bufferRecord(pr promisedRec, abortOnNewBatch bool) bool {
 	return true
 }
 
-// Stops lingering, potentially restarting it, and returns whether there is
-// more to drain.
-//
-// If lingering, if there are more than one batches ready, there is definitely
-// more to drain and we should not linger. Otherwise, if we cannot restart
-// lingering, then we are flushing and also indicate there is more to drain.
-func (recBuf *recBuf) tryStopLingerForDraining() bool {
-	recBuf.lockedStopLinger()
-	canLinger := recBuf.cl.cfg.linger != 0
-	moreToDrain := !canLinger && len(recBuf.batches) > recBuf.batchDrainIdx ||
-		canLinger && (len(recBuf.batches) > recBuf.batchDrainIdx+1 ||
-			len(recBuf.batches) == recBuf.batchDrainIdx+1 && !recBuf.lockedMaybeLinger())
-	return moreToDrain
+// Maybe drains, maybe starts a linger.
+// Must be called while locked.
+func (recBuf *recBuf) maybeTriggerDrain() {
+	if recBuf.checkIfShouldDrainOrStartLinger() {
+		recBuf.lockedStopLinger()
+		recBuf.sink.maybeDrain()
+	}
+}
+
+// Checks and returns if we should drain; if not, this potentially starts a
+// linger timer.
+func (recBuf *recBuf) checkIfShouldDrainOrStartLinger() bool {
+	nbufBatches := len(recBuf.batches) - recBuf.batchDrainIdx
+	return recBuf.cl.cfg.linger == 0 && nbufBatches > 0 || // no lingering, and any batch exists? drain
+		nbufBatches > 1 || // lingering, and more than one batch exists? drain -- one is full
+		nbufBatches == 1 && !recBuf.lockedMaybeLinger() // only one batch; if we cannot start lingering, we are being flushed or have too much buffered and have to drain
 }
 
 // Begins a linger timer unless the producer is being flushed.
@@ -1584,9 +1572,7 @@ func (recBuf *recBuf) clearFailing() {
 	defer recBuf.mu.Unlock()
 
 	recBuf.failing = false
-	if len(recBuf.batches) != recBuf.batchDrainIdx {
-		recBuf.sink.maybeDrain()
-	}
+	recBuf.maybeTriggerDrain()
 }
 
 func (recBuf *recBuf) resetBatchDrainIdx() {
@@ -1763,14 +1749,7 @@ func (b *recBatch) decInflight() {
 		return
 	}
 	recBuf.inflightOnSink = nil
-
-	nbufBatches := len(recBuf.batches) - recBuf.batchDrainIdx
-	if recBuf.cl.cfg.linger == 0 && nbufBatches > 0 ||
-		nbufBatches > 1 ||
-		nbufBatches == 1 && !recBuf.lockedMaybeLinger() {
-		recBuf.lockedStopLinger()
-		recBuf.sink.maybeDrain()
-	}
+	recBuf.maybeTriggerDrain()
 }
 
 ////////////////////
