Fix batch race

Author: nodece

pulsar/producer_partition.go

@@ -112,6 +112,7 @@ type partitionProducer struct {
 	// Channel where app is posting messages to be published
 	dataChan             chan *pendingItem
 	cmdChan              chan interface{}
+	batchChan            chan interface{}
 	connectClosedCh      chan *connectionClosed
 	publishSemaphore     internal.Semaphore
 	pendingQueue         internal.BlockingQueue
@@ -228,6 +229,7 @@ func newPartitionProducer(client *client, topic string, options *ProducerOptions
 	p.setProducerState(producerReady)
 
 	if !p.options.DisableBatching {
+		p.batchChan = make(chan interface{})
 		p.batchFlushTicker = time.NewTicker(batchingMaxPublishDelay)
 		go p.listenBatch()
 	}
@@ -562,7 +564,47 @@ func (p *partitionProducer) listenBatch() {
 	for {
 		select {
 		case <-p.batchFlushTicker.C:
-			p.internalFlushCurrentBatch()
+			p.internalFlushCurrentBatch(false)
+		case req, ok := <-p.batchChan:
+			if !ok {
+				return
+			}
+
+			switch req.(type) {
+			case *sendRequest:
+				sr := req.(*sendRequest)
+				smm := p.genSingleMessageMetadataInBatch(sr.msg, int(sr.uncompressedSize))
+				multiSchemaEnabled := !p.options.DisableMultiSchema
+
+				added := addRequestToBatch(
+					smm, p, sr.uncompressedPayload, sr, sr.msg, sr.deliverAt, sr.schemaVersion, multiSchemaEnabled)
+				if !added {
+					// The current batch is full. flush it and retry
+					p.internalFlushCurrentBatch(false)
+
+					// after flushing try again to add the current payload
+					ok := addRequestToBatch(
+						smm, p, sr.uncompressedPayload, sr, sr.msg, sr.deliverAt, sr.schemaVersion, multiSchemaEnabled)
+					if !ok {
+						p.log.WithField("size", sr.uncompressedSize).
+							WithField("properties", sr.msg.Properties).
+							Error("unable to add message to batch")
+						sr.done(nil, ErrFailAddToBatch)
+						continue
+					}
+				}
+				if sr.flushImmediately {
+					p.internalFlushCurrentBatch(false)
+					continue
+				}
+			case *flushRequest:
+				fr := req.(*flushRequest)
+				// Flush request happened on the event loop.
+				p.internalFlushCurrentBatch(true)
+				fr.doneCh <- struct{}{}
+			default:
+				p.log.Error("Unknown request type")
+			}
 		}
 	}
 }
@@ -615,30 +657,7 @@ func (p *partitionProducer) internalSend(sr *sendRequest) {
 	p.log.Debug("Received send request: ", *sr.msg)
 
 	if sr.sendAsBatch {
-		smm := p.genSingleMessageMetadataInBatch(sr.msg, int(sr.uncompressedSize))
-		multiSchemaEnabled := !p.options.DisableMultiSchema
-
-		added := addRequestToBatch(
-			smm, p, sr.uncompressedPayload, sr, sr.msg, sr.deliverAt, sr.schemaVersion, multiSchemaEnabled)
-		if !added {
-			// The current batch is full. flush it and retry
-			p.internalFlushCurrentBatch()
-
-			// after flushing try again to add the current payload
-			ok := addRequestToBatch(
-				smm, p, sr.uncompressedPayload, sr, sr.msg, sr.deliverAt, sr.schemaVersion, multiSchemaEnabled)
-			if !ok {
-				p.log.WithField("size", sr.uncompressedSize).
-					WithField("properties", sr.msg.Properties).
-					Error("unable to add message to batch")
-				sr.done(nil, ErrFailAddToBatch)
-				return
-			}
-		}
-
-		if sr.flushImmediately {
-			p.internalFlushCurrentBatch()
-		}
+		p.batchChan <- sr
 		return
 	}
 
@@ -843,7 +862,7 @@ func (p *partitionProducer) internalSingleSend(
 		return
 	}
 
-	p.writeData(buffer, sid, []interface{}{sr})
+	p.writeData(buffer, sid, []interface{}{sr}, false)
 }
 
 type pendingItem struct {
@@ -859,7 +878,7 @@ type pendingItem struct {
 	flushCallback func(err error)
 }
 
-func (p *partitionProducer) internalFlushCurrentBatch() {
+func (p *partitionProducer) internalFlushCurrentBatch(inEventLoop bool) {
 	if p.batchBuilder == nil {
 		// batch is not enabled
 		// the batch flush ticker should be stopped but it might still called once
@@ -868,7 +887,7 @@ func (p *partitionProducer) internalFlushCurrentBatch() {
 		return
 	}
 	if p.batchBuilder.IsMultiBatches() {
-		p.internalFlushCurrentBatches()
+		p.internalFlushCurrentBatches(inEventLoop)
 		return
 	}
 
@@ -894,10 +913,11 @@ func (p *partitionProducer) internalFlushCurrentBatch() {
 		return
 	}
 
-	p.writeData(batchData, sequenceID, callbacks)
+	p.writeData(batchData, sequenceID, callbacks, inEventLoop)
 }
 
-func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64, callbacks []interface{}) {
+func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64, callbacks []interface{},
+	inEventLoop bool) {
 	select {
 	case <-p.ctx.Done():
 		for _, cb := range callbacks {
@@ -919,7 +939,12 @@ func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64,
 			sendRequests: callbacks,
 		}
 		p.pendingQueue.Put(&item)
-		p.dataChan <- &item
+		if inEventLoop {
+			// When the flush request is called on the event loop, we can directly write the data.
+			p.internalWriteData(&item)
+		} else {
+			p.dataChan <- &item
+		}
 	}
 }
 
@@ -1015,7 +1040,7 @@ func (p *partitionProducer) failTimeoutMessages() {
 	}
 }
 
-func (p *partitionProducer) internalFlushCurrentBatches() {
+func (p *partitionProducer) internalFlushCurrentBatches(inEventLoop bool) {
 	flushBatches := p.batchBuilder.FlushBatches()
 	if flushBatches == nil {
 		return
@@ -1041,7 +1066,7 @@ func (p *partitionProducer) internalFlushCurrentBatches() {
 		if b.BatchData == nil {
 			continue
 		}
-		p.writeData(b.BatchData, b.SequenceID, b.Callbacks)
+		p.writeData(b.BatchData, b.SequenceID, b.Callbacks, inEventLoop)
 	}
 
 }
@@ -1050,7 +1075,17 @@ func (p *partitionProducer) internalFlush(fr *flushRequest) {
 	p.clearPendingSendRequests()
 
 	if !p.options.DisableBatching {
-		p.internalFlushCurrentBatch()
+		// Flush batch
+		batchFlushRequest := &flushRequest{
+			doneCh: make(chan struct{}),
+		}
+		p.batchChan <- batchFlushRequest
+		<-batchFlushRequest.doneCh
+		if batchFlushRequest.err != nil {
+			fr.err = batchFlushRequest.err
+			close(fr.doneCh)
+			return
+		}
 	}
 
 	pi, ok := p.pendingQueue.PeekLast().(*pendingItem)