[Issue 1395][Producer] Fix deduplication of batches that mix resent and new messages

pulsar/internal/batch_builder.go

@@ -71,6 +71,9 @@ type FlushBatch struct {
 	SequenceID uint64
 	Callbacks  []interface{}
 	Error      error
+
+	// HighestSequenceID is the largest sequence ID of the messages in the batch.
+	HighestSequenceID uint64
 }
 
 // batchContainer wraps the objects needed to a batch.
@@ -98,6 +101,11 @@ type batchContainer struct {
 	msgMetadata *pb.MessageMetadata
 	callbacks   []interface{}
 
+	// highestSequenceID is the largest sequence ID added to the current batch.
+	// It is sent as highest_sequence_id so that broker-side deduplication tracks
+	// the whole batch instead of only its first message.
+	highestSequenceID uint64
+
 	compressionProvider compression.Provider
 	buffersPool         BuffersPool
 	metrics             *Metrics
@@ -234,6 +242,9 @@ func (bc *batchContainer) Add(
 			bc.cmdSend.Send.TxnidMostBits = proto.Uint64(mostSigBits)
 			bc.cmdSend.Send.TxnidLeastBits = proto.Uint64(leastSigBits)
 		}
+		bc.highestSequenceID = sequenceID
+	} else if metadata.SequenceId != nil && *metadata.SequenceId > bc.highestSequenceID {
+		bc.highestSequenceID = *metadata.SequenceId
 	}
 	addSingleMessageToBatch(bc.buffer, metadata, payload)
 
