Repository: apache/pulsar-client-go
Pull request #988 (2023-03-09 05:38:49)

Optimize batch index ACK performance
### Motivation

Currently, when `EnableBatchIndexAck` is true, the ACK performance is very poor. There are two main reasons:
1. Acknowledgment by list is not supported. It means that even N MessageIDs are grouped, there are still N ACK requests to send.
2. The implementation of ACK grouping tracker is wrong. Give a batch that has N messages, when batch index ACK is enabled, each MessageID is cached. However, after all these N MessageIDs arrived, the current implementation does not clear them.

### Modifications
- Add a `func(id []*pb.MessageIdData)` to the ACK grouping tracker. When flushing individual ACKs, construct the slice and wrap the slice to `CommandAck` directly.
- Refactor the implementation of the ACK grouping tracker to have these two data structures:
  - `pendingAcks`: Cache the non-batched MessageIDs
  - `pendingBatchAcks`: Cache the batched MessageIDs, once all messages in a batch have been added, remove them from it and add the non-batched MessageID to `pendingAcks`
- Add `TestBatchIndexAckAllMessages` to verify the new behavior.

After this change, the ACK order cannot be guaranteed, sort the acknowledged MessageIDs in the `ack_grouping_tracker_test.go`.

BewareMyPower: ## Performance test setup

Run a Pulsar 2.11.0 standalone locally with `acknowledgmentAtBatchIndexLevelEnabled=true`.

Build the `perf` binary:

```bash
cd perf
go build
```

Run the consumer in terminal 1:

```bash
./perf consume --profile \
    -s sub \
    --enable-batch-index-ack \
    my-topic
```

Run the producer in terminal, the throughput is about 20 MB/s:

```bash
./perf produce \
  --batching-num-messages=10 \
  --size 105 \
  --rate 200000 \
  my-topic
```

## Test result

### Before this patch

```
INFO[13:31:17.584] Stats - Publish rate: 150038.6 msg/s -  120.2 Mbps -
                                Latency ms: 50%   3.1 -95%  13.0 - 99%  48.9 - 99.9%  48.9 - max   48.9
INFO[13:31:27.581] Stats - Publish rate: 179909.8 msg/s -  144.1 Mbps -
                                Latency ms: 50%   2.4 -95%   8.7 - 99%  26.2 - 99.9%  26.2 - max   26.2
INFO[13:31:37.581] Stats - Publish rate: 188957.1 msg/s -  151.4 Mbps -
                                Latency ms: 50%   2.3 -95%   7.8 - 99%  15.9 - 99.9%  15.9 - max   15.9
```

```
INFO[13:31:11.592] Stats - Consume rate: 22956.0 msg/s -   18.4 Mbps
INFO[13:31:21.592] Stats - Consume rate: 46805.3 msg/s -   37.5 Mbps
INFO[13:31:31.593] Stats - Consume rate: 47494.7 msg/s -   38.0 Mbps
INFO[13:31:41.594] Stats - Consume rate: 48100.3 msg/s -   38.5 Mbps
INFO[13:31:51.592] Stats - Consume rate: 59705.8 msg/s -   47.8 Mbps
INFO[13:32:01.597] Stats - Consume rate: 61694.0 msg/s -   49.4 Mbps
INFO[13:32:11.592] Stats - Consume rate: 58100.2 msg/s -   46.5 Mbps
INFO[13:32:21.592] Stats - Consume rate: 60818.2 msg/s -   48.7 Mbps
INFO[13:32:31.596] Stats - Consume rate: 60681.5 msg/s -   48.6 Mbps
```

The consumer can not catch up the producer even if the produce rate is only 20 MB/s. So it fell back to catch up read and the broker reads messages directly with BK.

### After this patch

```
INFO[13:27:33.212] Stats - Publish rate: 165290.5 msg/s -  132.4 Mbps -
                                Latency ms: 50%   2.9 -95%  11.3 - 99%  58.2 - 99.9%  58.2 - max   58.2
INFO[13:27:43.213] Stats - Publish rate: 197230.3 msg/s -  158.0 Mbps -
                                Latency ms: 50%   2.4 -95%   5.5 - 99%  22.2 - 99.9%  22.2 - max   22.2
INFO[13:27:53.212] Stats - Publish rate: 193150.6 msg/s -  154.7 Mbps -
                                Latency ms: 50%   2.5 -95%   5.9 - 99%  17.6 - 99.9%  17.6 - max   17.6
INFO[13:28:03.212] Stats - Publish rate: 170759.9 msg/s -  136.8 Mbps -
                                Latency ms: 50%   2.6 -95%   5.7 - 99%  23.4 - 99.9%  23.4 - max   23.4
INFO[13:28:13.212] Stats - Publish rate: 149780.3 msg/s -  120.0 Mbps -
                                Latency ms: 50%   2.9 -95%   6.4 - 99%  61.2 - 99.9%  61.2 - max   61.2
INFO[13:28:23.212] Stats - Publish rate: 157285.9 msg/s -  126.0 Mbps -
                                Latency ms: 50%   2.7 -95%   6.2 - 99%  14.2 - 99.9%  14.2 - max   14.2
```

