[fix][consumer] Initialize batch ack tracker from the broker ack set on redelivery

pulsar/consumer_partition.go

@@ -1338,13 +1338,6 @@ func (pc *partitionConsumer) MessageReceived(response *pb.CommandMessage, header
 		numMsgs = int(msgMeta.GetNumMessagesInBatch())
 	}
 
-	messages := make([]*message, 0)
-	var ackTracker *ackTracker
-	// are there multiple messages in this batch?
-	if numMsgs > 1 {
-		ackTracker = newAckTracker(uint(numMsgs))
-	}
-
 	var ackSet *bitset.BitSet
 	if response.GetAckSet() != nil {
 		ackSetFromResponse := response.GetAckSet()
@@ -1355,6 +1348,13 @@ func (pc *partitionConsumer) MessageReceived(response *pb.CommandMessage, header
 		ackSet = bitset.From(buf)
 	}
 
+	messages := make([]*message, 0)
+	var ackTracker *ackTracker
+	// are there multiple messages in this batch?
+	if numMsgs > 1 {
+		ackTracker = newAckTracker(uint(numMsgs), ackSet)
+	}
+
 	pc.metrics.MessagesReceived.Add(float64(numMsgs))
 	pc.metrics.PrefetchedMessages.Add(float64(numMsgs))
 

pulsar/consumer_test.go

@@ -4822,6 +4822,144 @@ func TestBatchIndexAck(t *testing.T) {
 	}
 }
 
