Repository: apache/pulsar-client-go
Pull request #1311 (2024-11-26 07:45:22)

Revert #1283 and to fix possible deadlock when queueCh is full
### Motivation

For the issue please refer to the PR description: #1310

Here have an analysis: https://github.com/apache/pulsar-client-go/pull/1310#issuecomment-2499888281

### Modifications
- Rever #1283 
- Add test from #1310 to cover `queueCh` never not full.

### Does this pull request potentially affect one of the following parts:

*If `yes` was chosen, please highlight the changes*

  - Dependencies (does it add or upgrade a dependency): (yes / no)
  - The public API: (yes / no)
  - The schema: (yes / no / don't know)
  - The default values of configurations: (yes / no)
  - The wire protocol: (yes / no)

### Documentation

  - Does this pull request introduce a new feature? (yes / no)
  - If yes, how is the feature documented? (not applicable / docs / GoDocs / not documented)
  - If a feature is not applicable for documentation, explain why?
  - If a feature is not documented yet in this PR, please create a followup issue for adding the documentation



Patch:
diff --git a/pulsar/consumer_partition.go b/pulsar/consumer_partition.go
index ae5e43a..983a6e2 100644
--- a/pulsar/consumer_partition.go
+++ b/pulsar/consumer_partition.go
@@ -156,7 +156,7 @@ type partitionConsumer struct {
 
 	// the size of the queue channel for buffering messages
 	maxQueueSize    int32
-	queueCh         chan *message
+	queueCh         chan []*message
 	startMessageID  atomicMessageID
 	lastDequeuedMsg *trackingMessageID
 
@@ -354,7 +354,7 @@ func newPartitionConsumer(parent Consumer, client *client, options *partitionCon
 		partitionIdx:               int32(options.partitionIdx),
 		eventsCh:                   make(chan interface{}, 10),
 		maxQueueSize:               int32(options.receiverQueueSize),
-		queueCh:                    make(chan *message, options.receiverQueueSize),
+		queueCh:                    make(chan []*message, options.receiverQueueSize),
 		startMessageID:             atomicMessageID{msgID: options.startMessageID},
 		connectedCh:                make(chan struct{}),
 		messageCh:                  messageCh,
@@ -1166,33 +1166,36 @@ func (pc *partitionConsumer) MessageReceived(response *pb.CommandMessage, header
 			return fmt.Errorf("discarding message on decryption error :%v", err)
 		case crypto.ConsumerCryptoFailureActionConsume:
 			pc.log.Warnf("consuming encrypted message due to error in decryption :%v", err)
+			messages := []*message{
+				{
+					publishTime:  timeFromUnixTimestampMillis(msgMeta.GetPublishTime()),
+					eventTime:    timeFromUnixTimestampMillis(msgMeta.GetEventTime()),
+					key:          msgMeta.GetPartitionKey(),
+					producerName: msgMeta.GetProducerName(),
+					properties:   internal.ConvertToStringMap(msgMeta.GetProperties()),
+					topic:        pc.topic,
+					msgID: newMessageID(
+						int64(pbMsgID.GetLedgerId()),
+						int64(pbMsgID.GetEntryId()),
+						pbMsgID.GetBatchIndex(),
+						pc.partitionIdx,
+						pbMsgID.GetBatchSize(),
+					),
+					payLoad:             headersAndPayload.ReadableSlice(),
+					schema:              pc.options.schema,
+					replicationClusters: msgMeta.GetReplicateTo(),
+					replicatedFrom:      msgMeta.GetReplicatedFrom(),
+					redeliveryCount:     response.GetRedeliveryCount(),
+					encryptionContext:   createEncryptionContext(msgMeta),
+					orderingKey:         string(msgMeta.OrderingKey),
+				},
+			}
 			if pc.options.autoReceiverQueueSize {
 				pc.incomingMessages.Inc()
 				pc.markScaleIfNeed()
 			}
 
-			pc.queueCh <- &message{
-				publishTime:  timeFromUnixTimestampMillis(msgMeta.GetPublishTime()),
-				eventTime:    timeFromUnixTimestampMillis(msgMeta.GetEventTime()),
-				key:          msgMeta.GetPartitionKey(),
-				producerName: msgMeta.GetProducerName(),
-				properties:   internal.ConvertToStringMap(msgMeta.GetProperties()),
-				topic:        pc.topic,
-				msgID: newMessageID(
-					int64(pbMsgID.GetLedgerId()),
-					int64(pbMsgID.GetEntryId()),
-					pbMsgID.GetBatchIndex(),
-					pc.partitionIdx,
-					pbMsgID.GetBatchSize(),
-				),
-				payLoad:             headersAndPayload.ReadableSlice(),
-				schema:              pc.options.schema,
-				replicationClusters: msgMeta.GetReplicateTo(),
-				replicatedFrom:      msgMeta.GetReplicatedFrom(),
-				redeliveryCount:     response.GetRedeliveryCount(),
-				encryptionContext:   createEncryptionContext(msgMeta),
-				orderingKey:         string(msgMeta.OrderingKey),
-			}
+			pc.queueCh <- messages
 			return nil
 		}
 	}
@@ -1228,6 +1231,7 @@ 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 {
@@ -1248,6 +1252,7 @@ func (pc *partitionConsumer) MessageReceived(response *pb.CommandMessage, header
 	pc.metrics.PrefetchedMessages.Add(float64(numMsgs))
 
 	var (
+		bytesReceived   int
 		skippedMessages int32
 	)
 	for i := 0; i < numMsgs; i++ {
@@ -1366,20 +1371,21 @@ func (pc *partitionConsumer) MessageReceived(response *pb.CommandMessage, header
 			Consumer: pc.parentConsumer,
 			Message:  msg,
 		})
-
+		messages = append(messages, msg)
+		bytesReceived += msg.size()
 		if pc.options.autoReceiverQueueSize {
-			pc.client.memLimit.ForceReserveMemory(int64(msg.size()))
+			pc.client.memLimit.ForceReserveMemory(int64(bytesReceived))
 			pc.incomingMessages.Add(int32(1))
 			pc.markScaleIfNeed()
 		}
-
-		pc.queueCh <- msg
 	}
 
 	if skippedMessages > 0 {
 		pc.availablePermits.add(skippedMessages)
 	}
 
+	// send messages to the dispatcher
+	pc.queueCh <- messages
 	return nil
 }
 
@@ -1535,19 +1541,20 @@ func (pc *partitionConsumer) dispatcher() {
 	defer func() {
 		pc.log.Debug("exiting dispatch loop")
 	}()
-	var queueMsg *message
+	var messages []*message
 	for {
-		var queueCh chan *message
+		var queueCh chan []*message
 		var messageCh chan ConsumerMessage
 		var nextMessage ConsumerMessage
 		var nextMessageSize int
 
-		if queueMsg != nil {
+		// are there more messages to send?
+		if len(messages) > 0 {
 			nextMessage = ConsumerMessage{
 				Consumer: pc.parentConsumer,
-				Message:  queueMsg,
+				Message:  messages[0],
 			}
-			nextMessageSize = queueMsg.size()
+			nextMessageSize = messages[0].size()
 
 			if !pc.isSeeking.Load() {
 				if pc.dlq.shouldSendToDlq(&nextMessage) {
@@ -1559,7 +1566,7 @@ func (pc *partitionConsumer) dispatcher() {
 					messageCh = pc.messageCh
 				}
 				pc.metrics.PrefetchedMessages.Dec()
-				pc.metrics.PrefetchedBytes.Sub(float64(len(queueMsg.payLoad)))
+				pc.metrics.PrefetchedBytes.Sub(float64(len(messages[0].payLoad)))
 			} else {
 				pc.log.Debug("skip dispatching messages when seeking")
 			}
@@ -1577,7 +1584,7 @@ func (pc *partitionConsumer) dispatcher() {
 			}
 			pc.log.Debug("dispatcher received connection event")
 
-			queueMsg = nil
+			messages = nil
 
 			// reset available permits
 			pc.availablePermits.reset()
@@ -1602,16 +1609,18 @@ func (pc *partitionConsumer) dispatcher() {
 			pc.log.Debug("received dispatcherSeekingControlCh, set isSeek to true")
 			pc.isSeeking.Store(true)
 
-		case msg, ok := <-queueCh:
+		case msgs, ok := <-queueCh:
 			if !ok {
 				return
 			}
 
-			queueMsg = msg
+			messages = msgs
 
 		// if the messageCh is nil or the messageCh is full this will not be selected
 		case messageCh <- nextMessage:
-			queueMsg = nil
+			// allow this message to be garbage collected
+			messages[0] = nil
+			messages = messages[1:]
 
 			pc.availablePermits.inc()
 
@@ -1634,14 +1643,14 @@ func (pc *partitionConsumer) dispatcher() {
 				if m == nil {
 					break
 				} else if nextMessageInQueue == nil {
-					nextMessageInQueue = toTrackingMessageID(m.msgID)
+					nextMessageInQueue = toTrackingMessageID(m[0].msgID)
 				}
 				if pc.options.autoReceiverQueueSize {
-					pc.incomingMessages.Sub(int32(1))
+					pc.incomingMessages.Sub(int32(len(m)))
 				}
 			}
 
-			queueMsg = nil
+			messages = nil
 
 			clearQueueCb(nextMessageInQueue)
 		}


Test patch:
diff --git a/pulsar/consumer_partition_test.go b/pulsar/consumer_partition_test.go
index 746e17a..4594c14 100644
--- a/pulsar/consumer_partition_test.go
+++ b/pulsar/consumer_partition_test.go
@@ -30,7 +30,7 @@ import (
 func TestSingleMessageIDNoAckTracker(t *testing.T) {
 	eventsCh := make(chan interface{}, 1)
 	pc := partitionConsumer{
-		queueCh:              make(chan *message, 1),
+		queueCh:              make(chan []*message, 1),
 		eventsCh:             eventsCh,
 		compressionProviders: sync.Map{},
 		options:              &partitionConsumerOpts{},
@@ -47,12 +47,13 @@ func TestSingleMessageIDNoAckTracker(t *testing.T) {
 	}
 
 	// ensure the tracker was set on the message id
-	message := <-pc.queueCh
-	id := message.ID().(*trackingMessageID)
-	assert.Nil(t, id.tracker)
+	messages := <-pc.queueCh
+	for _, m := range messages {
+		assert.Nil(t, m.ID().(*trackingMessageID).tracker)
+	}
 
 	// ack the message id
-	pc.AckID(id)
+	pc.AckID(messages[0].msgID.(*trackingMessageID))
 
 	select {
 	case <-eventsCh:
@@ -68,7 +69,7 @@ func newTestMetrics() *internal.LeveledMetrics {
 func TestBatchMessageIDNoAckTracker(t *testing.T) {
 	eventsCh := make(chan interface{}, 1)
 	pc := partitionConsumer{
-		queueCh:              make(chan *message, 1),
+		queueCh:              make(chan []*message, 1),
 		eventsCh:             eventsCh,
 		compressionProviders: sync.Map{},
 		options:              &partitionConsumerOpts{},
@@ -85,12 +86,13 @@ func TestBatchMessageIDNoAckTracker(t *testing.T) {
 	}
 
 	// ensure the tracker was set on the message id
-	message := <-pc.queueCh
-	id := message.ID().(*trackingMessageID)
-	assert.Nil(t, id.tracker)
+	messages := <-pc.queueCh
+	for _, m := range messages {
+		assert.Nil(t, m.ID().(*trackingMessageID).tracker)
+	}
 
 	// ack the message id
-	err := pc.AckID(id)
+	err := pc.AckID(messages[0].msgID.(*trackingMessageID))
 	assert.Nil(t, err)
 
 	select {
@@ -103,7 +105,7 @@ func TestBatchMessageIDNoAckTracker(t *testing.T) {
 func TestBatchMessageIDWithAckTracker(t *testing.T) {
 	eventsCh := make(chan interface{}, 1)
 	pc := partitionConsumer{
-		queueCh:              make(chan *message, 10),
+		queueCh:              make(chan []*message, 10),
 		eventsCh:             eventsCh,
 		compressionProviders: sync.Map{},
 		options:              &partitionConsumerOpts{},
@@ -120,21 +122,14 @@ func TestBatchMessageIDWithAckTracker(t *testing.T) {
 	}
 
 	// ensure the tracker was set on the message id
-	var messageIDs []*trackingMessageID
-	for i := 0; i < 10; i++ {
-		select {
-		case m := <-pc.queueCh:
-			id := m.ID().(*trackingMessageID)
-			assert.NotNil(t, id.tracker)
-			messageIDs = append(messageIDs, id)
-		default:
-			break
-		}
+	messages := <-pc.queueCh
+	for _, m := range messages {
+		assert.NotNil(t, m.ID().(*trackingMessageID).tracker)
 	}
 
 	// ack all message ids except the last one
 	for i := 0; i < 9; i++ {
-		err := pc.AckID(messageIDs[i])
+		err := pc.AckID(messages[i].msgID.(*trackingMessageID))
 		assert.Nil(t, err)
 	}
 
@@ -145,7 +140,7 @@ func TestBatchMessageIDWithAckTracker(t *testing.T) {
 	}
 
 	// ack last message
-	err := pc.AckID(messageIDs[9])
+	err := pc.AckID(messages[9].msgID.(*trackingMessageID))
 	assert.Nil(t, err)
 
 	select {
diff --git a/pulsar/consumer_test.go b/pulsar/consumer_test.go
index 83521dd..e3bf9bf 100644
--- a/pulsar/consumer_test.go
+++ b/pulsar/consumer_test.go
@@ -4877,3 +4877,66 @@ func receiveMessages(t *testing.T, consumer Consumer, numMessages int) []Message
 	assert.Equal(t, numMessages, len(msgs))
 	return msgs
 }
+
+func TestAckResponseNotBlocked(t *testing.T) {
+	client, err := NewClient(ClientOptions{
+		URL:              lookupURL,
+		OperationTimeout: 5 * time.Second,
+	})
+	assert.Nil(t, err)
+	defer client.Close()
+
+	topic := fmt.Sprintf("test-ack-response-not-blocked-%v", time.Now().Nanosecond())
+	assert.Nil(t, createPartitionedTopic(topic, 10))
+
+	producer, err := client.CreateProducer(ProducerOptions{
+		Topic: topic,
+	})
+	assert.Nil(t, err)
+
+	ctx := context.Background()
+	numMessages := 1000
+	for i := 0; i < numMessages; i++ {
+		producer.SendAsync(ctx, &ProducerMessage{
+			Payload: []byte(fmt.Sprintf("value-%d", i)),
+		}, func(_ MessageID, _ *ProducerMessage, err error) {
+			if err != nil {
+				t.Fatal(err)
+			}
+		})
+		if i%100 == 99 {
+			assert.Nil(t, producer.Flush())
+		}
+	}
+	producer.Flush()
+	producer.Close()
+
+	// Set a small receiver queue size to trigger ack response blocking if the internal `queueCh`
+	// is a channel with the same size
+	consumer, err := client.Subscribe(ConsumerOptions{
+		Topic:                          topic,
+		SubscriptionName:               "my-sub",
+		SubscriptionInitialPosition:    SubscriptionPositionEarliest,
+		Type:                           KeyShared,
+		EnableBatchIndexAcknowledgment: true,
+		AckWithResponse:                true,
+		ReceiverQueueSize:              5,
+	})
+	assert.Nil(t, err)
+	msgIDs := make([]MessageID, 0)
+	for i := 0; i < numMessages; i++ {
+		if msg, err := consumer.Receive(context.Background()); err != nil {
+			t.Fatal(err)
+		} else {
+			msgIDs = append(msgIDs, msg.ID())
+			if len(msgIDs) >= 10 {
+				if err := consumer.AckIDList(msgIDs); err != nil {
+					t.Fatal("Failed to acked messages: ", msgIDs, " ", err)
+				} else {
+					t.Log("Acked messages: ", msgIDs)
+				}
+				msgIDs = msgIDs[:0]
+			}
+		}
+	}
+}