```
INFO[13:27:34.523] Stats - Consume rate: 180047.6 msg/s -  144.2 Mbps
INFO[13:27:44.524] Stats - Consume rate: 197120.2 msg/s -  157.9 Mbps
INFO[13:27:54.523] Stats - Consume rate: 193490.8 msg/s -  155.0 Mbps
INFO[13:28:04.523] Stats - Consume rate: 165149.2 msg/s -  132.3 Mbps
INFO[13:28:14.523] Stats - Consume rate: 150062.6 msg/s -  120.2 Mbps
INFO[13:28:24.523] Stats - Consume rate: 154697.5 msg/s -  123.9 Mbps
```

### Flame Graph

Before this patch:

![image](https://user-images.githubusercontent.com/18204803/223931561-3b71c01a-5305-4cef-93ce-7977b0021994.png)

![image](https://user-images.githubusercontent.com/18204803/223931801-71892bde-0a2a-4ef0-8b42-364783dd97d8.png)

We can see `internalSendRequest` takes about 40% of the total, which also leads to the result that `runEventLoop` method also takes 8.87% of the total.

After this patch:

![image](https://user-images.githubusercontent.com/18204803/223932285-1e8adb5e-02b4-43a7-9ba9-beb854706fe3.png)

Patch:
diff --git a/pulsar/ack_grouping_tracker.go b/pulsar/ack_grouping_tracker.go
index 22aba11..dd05186 100644
--- a/pulsar/ack_grouping_tracker.go
+++ b/pulsar/ack_grouping_tracker.go
@@ -19,8 +19,10 @@ package pulsar
 
 import (
 	"sync"
+	"sync/atomic"
 	"time"
 
+	pb "github.com/apache/pulsar-client-go/pulsar/internal/pulsar_proto"
 	"github.com/bits-and-blooms/bitset"
 )
 
@@ -40,7 +42,8 @@ type ackGroupingTracker interface {
 
 func newAckGroupingTracker(options *AckGroupingOptions,
 	ackIndividual func(id MessageID),
-	ackCumulative func(id MessageID)) ackGroupingTracker {
+	ackCumulative func(id MessageID),
+	ackList func(ids []*pb.MessageIdData)) ackGroupingTracker {
 	if options == nil {
 		options = &AckGroupingOptions{
 			MaxSize: 1000,
@@ -56,38 +59,28 @@ func newAckGroupingTracker(options *AckGroupingOptions,
 	}
 
 	t := &timedAckGroupingTracker{
-		singleAcks:        make([]MessageID, options.MaxSize),
-		pendingAcks:       make(map[int64]*bitset.BitSet),
-		lastCumulativeAck: EarliestMessageID(),
-		ackIndividual:     ackIndividual,
+		maxNumAcks:        int(options.MaxSize),
 		ackCumulative:     ackCumulative,
-		ackList: func(ids []MessageID) {
-			// TODO: support ack a list of MessageIDs
-			for _, id := range ids {
-				ackIndividual(id)
-			}
-		},
-		options: *options,
-		tick:    time.NewTicker(time.Hour),
-		donCh:   make(chan struct{}),
+		ackList:           ackList,
+		pendingAcks:       make(map[[2]uint64]*bitset.BitSet),
+		lastCumulativeAck: EarliestMessageID(),
 	}
 
 	if options.MaxTime > 0 {
-		t.tick = time.NewTicker(options.MaxTime)
-	} else {
-		t.tick.Stop()
+		t.ticker = time.NewTicker(options.MaxTime)
+		t.exitCh = make(chan struct{})
+		go func() {
+			for {
+				select {
+				case <-t.exitCh:
+					return
+				case <-t.ticker.C:
+					t.flush()
+				}
+			}
+		}()
 	}
 
-	go func() {
-		for {
-			select {
-			case <-t.donCh:
-				return
-			case <-t.tick.C:
-				t.flush()
-			}
-		}
-	}()
 	return t
 }
 
@@ -117,157 +110,157 @@ func (i *immediateAckGroupingTracker) flushAndClean() {
 func (i *immediateAckGroupingTracker) close() {
 }
 
-func (t *timedAckGroupingTracker) addAndCheckIfFull(id MessageID) bool {
-	t.mutex.Lock()
-	defer t.mutex.Unlock()
-	t.singleAcks[t.index] = id
-	t.index++
-	key := messageIDHash(id)
-	ackSet, found := t.pendingAcks[key]
-	if !found {
-		if messageIDIsBatch(id) {
-			ackSet = bitset.New(uint(id.BatchSize()))
-			for i := 0; i < int(id.BatchSize()); i++ {
-				ackSet.Set(uint(i))
-			}
-			t.pendingAcks[key] = ackSet
-		} else {
-			t.pendingAcks[key] = nil
-		}
-	}
-	if ackSet != nil {
-		ackSet.Clear(uint(id.BatchIdx()))
-	}
-	return t.index == len(t.singleAcks)
-}
-
-func (t *timedAckGroupingTracker) tryUpdateLastCumulativeAck(id MessageID) {
-	t.mutex.Lock()
-	defer t.mutex.Unlock()
-	if messageIDCompare(t.lastCumulativeAck, id) < 0 {
-		t.lastCumulativeAck = id
-		t.cumulativeAckRequired = true
-	}
-}
-
-func (t *timedAckGroupingTracker) flushIndividualAcks() {
-	t.mutex.Lock()
-	defer t.mutex.Unlock()
-	if t.index > 0 {
-		t.ackList(t.singleAcks[0:t.index])
-		for _, id := range t.singleAcks[0:t.index] {
-			key := messageIDHash(id)
-			ackSet, found := t.pendingAcks[key]
-			if !found {
-				continue
-			}
-			if ackSet == nil {
-				delete(t.pendingAcks, key)
-			} else {
-				ackSet.Clear(uint(id.BatchIdx()))
-				if ackSet.None() { // all messages have been acknowledged
-					delete(t.pendingAcks, key)
-				}
-			}
-			delete(t.pendingAcks, messageIDHash(id))
-		}
-		t.index = 0
-	}
-}
-
-func (t *timedAckGroupingTracker) flushCumulativeAck() {
-	t.mutex.Lock()
-	defer t.mutex.Unlock()
-	if t.cumulativeAckRequired {
-		t.ackCumulative(t.lastCumulativeAck)
-		t.cumulativeAckRequired = false
-	}
-}
-
-func (t *timedAckGroupingTracker) clean() {
-	t.mutex.Lock()
-	defer t.mutex.Unlock()
-	maxSize := len(t.singleAcks)
-	t.singleAcks = make([]MessageID, maxSize)
-	t.index = 0
-	t.pendingAcks = make(map[int64]*bitset.BitSet)
-	t.lastCumulativeAck = EarliestMessageID()
-	t.cumulativeAckRequired = false
-}
-
 type timedAckGroupingTracker struct {
-	singleAcks []MessageID
-	index      int
+	sync.RWMutex
 
-	// Key is the hash code of the ledger id and the netry id,
+	maxNumAcks    int
+	ackCumulative func(id MessageID)
+	ackList       func(ids []*pb.MessageIdData)
+	ticker        *time.Ticker
+
+	// Key is the pair of the ledger id and the entry id,
 	// Value is the bit set that represents which messages are acknowledged if the entry stores a batch.
 	// The bit 1 represents the message has been acknowledged, i.e. the bits "111" represents all messages
 	// in the batch whose batch size is 3 are not acknowledged.
 	// After the 1st message (i.e. batch index is 0) is acknowledged, the bits will become "011".
 	// Value is nil if the entry represents a single message.
-	pendingAcks map[int64]*bitset.BitSet
+	pendingAcks map[[2]uint64]*bitset.BitSet
 
 	lastCumulativeAck     MessageID
-	cumulativeAckRequired bool
-
-	ackIndividual func(id MessageID)
-	ackCumulative func(id MessageID)
-	ackList       func(ids []MessageID)
-
-	options AckGroupingOptions
-	donCh   chan struct{}
-	tick    *time.Ticker
+	cumulativeAckRequired int32
 
-	mutex sync.RWMutex
+	exitCh chan struct{}
 }
 
 func (t *timedAckGroupingTracker) add(id MessageID) {
-	if t.addAndCheckIfFull(id) {
-		t.flushIndividualAcks()
-		if t.options.MaxTime > 0 {
-			t.tick.Reset(t.options.MaxTime)
+	if acks := t.tryAddIndividual(id); acks != nil {
+		t.flushIndividual(acks)
+	}
+}
+
+func (t *timedAckGroupingTracker) tryAddIndividual(id MessageID) map[[2]uint64]*bitset.BitSet {
+	t.Lock()
+	defer t.Unlock()
+	key := [2]uint64{uint64(id.LedgerID()), uint64(id.EntryID())}
+
+	batchIdx := id.BatchIdx()
+	batchSize := id.BatchSize()
+
+	if batchIdx >= 0 && batchSize > 0 {
+		bs, found := t.pendingAcks[key]
+		if !found {
+			if batchSize > 1 {
+				bs = bitset.New(uint(batchSize))
+				for i := uint(0); i < uint(batchSize); i++ {
+					bs.Set(i)
+				}
+			}
+			t.pendingAcks[key] = bs
 		}
+		if bs != nil {
+			bs.Clear(uint(batchIdx))
+		}
+	} else {
+		t.pendingAcks[key] = nil
 	}
+
+	if len(t.pendingAcks) >= t.maxNumAcks {
+		pendingAcks := t.pendingAcks
+		t.pendingAcks = make(map[[2]uint64]*bitset.BitSet)
+		return pendingAcks
+	}
+	return nil
 }
 
 func (t *timedAckGroupingTracker) addCumulative(id MessageID) {
-	t.tryUpdateLastCumulativeAck(id)
-	if t.options.MaxTime <= 0 {
-		t.flushCumulativeAck()
+	if t.tryUpdateCumulative(id) && t.ticker == nil {
+		t.ackCumulative(id)
 	}
 }
 
+func (t *timedAckGroupingTracker) tryUpdateCumulative(id MessageID) bool {
+	t.Lock()
+	defer t.Unlock()
+	if messageIDCompare(t.lastCumulativeAck, id) < 0 {
+		t.lastCumulativeAck = id
+		atomic.StoreInt32(&t.cumulativeAckRequired, 1)
+		return true
+	}
+	return false
+}
+
 func (t *timedAckGroupingTracker) isDuplicate(id MessageID) bool {
-	t.mutex.RLock()
+	t.RLock()
+	defer t.RUnlock()
 	if messageIDCompare(t.lastCumulativeAck, id) >= 0 {
-		t.mutex.RUnlock()
 		return true
 	}
-	ackSet, found := t.pendingAcks[messageIDHash(id)]
-	if !found {
-		t.mutex.RUnlock()
-		return false
-	}
-	t.mutex.RUnlock()
-	if ackSet == nil || !messageIDIsBatch(id) {
-		// NOTE: should we panic when ackSet != nil and messageIDIsBatch(id) is true?
-		return true
+	key := [2]uint64{uint64(id.LedgerID()), uint64(id.EntryID())}
+	if bs, found := t.pendingAcks[key]; found {
+		if bs == nil {
+			return true
+		}
+		if !bs.Test(uint(id.BatchIdx())) {
+			return true
+		}
 	}
-	// 0 represents the message has been acknowledged
-	return !ackSet.Test(uint(id.BatchIdx()))
+	return false
 }
 
 func (t *timedAckGroupingTracker) flush() {
-	t.flushIndividualAcks()
-	t.flushCumulativeAck()
+	if acks := t.clearPendingAcks(); len(acks) > 0 {
+		t.flushIndividual(acks)
+	}
+	if atomic.CompareAndSwapInt32(&t.cumulativeAckRequired, 1, 0) {
+		t.RLock()
+		id := t.lastCumulativeAck
+		t.RUnlock()
+		t.ackCumulative(id)
+	}
 }
 
 func (t *timedAckGroupingTracker) flushAndClean() {
-	t.flush()
-	t.clean()
+	if acks := t.clearPendingAcks(); len(acks) > 0 {
+		t.flushIndividual(acks)
+	}
+	if atomic.CompareAndSwapInt32(&t.cumulativeAckRequired, 1, 0) {
+		t.Lock()
+		id := t.lastCumulativeAck
+		t.lastCumulativeAck = EarliestMessageID()
+		t.Unlock()
+		t.ackCumulative(id)
+	}
+}
+
+func (t *timedAckGroupingTracker) clearPendingAcks() map[[2]uint64]*bitset.BitSet {
+	t.Lock()
+	defer t.Unlock()
+	pendingAcks := t.pendingAcks
+	t.pendingAcks = make(map[[2]uint64]*bitset.BitSet)
+	return pendingAcks
 }
 
 func (t *timedAckGroupingTracker) close() {
 	t.flushAndClean()
-	close(t.donCh)
+	if t.exitCh != nil {
+		close(t.exitCh)
+	}
+}
+
+func (t *timedAckGroupingTracker) flushIndividual(pendingAcks map[[2]uint64]*bitset.BitSet) {
+	msgIDs := make([]*pb.MessageIdData, 0, len(pendingAcks))
+	for k, v := range pendingAcks {
+		ledgerID := k[0]
+		entryID := k[1]
+		msgID := &pb.MessageIdData{LedgerId: &ledgerID, EntryId: &entryID}
+		if v != nil && !v.None() {
+			bytes := v.Bytes()
+			msgID.AckSet = make([]int64, len(bytes))
+			for i := 0; i < len(bytes); i++ {
+				msgID.AckSet[i] = int64(bytes[i])
+			}
+		}
+		msgIDs = append(msgIDs, msgID)
+	}
+	t.ackList(msgIDs)
 }
diff --git a/pulsar/consumer_partition.go b/pulsar/consumer_partition.go
index a3dac19..18100a9 100644
--- a/pulsar/consumer_partition.go
+++ b/pulsar/consumer_partition.go
@@ -310,7 +310,8 @@ func newPartitionConsumer(parent Consumer, client *client, options *partitionCon
 	pc.unAckChunksTracker = newUnAckChunksTracker(pc)
 	pc.ackGroupingTracker = newAckGroupingTracker(options.ackGroupingOptions,
 		func(id MessageID) { pc.sendIndividualAck(id) },
-		func(id MessageID) { pc.sendCumulativeAck(id) })
+		func(id MessageID) { pc.sendCumulativeAck(id) },
+		func(ids []*pb.MessageIdData) { pc.eventsCh <- ids })
 	pc.setConsumerState(consumerInit)
 	pc.log = client.log.SubLogger(log.Fields{
 		"name":         pc.name,
@@ -837,6 +838,14 @@ func (pc *partitionConsumer) internalAck(req *ackRequest) {
 	}
 }
 
+func (pc *partitionConsumer) internalAckList(msgIDs []*pb.MessageIdData) {
+	pc.client.rpcClient.RequestOnCnxNoWait(pc._getConn(), pb.BaseCommand_ACK, &pb.CommandAck{
+		AckType:    pb.CommandAck_Individual.Enum(),
+		ConsumerId: proto.Uint64(pc.consumerID),
+		MessageId:  msgIDs,
+	})
+}
+
 func (pc *partitionConsumer) MessageReceived(response *pb.CommandMessage, headersAndPayload internal.Buffer) error {
 	pbMsgID := response.GetMessageId()
 
@@ -1364,6 +1373,8 @@ func (pc *partitionConsumer) runEventsLoop() {
 			switch v := i.(type) {
 			case *ackRequest:
 				pc.internalAck(v)
+			case []*pb.MessageIdData:
+				pc.internalAckList(v)
 			case *redeliveryRequest:
 				pc.internalRedeliver(v)
 			case *unsubscribeRequest:
diff --git a/pulsar/message.go b/pulsar/message.go
index c44957d..98190e9 100644
--- a/pulsar/message.go
+++ b/pulsar/message.go
@@ -212,11 +212,3 @@ func messageIDCompare(lhs MessageID, rhs MessageID) int {
 	}
 	return 0
 }
-
-func messageIDHash(id MessageID) int64 {
-	return id.LedgerID() + 31*id.EntryID()
-}
-
-func messageIDIsBatch(id MessageID) bool {
-	return id.BatchIdx() >= 0 && id.BatchSize() > 0
-}
diff --git a/pulsar/reader.go b/pulsar/reader.go
index e4679ab..d58d06f 100644
--- a/pulsar/reader.go
+++ b/pulsar/reader.go
@@ -92,6 +92,16 @@ type ReaderOptions struct {
 	// BackoffPolicy parameterize the following options in the reconnection logic to
 	// allow users to customize the reconnection logic (minBackoff, maxBackoff and jitterPercentage)
 	BackoffPolicy internal.BackoffPolicy
+
+	// MaxPendingChunkedMessage sets the maximum pending chunked messages. (default: 100)
+	MaxPendingChunkedMessage int
+
+	// ExpireTimeOfIncompleteChunk sets the expiry time of discarding incomplete chunked message. (default: 60 seconds)
+	ExpireTimeOfIncompleteChunk time.Duration
+
+	// AutoAckIncompleteChunk sets whether reader auto acknowledges incomplete chunked message when it should
+	// be removed (e.g.the chunked message pending queue is full). (default: false)
+	AutoAckIncompleteChunk bool
 }
 
 // Reader can be used to scan through all the messages currently available in a topic.
diff --git a/pulsar/reader_impl.go b/pulsar/reader_impl.go
index c7620ad..36b492a 100644
--- a/pulsar/reader_impl.go
+++ b/pulsar/reader_impl.go
@@ -90,22 +90,33 @@ func newReader(client *client, options ReaderOptions) (Reader, error) {
 		options.Decryption.MessageCrypto = messageCrypto
 	}
 
+	if options.MaxPendingChunkedMessage == 0 {
+		options.MaxPendingChunkedMessage = 100
+	}
+
+	if options.ExpireTimeOfIncompleteChunk == 0 {
+		options.ExpireTimeOfIncompleteChunk = time.Minute
+	}
+
 	consumerOptions := &partitionConsumerOpts{
-		topic:                      options.Topic,
-		consumerName:               options.Name,
-		subscription:               subscriptionName,
-		subscriptionType:           Exclusive,
-		receiverQueueSize:          receiverQueueSize,
-		startMessageID:             startMessageID,
-		startMessageIDInclusive:    options.StartMessageIDInclusive,
-		subscriptionMode:           nonDurable,
-		readCompacted:              options.ReadCompacted,
-		metadata:                   options.Properties,
-		nackRedeliveryDelay:        defaultNackRedeliveryDelay,
-		replicateSubscriptionState: false,
-		decryption:                 options.Decryption,
-		schema:                     options.Schema,
-		backoffPolicy:              options.BackoffPolicy,
+		topic:                       options.Topic,
+		consumerName:                options.Name,
+		subscription:                subscriptionName,
+		subscriptionType:            Exclusive,
+		receiverQueueSize:           receiverQueueSize,
+		startMessageID:              startMessageID,
+		startMessageIDInclusive:     options.StartMessageIDInclusive,
+		subscriptionMode:            nonDurable,
+		readCompacted:               options.ReadCompacted,
+		metadata:                    options.Properties,
+		nackRedeliveryDelay:         defaultNackRedeliveryDelay,
+		replicateSubscriptionState:  false,
+		decryption:                  options.Decryption,
+		schema:                      options.Schema,
+		backoffPolicy:               options.BackoffPolicy,
+		maxPendingChunkedMessage:    options.MaxPendingChunkedMessage,
+		expireTimeOfIncompleteChunk: options.ExpireTimeOfIncompleteChunk,
+		autoAckIncompleteChunk:      options.AutoAckIncompleteChunk,
 	}
 
 	reader := &reader{


Test patch:
diff --git a/pulsar/ack_grouping_tracker_test.go b/pulsar/ack_grouping_tracker_test.go
index 41d24d4..0a794f6 100644
--- a/pulsar/ack_grouping_tracker_test.go
+++ b/pulsar/ack_grouping_tracker_test.go
@@ -19,11 +19,13 @@ package pulsar
 
 import (
 	"fmt"
+	"sort"
 	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
 
+	pb "github.com/apache/pulsar-client-go/pulsar/internal/pulsar_proto"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -45,7 +47,8 @@ func TestNoCacheTracker(t *testing.T) {
 				ledgerID1 := int64(-1)
 				tracker := newAckGroupingTracker(&option,
 					func(id MessageID) { ledgerID0 = id.LedgerID() },
-					func(id MessageID) { ledgerID1 = id.LedgerID() })
+					func(id MessageID) { ledgerID1 = id.LedgerID() },
+					nil)
 
 				tracker.add(&messageID{ledgerID: 1})
 				assert.Equal(t, atomic.LoadInt64(&ledgerID0), int64(1))
@@ -61,10 +64,12 @@ type mockAcker struct {
 	cumulativeLedgerID int64
 }
 
-func (a *mockAcker) ack(id MessageID) {
+func (a *mockAcker) ack(ids []*pb.MessageIdData) {
 	defer a.Unlock()
 	a.Lock()
-	a.ledgerIDs = append(a.ledgerIDs, id.LedgerID())
+	for _, id := range ids {
+		a.ledgerIDs = append(a.ledgerIDs, int64(*id.LedgerId))
+	}
 }
 
 func (a *mockAcker) ackCumulative(id MessageID) {
@@ -74,6 +79,8 @@ func (a *mockAcker) ackCumulative(id MessageID) {
 func (a *mockAcker) getLedgerIDs() []int64 {
 	defer a.Unlock()
 	a.Lock()
+
+	sort.Slice(a.ledgerIDs, func(i, j int) bool { return a.ledgerIDs[i] < a.ledgerIDs[j] })
 	return a.ledgerIDs
 }
 
@@ -88,8 +95,8 @@ func (a *mockAcker) reset() {
 
 func TestCachedTracker(t *testing.T) {
 	var acker mockAcker
-	tracker := newAckGroupingTracker(&AckGroupingOptions{MaxSize: 3, MaxTime: 0},
-		func(id MessageID) { acker.ack(id) }, func(id MessageID) { acker.ackCumulative(id) })
+	tracker := newAckGroupingTracker(&AckGroupingOptions{MaxSize: 3, MaxTime: 0}, nil,
+		func(id MessageID) { acker.ackCumulative(id) }, func(ids []*pb.MessageIdData) { acker.ack(ids) })
 
 	tracker.add(&messageID{ledgerID: 1})
 	tracker.add(&messageID{ledgerID: 2})
@@ -126,7 +133,8 @@ func TestCachedTracker(t *testing.T) {
 func TestTimedTrackerIndividualAck(t *testing.T) {
 	var acker mockAcker
 	// MaxSize: 1000, MaxTime: 100ms
-	tracker := newAckGroupingTracker(nil, func(id MessageID) { acker.ack(id) }, nil)
+	tracker := newAckGroupingTracker(nil, nil,
+		func(id MessageID) { acker.ackCumulative(id) }, func(ids []*pb.MessageIdData) { acker.ack(ids) })
 
 	expected := make([]int64, 0)
 	for i := 0; i < 999; i++ {
@@ -161,7 +169,7 @@ func TestTimedTrackerIndividualAck(t *testing.T) {
 func TestTimedTrackerCumulativeAck(t *testing.T) {
 	var acker mockAcker
 	// MaxTime is 100ms
-	tracker := newAckGroupingTracker(nil, nil, func(id MessageID) { acker.ackCumulative(id) })
+	tracker := newAckGroupingTracker(nil, nil, func(id MessageID) { acker.ackCumulative(id) }, nil)
 
 	// case 1: flush because of the timeout
 	tracker.addCumulative(&messageID{ledgerID: 1})
@@ -182,7 +190,8 @@ func TestTimedTrackerCumulativeAck(t *testing.T) {
 }
 
 func TestTimedTrackerIsDuplicate(t *testing.T) {
-	tracker := newAckGroupingTracker(nil, func(id MessageID) {}, func(id MessageID) {})
+	tracker := newAckGroupingTracker(nil, func(id MessageID) {}, func(id MessageID) {},
+		func(id []*pb.MessageIdData) {})
 
 	tracker.add(&messageID{batchIdx: 0, batchSize: 3})
 	tracker.add(&messageID{batchIdx: 2, batchSize: 3})
@@ -198,8 +207,8 @@ func TestTimedTrackerIsDuplicate(t *testing.T) {
 
 func TestDuplicateAfterClose(t *testing.T) {
 	var acker mockAcker
-	tracker := newAckGroupingTracker(&AckGroupingOptions{MaxSize: 3, MaxTime: 0},
-		func(id MessageID) { acker.ack(id) }, func(id MessageID) { acker.ackCumulative(id) })
+	tracker := newAckGroupingTracker(&AckGroupingOptions{MaxSize: 3, MaxTime: 0}, nil,
+		func(id MessageID) { acker.ackCumulative(id) }, func(ids []*pb.MessageIdData) { acker.ack(ids) })
 
 	tracker.add(&messageID{ledgerID: 1})
 	assert.True(t, tracker.isDuplicate(&messageID{ledgerID: 1}))
@@ -207,3 +216,28 @@ func TestDuplicateAfterClose(t *testing.T) {
 	tracker.close()
 	assert.False(t, tracker.isDuplicate(&messageID{ledgerID: 1}))
 }
+
+func TestTrackerPendingAcks(t *testing.T) {
+	m := make(map[uint64][]int64)
+	tracker := newAckGroupingTracker(&AckGroupingOptions{MaxSize: 3, MaxTime: 0}, nil, nil,
+		func(ids []*pb.MessageIdData) {
+			for _, id := range ids {
+				m[*id.LedgerId] = id.AckSet
+			}
+		})
+	tracker.add(&messageID{ledgerID: 0, batchIdx: 0, batchSize: 30})
+	for i := 0; i < 10; i++ {
+		tracker.add(&messageID{ledgerID: 1, batchIdx: int32(i), batchSize: 10})
+	}
+	assert.Equal(t, 0, len(m)) // the number of entries is 2, so it's not flushed
+	tracker.flush()
+	assert.Equal(t, 2, len(m))
+
+	ackSet, found := m[0]
+	assert.True(t, found)
+	assert.Greater(t, len(ackSet), 0)
+
+	ackSet, found = m[1]
+	assert.True(t, found)
+	assert.Equal(t, 0, len(ackSet)) // all messages in the batch are acknowledged
+}
diff --git a/pulsar/consumer_partition_test.go b/pulsar/consumer_partition_test.go
index 16c4399..21bc0e3 100644
--- a/pulsar/consumer_partition_test.go
+++ b/pulsar/consumer_partition_test.go
@@ -38,7 +38,7 @@ func TestSingleMessageIDNoAckTracker(t *testing.T) {
 		decryptor:            crypto.NewNoopDecryptor(),
 	}
 	pc.ackGroupingTracker = newAckGroupingTracker(&AckGroupingOptions{MaxSize: 0},
-		func(id MessageID) { pc.sendIndividualAck(id) }, nil)
+		func(id MessageID) { pc.sendIndividualAck(id) }, nil, nil)
 
 	headersAndPayload := internal.NewBufferWrapper(rawCompatSingleMessage)
 	if err := pc.MessageReceived(nil, headersAndPayload); err != nil {
@@ -76,7 +76,7 @@ func TestBatchMessageIDNoAckTracker(t *testing.T) {
 		decryptor:            crypto.NewNoopDecryptor(),
 	}
 	pc.ackGroupingTracker = newAckGroupingTracker(&AckGroupingOptions{MaxSize: 0},
-		func(id MessageID) { pc.sendIndividualAck(id) }, nil)
+		func(id MessageID) { pc.sendIndividualAck(id) }, nil, nil)
 
 	headersAndPayload := internal.NewBufferWrapper(rawBatchMessage1)
 	if err := pc.MessageReceived(nil, headersAndPayload); err != nil {
@@ -111,7 +111,7 @@ func TestBatchMessageIDWithAckTracker(t *testing.T) {
 		decryptor:            crypto.NewNoopDecryptor(),
 	}
 	pc.ackGroupingTracker = newAckGroupingTracker(&AckGroupingOptions{MaxSize: 0},
-		func(id MessageID) { pc.sendIndividualAck(id) }, nil)
+		func(id MessageID) { pc.sendIndividualAck(id) }, nil, nil)
 
 	headersAndPayload := internal.NewBufferWrapper(rawBatchMessage10)
 	if err := pc.MessageReceived(nil, headersAndPayload); err != nil {
diff --git a/pulsar/reader_test.go b/pulsar/reader_test.go
index 3543187..0a1b2a1 100644
--- a/pulsar/reader_test.go
+++ b/pulsar/reader_test.go
@@ -70,6 +70,45 @@ func TestReaderConfigSubscribeName(t *testing.T) {
 	assert.NotNil(t, consumer)
 }
 
+func TestReaderConfigChunk(t *testing.T) {
+	client, err := NewClient(ClientOptions{
+		URL: lookupURL,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+
+	r1, err := client.CreateReader(ReaderOptions{
+		Topic:                       "my-topic1",
+		StartMessageID:              EarliestMessageID(),
+		MaxPendingChunkedMessage:    50,
+		ExpireTimeOfIncompleteChunk: 30 * time.Second,
+		AutoAckIncompleteChunk:      true,
+	})
+	assert.Nil(t, err)
+	defer r1.Close()
+
+	// verify specified chunk options
+	pcOpts := r1.(*reader).pc.options
+	assert.Equal(t, 50, pcOpts.maxPendingChunkedMessage)
+	assert.Equal(t, 30*time.Second, pcOpts.expireTimeOfIncompleteChunk)
+	assert.True(t, pcOpts.autoAckIncompleteChunk)
+
+	r2, err := client.CreateReader(ReaderOptions{
+		Topic:          "my-topic2",
+		StartMessageID: EarliestMessageID(),
+	})
+	assert.Nil(t, err)
+	defer r2.Close()
+
+	// verify default chunk options
+	pcOpts = r2.(*reader).pc.options
+	assert.Equal(t, 100, pcOpts.maxPendingChunkedMessage)
+	assert.Equal(t, time.Minute, pcOpts.expireTimeOfIncompleteChunk)
+	assert.False(t, pcOpts.autoAckIncompleteChunk)
+}
+
 func TestReader(t *testing.T) {
 	client, err := NewClient(ClientOptions{
 		URL: lookupURL,