Repository: apache/pulsar-client-go
Pull request #976 (2023-03-03 09:17:42)

[feat]: Support auto scaled consumer receiver queue
Master Issue: #927 

### Motivation
**Note: This is part of the work for [PIP 74](https://github.com/apache/pulsar/wiki/PIP-74%3A-Pulsar-client-memory-limits) in go client.**

The memory limit in consumer side relys on flow control. So we need to firstly support auto scaled consumer receiver queue before we implement the consumer side memory limitation.

#### Why should we support auto scaled consumer first before supporting consumer client memory limit?

The main memory consumption on the consumer side comes from the messages that the consumer prefetches from the broker. Currently, the number of prefetched messages depends on the "current available permits" which represents the messages that have been delivered to the user. When the `availablePermits` reach a certain threshold, the client will send a `Flow` request to the broker to prefetch messages. The current threshold is half of the "receiverQueueSize" (default 1000).

https://github.com/apache/pulsar-client-go/blob/352c463194916cab7772c0729c9cdac14404e8bd/pulsar/consumer_partition.go#L194-L207

If we do not support auto scaled consumer receiver queue, the `Flow` threshold will be fixed. We cannot dynamically adjust the number of prefetched messages, resulting in the inability to control consumer memory. So auto scaled consumer receiver queue needs to be supported first.

#### How does "auto scaled consumer receiver queue" work?

In consumer side, all controls related to message pulling are actually controlling the `Flow`. For example, `receiverQueueSize` controls the `Flow` by controlling the size of the receiver queue.  ReceiverQueue `queueCh` can receive more messages than `receiverQueueSize` because queueCh is a `chan []*message`.

https://github.com/apache/pulsar-client-go/blob/352c463194916cab7772c0729c9cdac14404e8bd/pulsar/consumer_partition.go#L296

If we want to control the `Flow` threshold, we should make the `queueSize` can be auto scaled. There are mainly two problems to be solved.

- Under what condition is the queue considered to need to be scaled?

When the consumption rate required by the user is greater than the configured rate, we increase the configured rate. This condition can be expressed as, before triggering a new Flow to obtain new messages, the number of messages that the current consumer has received (i.e., current available permits + current messages in `queueCh`) is greater than the current queue size. 

- When is the timing to scale receiver queue size?

In the Java client, `expectMoreIncomingMessages` is triggered by "empty reveive".

https://github.com/apache/pulsar/blob/da30b2e211d0a097abfa8d0392f1fd4b13a46328/pulsar-client/src/main/java/org/apache/pulsar/client/impl/ConsumerImpl.java#L478-L482

In the Go Client, we can not get the "empty reveive" because the `messageCh` is exposed to user to async consume. So we `expectMoreIncomingMessages` when `messageCh` is not full.

https://github.com/apache/pulsar-client-go/blob/352c463194916cab7772c0729c9cdac14404e8bd/pulsar/consumer_partition.go#L1278-L1283

More details in #927 .

### Modifications

- Add `currentQueueSize` and `scaleReceiverQueueHint` to control `Flow` rate.
- Add `AutoScaledReceiverQueueSize` option for `Consumer`

### Verifying this change

- [x] Make sure that the change passes the CI checks.

Here is a test with pulsar-perf on local standalone server to confirm the effect of this PR.

Non-partitioned topic consumer

```
$  ./pulsar-perf produce my-test -r 1000 -s 128 -b 0 
$  ./pulsar-perf consume my-test -r 10000 --enable-auto-scaled-queue-size=true
```

| RATE  | Final receiver queue size |
| ----- | ------------------------- |
| 1     | 1                         |
| 5     | 2                         |
| 10    | 2                         |
| 100   | 8                         |
| 1000  | 64                        |
| 10000 | 1000                      |

3-partitioned topic consumer

```
$  ./pulsar-perf produce my-test-partitioned -r 1000 -s 128 -b 0 
$  ./pulsar-perf consume my-test-partitioned -r 10000 --enable-auto-scaled-queue-size=true
```

| RATE  | Final total receiver queue size | Sub-consumers receiver queue size |
| ----- | ------------------------------- | --------------------------------- |
| 1     | 4                               | 2,1,1                             |
| 10    | 8                               | 4,2,2                             |
| 100   | 16                              | 4,8,4                             |
| 1000  | 192                             | 64,64,64                          |
| 10000 | 1536                            | 512,512,512                       |
| 20000 | 3000                            | 1000,1000,1000                    |




Patch:
diff --git a/perf/perf-consumer.go b/perf/perf-consumer.go
index 6b6e411..825de62 100644
--- a/perf/perf-consumer.go
+++ b/perf/perf-consumer.go
@@ -31,10 +31,11 @@ import (
 
 // ConsumeArgs define the parameters required by consume
 type ConsumeArgs struct {
-	Topic               string
-	SubscriptionName    string
-	ReceiverQueueSize   int
-	EnableBatchIndexAck bool
+	Topic                             string
+	SubscriptionName                  string
+	ReceiverQueueSize                 int
+	EnableBatchIndexAck               bool
+	EnableAutoScaledReceiverQueueSize bool
 }
 
 func newConsumerCommand() *cobra.Command {
@@ -57,6 +58,8 @@ func newConsumerCommand() *cobra.Command {
 	flags.StringVarP(&consumeArgs.SubscriptionName, "subscription", "s", "sub", "Subscription name")
 	flags.IntVarP(&consumeArgs.ReceiverQueueSize, "receiver-queue-size", "r", 1000, "Receiver queue size")
 	flags.BoolVar(&consumeArgs.EnableBatchIndexAck, "enable-batch-index-ack", false, "Whether to enable batch index ACK")
+	flags.BoolVar(&consumeArgs.EnableAutoScaledReceiverQueueSize, "enable-auto-scaled-queue-size", false,
+		"Whether to enable auto scaled receiver queue size")
 
 	return cmd
 }
@@ -76,9 +79,10 @@ func consume(consumeArgs *ConsumeArgs, stop <-chan struct{}) {
 	defer client.Close()
 
 	consumer, err := client.Subscribe(pulsar.ConsumerOptions{
-		Topic:                          consumeArgs.Topic,
-		SubscriptionName:               consumeArgs.SubscriptionName,
-		EnableBatchIndexAcknowledgment: consumeArgs.EnableBatchIndexAck,
+		Topic:                             consumeArgs.Topic,
+		SubscriptionName:                  consumeArgs.SubscriptionName,
+		EnableBatchIndexAcknowledgment:    consumeArgs.EnableBatchIndexAck,
+		EnableAutoScaledReceiverQueueSize: consumeArgs.EnableAutoScaledReceiverQueueSize,
 	})
 
 	if err != nil {
diff --git a/pulsar/consumer.go b/pulsar/consumer.go
index 9576d7a..64a096d 100644
--- a/pulsar/consumer.go
+++ b/pulsar/consumer.go
@@ -160,6 +160,12 @@ type ConsumerOptions struct {
 	// Default value is `1000` messages and should be good for most use cases.
 	ReceiverQueueSize int
 
+	// EnableAutoScaledReceiverQueueSize, if enabled, the consumer receive queue will be auto-scaled
+	// by the consumer actual throughput. The ReceiverQueueSize will be the maximum size which consumer
+	// receive queue can be scaled.
+	// Default is false.
+	EnableAutoScaledReceiverQueueSize bool
+
 	// NackRedeliveryDelay specifies the delay after which to redeliver the messages that failed to be
 	// processed. Default is 1 min. (See `Consumer.Nack()`)
 	NackRedeliveryDelay time.Duration
diff --git a/pulsar/consumer_impl.go b/pulsar/consumer_impl.go
index d19e522..fd7fa57 100644
--- a/pulsar/consumer_impl.go
+++ b/pulsar/consumer_impl.go
@@ -399,6 +399,7 @@ func (c *consumer) internalTopicSubscribeToPartitions() error {
 				consumerEventListener:       c.options.EventListener,
 				enableBatchIndexAck:         c.options.EnableBatchIndexAcknowledgment,
 				ackGroupingOptions:          c.options.AckGroupingOptions,
+				autoReceiverQueueSize:       c.options.EnableAutoScaledReceiverQueueSize,
 			}
 			cons, err := newPartitionConsumer(c, c.client, opts, c.messageCh, c.dlq, c.metrics)
 			ch <- ConsumerError{
diff --git a/pulsar/consumer_partition.go b/pulsar/consumer_partition.go
index 18100a9..a08405a 100644
--- a/pulsar/consumer_partition.go
+++ b/pulsar/consumer_partition.go
@@ -25,7 +25,6 @@ import (
 	"math"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"google.golang.org/protobuf/proto"
@@ -77,6 +76,10 @@ const (
 	nonDurable
 )
 
+const (
+	initialReceiverQueueSize = 1
+)
+
 const (
 	noMessageEntry = -1
 )
@@ -89,6 +92,7 @@ type partitionConsumerOpts struct {
 	subscriptionInitPos         SubscriptionInitialPosition
 	partitionIdx                int
 	receiverQueueSize           int
+	autoReceiverQueueSize       bool
 	nackRedeliveryDelay         time.Duration
 	nackBackoffPolicy           NackBackoffPolicy
 	metadata                    map[string]string
@@ -142,11 +146,15 @@ type partitionConsumer struct {
 	availablePermits *availablePermits
 
 	// the size of the queue channel for buffering messages
-	queueSize       int32
+	maxQueueSize    int32
 	queueCh         chan []*message
 	startMessageID  atomicMessageID
 	lastDequeuedMsg *trackingMessageID
 
+	currentQueueSize       uAtomic.Int32
+	scaleReceiverQueueHint uAtomic.Bool
+	incomingMessages       uAtomic.Int32
+
 	eventsCh        chan interface{}
 	connectedCh     chan struct{}
 	connectClosedCh chan connectionClosed
@@ -181,22 +189,44 @@ func (pc *partitionConsumer) ActiveConsumerChanged(isActive bool) {
 }
 
 type availablePermits struct {
-	permits int32
+	permits uAtomic.Int32
 	pc      *partitionConsumer
 }
 
 func (p *availablePermits) inc() {
 	// atomic add availablePermits
-	ap := atomic.AddInt32(&p.permits, 1)
+	p.add(1)
+}
+
+func (p *availablePermits) add(delta int32) {
+	p.permits.Add(delta)
+	p.flowIfNeed()
+}
+
+func (p *availablePermits) reset() {
+	p.permits.Store(0)
+}
 
+func (p *availablePermits) get() int32 {
+	return p.permits.Load()
+}
+
+func (p *availablePermits) flowIfNeed() {
 	// TODO implement a better flow controller
 	// send more permits if needed
-	flowThreshold := int32(math.Max(float64(p.pc.queueSize/2), 1))
-	if ap >= flowThreshold {
-		availablePermits := ap
-		requestedPermits := ap
+	var flowThreshold int32
+	if p.pc.options.autoReceiverQueueSize {
+		flowThreshold = int32(math.Max(float64(p.pc.currentQueueSize.Load()/2), 1))
+	} else {
+		flowThreshold = int32(math.Max(float64(p.pc.maxQueueSize/2), 1))
+	}
+
+	current := p.get()
+	if current >= flowThreshold {
+		availablePermits := current
+		requestedPermits := current
 		// check if permits changed
-		if !atomic.CompareAndSwapInt32(&p.permits, ap, 0) {
+		if !p.permits.CAS(current, 0) {
 			return
 		}
 
@@ -207,10 +237,6 @@ func (p *availablePermits) inc() {
 	}
 }
 
-func (p *availablePermits) reset() {
-	atomic.StoreInt32(&p.permits, 0)
-}
-
 // atomicMessageID is a wrapper for trackingMessageID to make get and set atomic
 type atomicMessageID struct {
 	msgID *trackingMessageID
@@ -292,7 +318,7 @@ func newPartitionConsumer(parent Consumer, client *client, options *partitionCon
 		consumerID:           client.rpcClient.NewConsumerID(),
 		partitionIdx:         int32(options.partitionIdx),
 		eventsCh:             make(chan interface{}, 10),
-		queueSize:            int32(options.receiverQueueSize),
+		maxQueueSize:         int32(options.receiverQueueSize),
 		queueCh:              make(chan []*message, options.receiverQueueSize),
 		startMessageID:       atomicMessageID{msgID: options.startMessageID},
 		connectedCh:          make(chan struct{}),
@@ -305,6 +331,11 @@ func newPartitionConsumer(parent Consumer, client *client, options *partitionCon
 		metrics:              metrics,
 		schemaInfoCache:      newSchemaInfoCache(client, options.topic),
 	}
+	if pc.options.autoReceiverQueueSize {
+		pc.currentQueueSize.Store(initialReceiverQueueSize)
+	} else {
+		pc.currentQueueSize.Store(int32(pc.options.receiverQueueSize))
+	}
 	pc.availablePermits = &availablePermits{pc: pc}
 	pc.chunkedMsgCtxMap = newChunkedMsgCtxMap(options.maxPendingChunkedMessage, pc)
 	pc.unAckChunksTracker = newUnAckChunksTracker(pc)
@@ -904,6 +935,12 @@ func (pc *partitionConsumer) MessageReceived(response *pb.CommandMessage, header
 					orderingKey:         string(msgMeta.OrderingKey),
 				},
 			}
+
+			if pc.options.autoReceiverQueueSize {
+				pc.incomingMessages.Inc()
+				pc.markScaleIfNeed()
+			}
+
 			pc.queueCh <- messages
 			return nil
 		}
@@ -1073,6 +1110,11 @@ func (pc *partitionConsumer) MessageReceived(response *pb.CommandMessage, header
 		messages = append(messages, msg)
 	}
 
+	if pc.options.autoReceiverQueueSize {
+		pc.incomingMessages.Add(int32(len(messages)))
+		pc.markScaleIfNeed()
+	}
+
 	// send messages to the dispatcher
 	pc.queueCh <- messages
 	return nil
@@ -1240,7 +1282,6 @@ func (pc *partitionConsumer) dispatcher() {
 			pc.metrics.PrefetchedMessages.Dec()
 			pc.metrics.PrefetchedBytes.Sub(float64(len(messages[0].payLoad)))
 		} else {
-			// we are ready for more messages
 			queueCh = pc.queueCh
 		}
 
@@ -1258,7 +1299,13 @@ func (pc *partitionConsumer) dispatcher() {
 
 			// reset available permits
 			pc.availablePermits.reset()
-			initialPermits := uint32(pc.queueSize)
+
+			var initialPermits uint32
+			if pc.options.autoReceiverQueueSize {
+				initialPermits = uint32(pc.currentQueueSize.Load())
+			} else {
+				initialPermits = uint32(pc.maxQueueSize)
+			}
 
 			pc.log.Debugf("dispatcher requesting initial permits=%d", initialPermits)
 			// send initial permits
@@ -1282,6 +1329,11 @@ func (pc *partitionConsumer) dispatcher() {
 
 			pc.availablePermits.inc()
 
+			if pc.options.autoReceiverQueueSize {
+				pc.incomingMessages.Dec()
+				pc.expectMoreIncomingMessages()
+			}
+
 		case clearQueueCb := <-pc.clearQueueCh:
 			// drain the message queue on any new connection by sending a
 			// special nil message to the channel so we know when to stop dropping messages
@@ -1297,6 +1349,9 @@ func (pc *partitionConsumer) dispatcher() {
 				} else if nextMessageInQueue == nil {
 					nextMessageInQueue = toTrackingMessageID(m[0].msgID)
 				}
+				if pc.options.autoReceiverQueueSize {
+					pc.incomingMessages.Sub(int32(len(m)))
+				}
 			}
 
 			messages = nil
@@ -1671,6 +1726,32 @@ func getPreviousMessage(mid *trackingMessageID) *trackingMessageID {
 	}
 }
 
+func (pc *partitionConsumer) expectMoreIncomingMessages() {
+	if !pc.options.autoReceiverQueueSize {
+		return
+	}
+	if pc.scaleReceiverQueueHint.CAS(true, false) {
+		oldSize := pc.currentQueueSize.Load()
+		maxSize := int32(pc.options.receiverQueueSize)
+		newSize := int32(math.Min(float64(maxSize), float64(oldSize*2)))
+		if newSize > oldSize {
+			pc.currentQueueSize.CAS(oldSize, newSize)
+			pc.availablePermits.add(newSize - oldSize)
+			pc.log.Debugf("update currentQueueSize from %d -> %d", oldSize, newSize)
+		}
+	}
+}
+
+func (pc *partitionConsumer) markScaleIfNeed() {
+	// availablePermits + incomingMessages (messages in queueCh) is the number of prefetched messages
+	// The result of auto-scale we expected is currentQueueSize is slightly bigger than prefetched messages
+	prev := pc.scaleReceiverQueueHint.Swap(pc.availablePermits.get()+pc.incomingMessages.Load() >=
+		pc.currentQueueSize.Load())
+	if prev != pc.scaleReceiverQueueHint.Load() {
+		pc.log.Debugf("update scaleReceiverQueueHint from %t -> %t", prev, pc.scaleReceiverQueueHint.Load())
+	}
+}
+
 func (pc *partitionConsumer) Decompress(msgMeta *pb.MessageMetadata, payload internal.Buffer) (internal.Buffer, error) {
 	providerEntry, ok := pc.compressionProviders.Load(msgMeta.GetCompression())
 	if !ok {
diff --git a/pulsar/dlq_router.go b/pulsar/dlq_router.go
index 5ecd8f8..6d30003 100644
--- a/pulsar/dlq_router.go
+++ b/pulsar/dlq_router.go
@@ -109,13 +109,17 @@ func (r *dlqRouter) run() {
 				Properties:          properties,
 				EventTime:           msg.EventTime(),
 				ReplicationClusters: msg.replicationClusters,
-			}, func(MessageID, *ProducerMessage, error) {
-				r.log.WithField("msgID", msgID).Debug("Sent message to DLQ")
-
-				// The Producer ack might be coming from the connection go-routine that
-				// is also used by the consumer. In that case we would get a dead-lock
-				// if we'd try to ack.
-				go cm.Consumer.AckID(msgID)
+			}, func(messageID MessageID, producerMessage *ProducerMessage, err error) {
+				if err == nil {
+					r.log.WithField("msgID", msgID).Debug("Succeed to send message to DLQ")
+					// The Producer ack might be coming from the connection go-routine that
+					// is also used by the consumer. In that case we would get a dead-lock
+					// if we'd try to ack.
+					go cm.Consumer.AckID(msgID)
+				} else {
+					r.log.WithError(err).WithField("msgID", msgID).Debug("Failed to send message to DLQ")
+					go cm.Consumer.Nack(cm)
+				}
 			})
 
 		case <-r.closeCh:


Test patch:
diff --git a/pulsar/consumer_partition_test.go b/pulsar/consumer_partition_test.go
index 21bc0e3..21280fb 100644
--- a/pulsar/consumer_partition_test.go
+++ b/pulsar/consumer_partition_test.go
@@ -37,6 +37,7 @@ func TestSingleMessageIDNoAckTracker(t *testing.T) {
 		metrics:              newTestMetrics(),
 		decryptor:            crypto.NewNoopDecryptor(),
 	}
+	pc.availablePermits = &availablePermits{pc: &pc}
 	pc.ackGroupingTracker = newAckGroupingTracker(&AckGroupingOptions{MaxSize: 0},
 		func(id MessageID) { pc.sendIndividualAck(id) }, nil, nil)
 
@@ -75,6 +76,7 @@ func TestBatchMessageIDNoAckTracker(t *testing.T) {
 		metrics:              newTestMetrics(),
 		decryptor:            crypto.NewNoopDecryptor(),
 	}
+	pc.availablePermits = &availablePermits{pc: &pc}
 	pc.ackGroupingTracker = newAckGroupingTracker(&AckGroupingOptions{MaxSize: 0},
 		func(id MessageID) { pc.sendIndividualAck(id) }, nil, nil)
 
@@ -110,6 +112,7 @@ func TestBatchMessageIDWithAckTracker(t *testing.T) {
 		metrics:              newTestMetrics(),
 		decryptor:            crypto.NewNoopDecryptor(),
 	}
+	pc.availablePermits = &availablePermits{pc: &pc}
 	pc.ackGroupingTracker = newAckGroupingTracker(&AckGroupingOptions{MaxSize: 0},
 		func(id MessageID) { pc.sendIndividualAck(id) }, nil, nil)
 
diff --git a/pulsar/consumer_test.go b/pulsar/consumer_test.go
index 21fa7d0..a93dea9 100644
--- a/pulsar/consumer_test.go
+++ b/pulsar/consumer_test.go
@@ -3988,3 +3988,88 @@ func runBatchIndexAckTest(t *testing.T, ackWithResponse bool, cumulative bool, o
 
 	client.Close()
 }
+
+func TestConsumerWithAutoScaledQueueReceive(t *testing.T) {
+	client, err := NewClient(ClientOptions{
+		URL: lookupURL,
+	})
+
+	assert.Nil(t, err)
+	defer client.Close()
+
+	topic := newTopicName()
+
+	// create consumer
+	c, err := client.Subscribe(ConsumerOptions{
+		Topic:                             topic,
+		SubscriptionName:                  "my-sub",
+		Type:                              Exclusive,
+		ReceiverQueueSize:                 3,
+		EnableAutoScaledReceiverQueueSize: true,
+	})
+	assert.Nil(t, err)
+	pc := c.(*consumer).consumers[0]
+	assert.Equal(t, int32(1), pc.currentQueueSize.Load())
+	defer c.Close()
+
+	// create p
+	p, err := client.CreateProducer(ProducerOptions{
+		Topic:           topic,
+		DisableBatching: false,
+	})
+	assert.Nil(t, err)
+	defer p.Close()
+
+	// send message, it will update scaleReceiverQueueHint from false to true
+	_, err = p.Send(context.Background(), &ProducerMessage{
+		Payload: []byte("hello"),
+	})
+	assert.NoError(t, err)
+
+	// this will trigger receiver queue size expanding to 2 because we have prefetched 1 message >= currentSize 1.
+	_, err = c.Receive(context.Background())
+	assert.Nil(t, err)
+
+	// currentQueueSize should be doubled in size
+	retryAssert(t, 5, 200, func() {}, func(t assert.TestingT) bool {
+		return assert.Equal(t, 2, int(pc.currentQueueSize.Load()))
+	})
+
+	for i := 0; i < 5; i++ {
+		_, err = p.Send(context.Background(), &ProducerMessage{
+			Payload: []byte("hello"),
+		})
+		assert.NoError(t, err)
+
+		// waiting for prefetched message passing from queueCh to messageCh
+		retryAssert(t, 5, 200, func() {}, func(t assert.TestingT) bool {
+			return assert.Equal(t, 1, len(pc.messageCh))
+		})
+
+		_, err = p.Send(context.Background(), &ProducerMessage{
+			Payload: []byte("hello"),
+		})
+		assert.NoError(t, err)
+
+		// wait all the messages has been prefetched
+		_, err = c.Receive(context.Background())
+		assert.Nil(t, err)
+		_, err = c.Receive(context.Background())
+		assert.Nil(t, err)
+		// this will not trigger receiver queue size expanding because we have prefetched 2 message < currentSize 4.
+		assert.Equal(t, int32(2), pc.currentQueueSize.Load())
+	}
+
+	for i := 0; i < 5; i++ {
+		p.SendAsync(
+			context.Background(),
+			&ProducerMessage{Payload: []byte("hello")},
+			func(id MessageID, producerMessage *ProducerMessage, err error) {
+			},
+		)
+	}
+
+	retryAssert(t, 3, 300, func() {}, func(t assert.TestingT) bool {
+		return assert.Equal(t, 3, int(pc.currentQueueSize.Load()))
+	})
+}