@@ -244,6 +255,7 @@ func (bc *batchContainer) Add(
 
 func (bc *batchContainer) reset() {
 	bc.numMessages = 0
+	bc.highestSequenceID = 0
 	bc.buffer.Clear()
 	bc.callbacks = []interface{}{}
 	bc.msgMetadata.ReplicateTo = nil
@@ -263,6 +275,8 @@ func (bc *batchContainer) Flush() *FlushBatch {
 
 	bc.msgMetadata.NumMessagesInBatch = proto.Int32(int32(bc.numMessages))
 	bc.cmdSend.Send.NumMessages = proto.Int32(int32(bc.numMessages))
+	bc.msgMetadata.HighestSequenceId = proto.Uint64(bc.highestSequenceID)
+	bc.cmdSend.Send.HighestSequenceId = proto.Uint64(bc.highestSequenceID)
 
 	uncompressedSize := bc.buffer.ReadableBytes()
 	bc.msgMetadata.UncompressedSize = &uncompressedSize
@@ -273,21 +287,24 @@ func (bc *batchContainer) Flush() *FlushBatch {
 	buffer.SetReleaseCallback(func() { bufferCount.Dec() })
 
 	sequenceID := uint64(0)
+	highestSequenceID := uint64(0)
 	var err error
 	if err = serializeMessage(
 		buffer, bc.cmdSend, bc.msgMetadata, bc.buffer, bc.compressionProvider,
 		bc.encryptor, bc.maxMessageSize, true,
 	); err == nil { // no error in serializing Batch
 		sequenceID = bc.cmdSend.Send.GetSequenceId()
+		highestSequenceID = bc.cmdSend.Send.GetHighestSequenceId()
 	}
 
 	callbacks := bc.callbacks
 	bc.reset()
 	return &FlushBatch{
-		BatchData:  buffer,
-		SequenceID: sequenceID,
-		Callbacks:  callbacks,
-		Error:      err,
+		BatchData:         buffer,
+		SequenceID:        sequenceID,
+		HighestSequenceID: highestSequenceID,
+		Callbacks:         callbacks,
+		Error:             err,
 	}
 }
 

pulsar/internal/batch_builder_test.go

@@ -0,0 +1,132 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package internal
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+
+	"github.com/apache/pulsar-client-go/pulsar/internal/compression"
+	pb "github.com/apache/pulsar-client-go/pulsar/internal/pulsar_proto"
+	"github.com/apache/pulsar-client-go/pulsar/log"
+)
+
+// readBatchFrame decodes the CommandSend and the batch MessageMetadata of a flushed batch.
+func readBatchFrame(t *testing.T, buffer Buffer) (*pb.CommandSend, *pb.MessageMetadata) {
+	t.Helper()
+	frame := NewBufferWrapper(buffer.ReadableSlice())
+	frame.ReadUint32() // total size
+	var cmd pb.BaseCommand
+	require.NoError(t, proto.Unmarshal(frame.Read(frame.ReadUint32()), &cmd))
+	meta, err := NewMessageReader(frame).ReadMessageMetadata()
+	require.NoError(t, err)
+	return cmd.GetSend(), meta
+}
+
+func newTestBatchBuilder(t *testing.T, provider BatcherBuilderProvider) BatchBuilder {
+	t.Helper()
+	builder, err := provider(1000, 1<<20, 1<<20, "test", 1, pb.CompressionType_NONE, compression.Level(0),
+		&bufferPoolImpl{}, NewMetricsProvider(2, map[string]string{}, prometheus.NewRegistry()),
+		log.DefaultNopLogger(), &mockEncryptor{})
+	require.NoError(t, err)
+	return builder
+}
+
+func addWithSequenceID(t *testing.T, builder BatchBuilder, generator *uint64, key string, sequenceID *uint64) {
+	t.Helper()
+	metadata := &pb.SingleMessageMetadata{
+		PartitionKey: proto.String(key),
+		PayloadSize:  proto.Int32(4),
+		SequenceId:   sequenceID,
+	}
+	require.True(t, builder.Add(metadata, generator, []byte("test"), nil, nil, time.Now(),
+		nil, false, false, 0, 0))
+}
+
+func TestBatchContainerSetsHighestSequenceID(t *testing.T) {
+	builder := newTestBatchBuilder(t, NewBatchBuilder)
+	generator := uint64(0)
+
+	for _, sequenceID := range []uint64{5, 6, 9} {
+		addWithSequenceID(t, builder, &generator, "", proto.Uint64(sequenceID))
+	}
+	batch := builder.Flush()
+	require.NotNil(t, batch)
+	require.NoError(t, batch.Error)
+	assert.Equal(t, uint64(5), batch.SequenceID)
+	assert.Equal(t, uint64(9), batch.HighestSequenceID)
+
+	send, meta := readBatchFrame(t, batch.BatchData)
+	assert.Equal(t, uint64(5), send.GetSequenceId())
+	assert.Equal(t, uint64(9), send.GetHighestSequenceId())
+	assert.Equal(t, int32(3), send.GetNumMessages())
+	assert.Equal(t, uint64(5), meta.GetSequenceId())
+	assert.Equal(t, uint64(9), meta.GetHighestSequenceId())
+	assert.Zero(t, builder.(*batchContainer).highestSequenceID, "Flush must reset the highest sequence ID")
+
+	// The next batch must not inherit the previous highest sequence ID.
+	addWithSequenceID(t, builder, &generator, "", proto.Uint64(10))
+	batch = builder.Flush()
+	require.NotNil(t, batch)
+	assert.Equal(t, uint64(10), batch.SequenceID)
+	assert.Equal(t, uint64(10), batch.HighestSequenceID)
+	send, meta = readBatchFrame(t, batch.BatchData)
+	assert.Equal(t, uint64(10), send.GetHighestSequenceId())
+	assert.Equal(t, uint64(10), meta.GetHighestSequenceId())
+}
+
+func TestBatchContainerHighestSequenceIDWithGeneratedFirstID(t *testing.T) {
+	builder := newTestBatchBuilder(t, NewBatchBuilder)
+	generator := uint64(7)
+
+	// The first message takes its sequence ID from the generator inside Add.
+	addWithSequenceID(t, builder, &generator, "", nil)
+	batch := builder.Flush()
+	require.NotNil(t, batch)
+	assert.Equal(t, uint64(7), batch.SequenceID)
+	assert.Equal(t, uint64(7), batch.HighestSequenceID)
+	assert.Equal(t, uint64(8), generator)
+}
+
+func TestKeyBasedBatchContainerSetsHighestSequenceIDPerKey(t *testing.T) {
+	builder := newTestBatchBuilder(t, NewKeyBasedBatchBuilder)
+	generator := uint64(0)
+
+	sequenceID := uint64(0)
+	for i := 0; i < 9; i++ {
+		addWithSequenceID(t, builder, &generator, fmt.Sprintf("key-%d", i%3), proto.Uint64(sequenceID))
+		sequenceID++
+	}
+	batches := builder.FlushBatches()
+	require.Len(t, batches, 3)
+	for i, batch := range batches {
+		// key-i holds sequence IDs i, i+3 and i+6.
+		assert.Equal(t, uint64(i), batch.SequenceID)
+		assert.Equal(t, uint64(i+6), batch.HighestSequenceID)
+		send, meta := readBatchFrame(t, batch.BatchData)
+		assert.Equal(t, uint64(i), send.GetSequenceId())
+		assert.Equal(t, uint64(i+6), send.GetHighestSequenceId())
+		assert.Equal(t, uint64(i+6), meta.GetHighestSequenceId())
+	}
+}

pulsar/producer_partition.go

@@ -112,6 +112,13 @@ type partitionProducer struct {
 	encryptor                internalcrypto.Encryptor
 	compressionProvider      compression.Provider
 
+	// lastSequenceIDPushed is the highest sequence ID handed to the connection or
+	// added to a batch, and lastSequenceIDPotentialDuplicated tells whether the
+	// current batch holds a message whose sequence ID was not above it. Both are
+	// only accessed from the event loop, see internalSend.
+	lastSequenceIDPushed              int64
+	lastSequenceIDPotentialDuplicated bool
+
 	// Channel where app is posting messages to be published
 	dataChan             chan *sendRequest
 	cmdChan              chan interface{}
@@ -363,9 +370,15 @@ func (p *partitionProducer) grabCnx(assignedBrokerURL string) error {
 		p.encryptor = internalcrypto.NewNoopEncryptor()
 	}
 
+	// Only the first connection seeds the sequence state from the broker, as
+	// in the Java client. On a reconnect the local values already cover what
+	// this producer pushed, and grabCnx then runs on the reconnect goroutine
+	// concurrently with internalSend, which owns lastSequenceIDPushed.
 	if p.sequenceIDGenerator == nil {
-		nextSequenceID := uint64(res.Response.ProducerSuccess.GetLastSequenceId() + 1)
+		lastSequenceID := res.Response.ProducerSuccess.GetLastSequenceId()
+		nextSequenceID := uint64(lastSequenceID + 1)
 		p.sequenceIDGenerator = &nextSequenceID
+		p.lastSequenceIDPushed = lastSequenceID
 	}
 
 	schemaVersion := res.Response.ProducerSuccess.GetSchemaVersion()
@@ -641,6 +654,25 @@ func (p *partitionProducer) internalSend(sr *sendRequest) {
 		smm := p.genSingleMessageMetadataInBatch(sr.msg, int(sr.uncompressedSize))
 		multiSchemaEnabled := !p.options.DisableMultiSchema
 
+		// The broker deduplicates a batch by its lowest and highest sequence IDs only,
+		// so a batch that mixes a resent (potentially duplicated) message with new
+		// messages would be dropped or accepted as a whole. Keep potential duplicates
+		// in their own batch, as the Java client does.
+		sequenceID := int64(smm.GetSequenceId())
+		potentialDuplicate := sequenceID <= p.lastSequenceIDPushed
+		if potentialDuplicate {
+			if sequenceID <= p.LastSequenceID() {
+				p.log.Warnf("Message with sequence id %d is definitely a duplicate", sequenceID)
+			} else {
+				p.log.Infof("Message with sequence id %d might be a duplicate but cannot be determined at this time",
+					sequenceID)
+			}
+		}
+		if potentialDuplicate || p.lastSequenceIDPotentialDuplicated {
+			p.internalFlushCurrentBatch()
+		}
+		p.lastSequenceIDPotentialDuplicated = potentialDuplicate
+
 		added := addRequestToBatch(
 			smm, p, sr.uncompressedPayload, sr, sr.msg, sr.deliverAt, sr.schemaVersion, multiSchemaEnabled)
 		if !added {
@@ -658,13 +690,16 @@ func (p *partitionProducer) internalSend(sr *sendRequest) {
 				return
 			}
 		}
+		p.lastSequenceIDPushed = max(p.lastSequenceIDPushed, sequenceID)
 
 		if sr.flushImmediately {
 			p.internalFlushCurrentBatch()
 		}
 		return
 	}
 
+	p.lastSequenceIDPushed = max(p.lastSequenceIDPushed, int64(sr.mm.GetSequenceId()))
+
 	if sr.totalChunks <= 1 {
 		p.internalSingleSend(sr.mm, sr.compressedPayload, sr, uint32(sr.maxMessageSize))
 		return
@@ -870,7 +905,7 @@ func (p *partitionProducer) internalSingleSend(
 		return
 	}
 
-	p.writeData(buffer, sid, []interface{}{sr})
+	p.writeData(buffer, sid, sid, []interface{}{sr})
 }
 
 type pendingItem struct {
@@ -884,6 +919,10 @@ type pendingItem struct {
 	sendRequests   []interface{}
 	isDone         bool
 	flushCallbacks []func(err error)
+
+	// highestSequenceID is the largest sequence ID in this item. It differs from
+	// sequenceID only for batches.
+	highestSequenceID uint64
 }
 
 func (p *partitionProducer) internalFlushCurrentBatch() {
@@ -921,10 +960,11 @@ func (p *partitionProducer) internalFlushCurrentBatch() {
 		return
 	}
 
-	p.writeData(batchData, sequenceID, callbacks)
+	p.writeData(batchData, sequenceID, batch.HighestSequenceID, callbacks)
 }
 
-func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64, callbacks []interface{}) {
+func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID, highestSequenceID uint64,
+	callbacks []interface{}) {
 	select {
 	case <-p.ctx.Done():
 		buffer.Release()
@@ -949,13 +989,14 @@ func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64,
 			p.log.Debug("Skipping write to connection, producer state: ", p.getProducerState())
 		}
 		p.pendingQueue.PutUnsafe(&pendingItem{
-			ctx:          ctx,
-			cancel:       cancel,
-			createdAt:    now,
-			sentAt:       now,
-			buffer:       buffer,
-			sequenceID:   sequenceID,
-			sendRequests: callbacks,
+			ctx:               ctx,
+			cancel:            cancel,
+			createdAt:         now,
+			sentAt:            now,
+			buffer:            buffer,
+			sequenceID:        sequenceID,
+			highestSequenceID: highestSequenceID,
+			sendRequests:      callbacks,
 		})
 	}
 }
@@ -1077,7 +1118,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.HighestSequenceID, b.Callbacks)
 	}
 
 }
@@ -1453,9 +1494,11 @@ func (p *partitionProducer) ReceivedSendReceipt(response *pb.CommandSendReceipt)
 	}
 	p.metrics.PublishRPCLatency.Observe(float64(now-pi.sentAt.UnixNano()) / 1.0e9)
 	batchSize := int32(len(pi.sendRequests))
+	if highest := int64(max(pi.sequenceID, pi.highestSequenceID)); highest > atomic.LoadInt64(&p.lastSequenceID) {
+		atomic.StoreInt64(&p.lastSequenceID, highest)
+	}
 	for idx, i := range pi.sendRequests {
 		sr := i.(*sendRequest)
-		atomic.StoreInt64(&p.lastSequenceID, int64(pi.sequenceID))
 
 		msgID := newMessageID(
 			int64(response.MessageId.GetLedgerId()),