+// Go adaptation of testAckedBatchMessageNotSentToDeadLetterTopicOnFinalRedeliveryRound
+// from apache/pulsar (DeadLetterTopicTest): messages of a partially-acked batch that are
+// acked on their final redelivery round must not be routed to the DLQ.
+//
+// The final-round acks go through a consumer WITHOUT batch index ack. Such a consumer can
+// only ack a batch as a whole entry, once its ack tracker reports every index as acked (a
+// consumer with batch index ack would ack each index at the broker directly, and would not
+// depend on the tracker completing). If the tracker of a redelivered batch is not
+// initialized from the broker-provided ack set, it wrongly counts the already-acked indexes
+// as outstanding, so acking all remaining messages on their final redelivery round never
+// completes it: the acks are silently dropped, the broker redelivers the messages to the
+// next consumer session, and once their redeliveries are exhausted they are routed to the
+// DLQ -- even though the application acked every one of them in time.
+func TestAckedBatchMessageNotSentToDeadLetterTopicWithoutBatchIndexAck(t *testing.T) {
+	client, err := NewClient(ClientOptions{
+		URL: lookupURL,
+	})
+	require.NoError(t, err)
+	defer client.Close()
+
+	topic := newTopicName()
+	maxRedeliveryCount := 3
+	batchSize := 5
+	subscriptionName := "my-subscription"
+	dlqTopic := topic + "-" + subscriptionName + "-DLQ"
+
+	subscribe := func(batchIndexAck bool) Consumer {
+		consumer, err := client.Subscribe(ConsumerOptions{
+			Topic:                          topic,
+			SubscriptionName:               subscriptionName,
+			Type:                           Shared,
+			SubscriptionInitialPosition:    SubscriptionPositionEarliest,
+			EnableBatchIndexAcknowledgment: batchIndexAck,
+			AckWithResponse:                true,
+			NackRedeliveryDelay:            1 * time.Second,
+			DLQ: &DLQPolicy{
+				MaxDeliveries:   uint32(maxRedeliveryCount),
+				DeadLetterTopic: dlqTopic,
+			},
+		})
+		require.NoError(t, err)
+		return consumer
+	}
+
+	receive := func(consumer Consumer, timeout time.Duration) (Message, error) {
+		receiveCtx, cancel := context.WithTimeout(context.Background(), timeout)
+		defer cancel()
+		return consumer.Receive(receiveCtx)
+	}
+
+	deadLetterConsumer, err := client.Subscribe(ConsumerOptions{
+		Topic:                       dlqTopic,
+		SubscriptionName:            subscriptionName,
+		SubscriptionInitialPosition: SubscriptionPositionEarliest,
+	})
+	require.NoError(t, err)
+	defer deadLetterConsumer.Close()
+
+	producer, err := client.CreateProducer(ProducerOptions{
+		Topic:                   topic,
+		DisableBatching:         false,
+		BatchingMaxMessages:     uint(batchSize),
+		BatchingMaxPublishDelay: time.Hour,
+	})
+	require.NoError(t, err)
+	defer producer.Close()
+
+	ctx := context.Background()
+	sendErrCh := make(chan error, batchSize)
+	for i := 0; i < batchSize; i++ {
+		producer.SendAsync(ctx, &ProducerMessage{
+			Payload: []byte(fmt.Sprintf("message-%d", i)),
+		}, func(_ MessageID, _ *ProducerMessage, err error) {
+			sendErrCh <- err
+		})
+	}
+	require.NoError(t, producer.FlushWithCtx(ctx))
+	for i := 0; i < batchSize; i++ {
+		require.NoError(t, <-sendErrCh)
+	}
+
+	// Ack batch indexes 0, 3 and 4 with batch index ack enabled, so that the broker keeps a
+	// partial ack set for the batch, then close the consumer to get 1 and 2 redelivered.
+	consumer := subscribe(true)
+	messagesByBatchIdx := make(map[int32]Message)
+	for i := 0; i < batchSize; i++ {
+		message, err := receive(consumer, 5*time.Second)
+		require.NoError(t, err)
+		messagesByBatchIdx[message.ID().BatchIdx()] = message
+	}
+	require.Len(t, messagesByBatchIdx, batchSize)
+	for _, batchIndex := range []int32{0, 3, 4} {
+		require.NoError(t, consumer.Ack(messagesByBatchIdx[batchIndex]))
+	}
+	consumer.Close()
+
+	// Indexes 1 and 2 are redelivered to a consumer without batch index ack, along with the
+	// broker ack set. Nack them on every round but the final one (redeliveryCount ==
+	// maxRedeliveryCount-1; one more redelivery would route them to the DLQ instead of the
+	// application), where they are explicitly acked. Note that redelivery caused by closing
+	// a consumer does not increment redeliveryCount, only nacks do.
+	consumer = subscribe(false)
+	for acked := 0; acked < 2; {
+		message, err := receive(consumer, 5*time.Second)
+		require.NoError(t, err)
+		require.Contains(t, []string{"message-1", "message-2"}, string(message.Payload()))
+		if message.RedeliveryCount() < uint32(maxRedeliveryCount-1) {
+			// Let it be redelivered instead of acking.
+			consumer.Nack(message)
+			continue
+		}
+		require.NoError(t, consumer.Ack(message))
+		acked++
+	}
+	consumer.Close()
+
+	// Every message of the batch has been acked at or before its final allowed redelivery
+	// round, so a new consumer session must not receive anything. If the final-round acks
+	// were lost, the broker redelivers messages 1 and 2 here; the application cannot tell
+	// them apart from any other redelivered message, and nacking them exhausts their
+	// redeliveries and routes them to the DLQ.
+	consumer = subscribe(false)
+	defer consumer.Close()
+	for {
+		message, err := receive(consumer, 3*time.Second)
+		if err != nil {
+			break
+		}
+		consumer.Nack(message)
+	}
+
+	// No message should ever be routed to the DLQ.
+	receiveCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+	deadLetterMessage, err := deadLetterConsumer.Receive(receiveCtx)
+	require.Error(t, err, "no message should have been routed to the DLQ, but received: %v", deadLetterMessage)
+}
+
 func runBatchIndexAckTest(t *testing.T, ackWithResponse bool, cumulative bool, option *AckGroupingOptions) {
 	client, err := NewClient(ClientOptions{
 		URL: lookupURL,

pulsar/impl_message.go

@@ -409,10 +409,17 @@ func (msg *message) getConn() internal.Connection {
 	return msg.conn
 }
 
-func newAckTracker(size uint) *ackTracker {
-	batchIDs := bitset.New(size)
-	for i := uint(0); i < size; i++ {
-		batchIDs.Set(i)
+// newAckTracker creates a tracker for a batch of the given size. batchIDs is the set of
+// indexes still outstanding (e.g. the broker-provided ack set of a redelivered batch); if
+// nil, every index is considered outstanding.
+func newAckTracker(size uint, batchIDs *bitset.BitSet) *ackTracker {
+	if batchIDs == nil {
+		batchIDs = bitset.New(size)
+		for i := uint(0); i < size; i++ {
+			batchIDs.Set(i)
+		}
+	} else {
+		batchIDs = batchIDs.Clone()
 	}
 	return &ackTracker{
 		size:     size,

pulsar/impl_message_test.go

@@ -20,6 +20,7 @@ package pulsar
 import (
 	"testing"
 
+	"github.com/bits-and-blooms/bitset"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -57,11 +58,11 @@ func TestMessageIdGetFuncs(t *testing.T) {
 }
 
 func TestAckTracker(t *testing.T) {
-	tracker := newAckTracker(1)
+	tracker := newAckTracker(1, nil)
 	assert.Equal(t, true, tracker.ack(0))
 
 	// test 64
-	tracker = newAckTracker(64)
+	tracker = newAckTracker(64, nil)
 	for i := 0; i < 64; i++ {
 		if i < 63 {
 			assert.Equal(t, false, tracker.ack(i))
@@ -72,7 +73,7 @@ func TestAckTracker(t *testing.T) {
 	assert.Equal(t, true, tracker.completed())
 
 	// test large number 1000
-	tracker = newAckTracker(1000)
+	tracker = newAckTracker(1000, nil)
 	for i := 0; i < 1000; i++ {
 		if i < 999 {
 			assert.Equal(t, false, tracker.ack(i))
@@ -83,7 +84,7 @@ func TestAckTracker(t *testing.T) {
 	}
 	assert.Equal(t, true, tracker.completed())
 
-	tracker = newAckTracker(1000)
+	tracker = newAckTracker(1000, nil)
 	for i := 0; i < 1000; i++ {
 		if i < 999 {
 			assert.False(t, tracker.ackCumulative(i))
@@ -95,21 +96,65 @@ func TestAckTracker(t *testing.T) {
 	}
 
 	// test large number 1000 cumulative
-	tracker = newAckTracker(1000)
+	tracker = newAckTracker(1000, nil)
 
 	assert.Equal(t, true, tracker.ackCumulative(999))
 	assert.Equal(t, true, tracker.completed())
 }
 
+func TestAckTrackerWithBatchIDs(t *testing.T) {
+	// batchIDs mirrors the ack set the broker attaches when redelivering a partially-acked
+	// batch: bits 1 and 2 set means only indexes 1 and 2 are still outstanding.
+	newOutstanding12 := func() *bitset.BitSet {
+		batchIDs := bitset.New(5)
+		batchIDs.Set(1)
+		batchIDs.Set(2)
+		return batchIDs
+	}
+
+	tracker := newAckTracker(5, newOutstanding12())
+	assert.False(t, tracker.completed())
+	ackBitSet := tracker.getAckBitSet()
+	assert.Equal(t, uint(2), ackBitSet.Count())
+	assert.True(t, ackBitSet.Test(1))
+	assert.True(t, ackBitSet.Test(2))
+
+	// acking an index the broker already considers acked must not complete the tracker
+	assert.False(t, tracker.ack(0))
+	// acking the outstanding indexes completes the tracker
+	assert.False(t, tracker.ack(1))
+	assert.True(t, tracker.ack(2))
+	assert.True(t, tracker.completed())
+	assert.Nil(t, tracker.toAckSet())
+
+	// the tracker must own a copy: mutating the caller's bitset must not leak into it
+	batchIDs := newOutstanding12()
+	tracker = newAckTracker(5, batchIDs)
+	batchIDs.Set(3)
+	assert.False(t, tracker.ack(1))
+	assert.True(t, tracker.ack(2))
+	assert.True(t, tracker.completed())
+
+	// cumulative ack covering the outstanding indexes completes the tracker
+	tracker = newAckTracker(5, newOutstanding12())
+	assert.True(t, tracker.ackCumulative(2))
+	assert.True(t, tracker.completed())
+
+	// nil batchIDs keeps the previous behavior: every index starts outstanding
+	tracker = newAckTracker(5, nil)
+	assert.Equal(t, uint(5), tracker.getAckBitSet().Count())
+	assert.False(t, tracker.completed())
+}
+
 func TestAckingMessageIDBatchOne(t *testing.T) {
-	tracker := newAckTracker(1)
+	tracker := newAckTracker(1, nil)
 	msgID := newTrackingMessageID(1, 1, 0, 0, 0, tracker)
 	assert.Equal(t, true, msgID.ack())
 	assert.Equal(t, true, tracker.completed())
 }
 
 func TestAckingMessageIDBatchTwo(t *testing.T) {
-	tracker := newAckTracker(2)
+	tracker := newAckTracker(2, nil)
 	ids := []*trackingMessageID{
 		newTrackingMessageID(1, 1, 0, 0, 0, tracker),
 		newTrackingMessageID(1, 1, 1, 0, 0, tracker),
@@ -120,7 +165,7 @@ func TestAckingMessageIDBatchTwo(t *testing.T) {
 	assert.Equal(t, true, tracker.completed())
 
 	// try reverse order
-	tracker = newAckTracker(2)
+	tracker = newAckTracker(2, nil)
 	ids = []*trackingMessageID{
 		newTrackingMessageID(1, 1, 0, 0, 0, tracker),
 		newTrackingMessageID(1, 1, 1, 0, 0, tracker),