feat: support cumulative ack

pulsar/consumer.go

@@ -231,6 +231,14 @@ type Consumer interface {
 	// AckID the consumption of a single message, identified by its MessageID
 	AckID(MessageID) error
 
+	// CumulativeAck the reception of all the messages in the stream up to (and including)
+	// the provided message.
+	CumulativeAck(msg Message) error
+
+	// CumulativeAckID the reception of all the messages in the stream up to (and including)
+	// the provided message, identified by its MessageID
+	CumulativeAckID(msgID MessageID) error
+
 	// ReconsumeLater mark a message for redelivery after custom delay
 	ReconsumeLater(msg Message, delay time.Duration)
 

pulsar/consumer_impl.go

@@ -38,6 +38,8 @@ type acker interface {
 	// AckID does not handle errors returned by the Broker side, so no need to wait for doneCh to finish.
 	AckID(id MessageID) error
 	AckIDWithResponse(id MessageID) error
+	CumulativeAckID(msgID MessageID) error
+	CumulativeAckIDWithResponse(msgID MessageID) error
 	NackID(id MessageID)
 	NackMsg(msg Message)
 }
@@ -477,6 +479,26 @@ func (c *consumer) AckID(msgID MessageID) error {
 	return c.consumers[msgID.PartitionIdx()].AckID(msgID)
 }
 
+// CumulativeAck the reception of all the messages in the stream up to (and including)
+// the provided message, identified by its MessageID
+func (c *consumer) CumulativeAck(msg Message) error {
+	return c.CumulativeAckID(msg.ID())
+}
+
+// CumulativeAckID the reception of all the messages in the stream up to (and including)
+// the provided message, identified by its MessageID
+func (c *consumer) CumulativeAckID(msgID MessageID) error {
+	if err := c.checkMsgIDPartition(msgID); err != nil {
+		return err
+	}
+
+	if c.options.AckWithResponse {
+		return c.consumers[msgID.PartitionIdx()].CumulativeAckIDWithResponse(msgID)
+	}
+
+	return c.consumers[msgID.PartitionIdx()].CumulativeAckID(msgID)
+}
+
 // ReconsumeLater mark a message for redelivery after custom delay
 func (c *consumer) ReconsumeLater(msg Message, delay time.Duration) {
 	if delay < 0 {

pulsar/consumer_multitopic.go

@@ -143,6 +143,33 @@ func (c *multiTopicConsumer) AckID(msgID MessageID) error {
 	return mid.consumer.AckID(msgID)
 }
 
+// CumulativeAck the reception of all the messages in the stream up to (and including)
+// the provided message
+func (c *multiTopicConsumer) CumulativeAck(msg Message) error {
+	return c.CumulativeAckID(msg.ID())
+}
+
+// CumulativeAckID the reception of all the messages in the stream up to (and including)
+// the provided message, identified by its MessageID
+func (c *multiTopicConsumer) CumulativeAckID(msgID MessageID) error {
+	mid, ok := toTrackingMessageID(msgID)
+	if !ok {
+		c.log.Warnf("invalid message id type %T", msgID)
+		return errors.New("invalid message id type in multi_consumer")
+	}
+
+	if mid.consumer == nil {
+		c.log.Warnf("unable to ack messageID=%+v can not determine topic", msgID)
+		return errors.New("unable to ack message because consumer is nil")
+	}
+
+	if c.options.AckWithResponse {
+		return mid.consumer.CumulativeAckIDWithResponse(msgID)
+	}
+
+	return mid.consumer.CumulativeAckID(msgID)
+}
+
 func (c *multiTopicConsumer) ReconsumeLater(msg Message, delay time.Duration) {
 	names, err := validateTopicNames(msg.Topic())
 	if err != nil {

pulsar/consumer_partition.go

@@ -431,6 +431,7 @@ func (pc *partitionConsumer) AckIDWithResponse(msgID MessageID) error {
 
 	ackReq := new(ackRequest)
 	ackReq.doneCh = make(chan struct{})
+	ackReq.ackType = individualAck
 	if !trackingID.Undefined() && trackingID.ack() {
 		pc.metrics.AcksCounter.Inc()
 		pc.metrics.ProcessingTime.Observe(float64(time.Now().UnixNano()-trackingID.receivedTime.UnixNano()) / 1.0e9)
@@ -463,6 +464,7 @@ func (pc *partitionConsumer) AckID(msgID MessageID) error {
 
 	ackReq := new(ackRequest)
 	ackReq.doneCh = make(chan struct{})
+	ackReq.ackType = individualAck
 	if !trackingID.Undefined() && trackingID.ack() {
 		pc.metrics.AcksCounter.Inc()
 		pc.metrics.ProcessingTime.Observe(float64(time.Now().UnixNano()-trackingID.receivedTime.UnixNano()) / 1.0e9)
@@ -477,6 +479,72 @@ func (pc *partitionConsumer) AckID(msgID MessageID) error {
 	return nil
 }
 
+func (pc *partitionConsumer) CumulativeAckIDWithResponse(msgID MessageID) error {
+	if state := pc.getConsumerState(); state == consumerClosed || state == consumerClosing {
+		pc.log.WithField("state", state).Error("Failed to ack by closing or closed consumer")
+		return errors.New("consumer state is closed")
+	}
+
+	trackingID, ok := toTrackingMessageID(msgID)
+	if !ok {
+		return errors.New("failed to convert trackingMessageID")
+	}
+
+	ackReq := new(ackRequest)
+	ackReq.doneCh = make(chan struct{})
+	ackReq.ackType = cumulativeAck
+	if !trackingID.Undefined() && trackingID.ack() {
+		pc.metrics.AcksCounter.Inc()
+		pc.metrics.ProcessingTime.Observe(float64(time.Now().UnixNano()-trackingID.receivedTime.UnixNano()) / 1.0e9)
+		ackReq.msgID = trackingID
+		// send ack request to eventsCh
+		pc.eventsCh <- ackReq
+		// wait for the request to complete
+		<-ackReq.doneCh
+
+		pc.options.interceptors.OnAcknowledge(pc.parentConsumer, msgID)
+	}
+
+	if cmid, ok := toChunkedMessageID(msgID); ok {
+		pc.unAckChunksTracker.remove(cmid)
+	}
+
+	return nil
+}
+
+func (pc *partitionConsumer) CumulativeAckID(msgID MessageID) error {
+	if state := pc.getConsumerState(); state == consumerClosed || state == consumerClosing {
+		pc.log.WithField("state", state).Error("Failed to ack by closing or closed consumer")
+		return errors.New("consumer state is closed")
+	}
+
+	// chunk message id will be converted to tracking message id
+	trackingID, ok := toTrackingMessageID(msgID)
+	if !ok {
+		return errors.New("failed to convert trackingMessageID")
+	}
+
+	ackReq := new(ackRequest)
+	ackReq.doneCh = make(chan struct{})
+	ackReq.ackType = cumulativeAck
+	if !trackingID.Undefined() && trackingID.ack() {
+		pc.metrics.AcksCounter.Inc()
+		pc.metrics.ProcessingTime.Observe(float64(time.Now().UnixNano()-trackingID.receivedTime.UnixNano()) / 1.0e9)
+		ackReq.msgID = trackingID
+		// send ack request to eventsCh
+		pc.eventsCh <- ackReq
+		// No need to wait for ackReq.doneCh to finish
+
+		pc.options.interceptors.OnAcknowledge(pc.parentConsumer, msgID)
+	}
+
+	if cmid, ok := toChunkedMessageID(msgID); ok {
+		pc.unAckChunksTracker.remove(cmid)
+	}
+
+	return nil
+}
+
 func (pc *partitionConsumer) NackID(msgID MessageID) {
 	if cmid, ok := toChunkedMessageID(msgID); ok {
 		pc.unAckChunksTracker.nack(cmid)
@@ -691,7 +759,13 @@ func (pc *partitionConsumer) internalAck(req *ackRequest) {
 	cmdAck := &pb.CommandAck{
 		ConsumerId: proto.Uint64(pc.consumerID),
 		MessageId:  messageIDs,
-		AckType:    pb.CommandAck_Individual.Enum(),
+	}
+
+	switch req.ackType {
+	case individualAck:
+		cmdAck.AckType = pb.CommandAck_Individual.Enum()
+	case cumulativeAck:
+		cmdAck.AckType = pb.CommandAck_Cumulative.Enum()
 	}
 
 	if pc.options.ackWithResponse {
@@ -1169,10 +1243,16 @@ func (pc *partitionConsumer) dispatcher() {
 	}
 }
 
+const (
+	individualAck = iota
+	cumulativeAck
+)
+
 type ackRequest struct {
-	doneCh chan struct{}
-	msgID  trackingMessageID
-	err    error
+	doneCh  chan struct{}
+	msgID   trackingMessageID
+	ackType int
+	err     error
 }
 
 type unsubscribeRequest struct {

pulsar/consumer_regex.go

@@ -187,6 +187,33 @@ func (c *regexConsumer) AckID(msgID MessageID) error {
 	return mid.consumer.AckID(msgID)
 }
 
+// CumulativeAck the reception of all the messages in the stream up to (and including)
+// the provided message.
+func (c *regexConsumer) CumulativeAck(msg Message) error {
+	return c.CumulativeAckID(msg.ID())
+}
+
+// CumulativeAckID the reception of all the messages in the stream up to (and including)
+// the provided message, identified by its MessageID
+func (c *regexConsumer) CumulativeAckID(msgID MessageID) error {
+	mid, ok := toTrackingMessageID(msgID)
+	if !ok {
+		c.log.Warnf("invalid message id type %T", msgID)
+		return errors.New("invalid message id type")
+	}
+
+	if mid.consumer == nil {
+		c.log.Warnf("unable to ack messageID=%+v can not determine topic", msgID)
+		return errors.New("unable to ack message because consumer is nil")
+	}
+
+	if c.options.AckWithResponse {
+		return mid.consumer.CumulativeAckIDWithResponse(msgID)
+	}
+
+	return mid.consumer.CumulativeAckID(msgID)
+}
+
 func (c *regexConsumer) Nack(msg Message) {
 	if c.options.EnableDefaultNackBackoffPolicy || c.options.NackBackoffPolicy != nil {
 		msgID := msg.ID()

pulsar/consumer_test.go

@@ -731,6 +731,72 @@ func TestConsumerAck(t *testing.T) {
 	}
 }
 
+func TestConsumerCumulativeAck(t *testing.T) {
+	client, err := NewClient(ClientOptions{
+		URL: lookupURL,
+	})
+
+	assert.Nil(t, err)
+	defer client.Close()
+
+	topicName := newTopicName()
+	ctx := context.Background()
+
+	producer, err := client.CreateProducer(ProducerOptions{
+		Topic: topicName,
+	})
+	assert.Nil(t, err)
+	defer producer.Close()
+
+	consumer, err := client.Subscribe(ConsumerOptions{
+		Topic:            topicName,
+		SubscriptionName: "sub-1",
+		Type:             Exclusive,
+	})
+	assert.Nil(t, err)
+
+	const N = 100
+
+	for i := 0; i < N; i++ {
+		if _, err := producer.Send(ctx, &ProducerMessage{
+			Payload: []byte(fmt.Sprintf("msg-content-%d", i)),
+		}); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	for i := 0; i < N; i++ {
+		msg, err := consumer.Receive(ctx)
+		assert.Nil(t, err)
+		assert.Equal(t, fmt.Sprintf("msg-content-%d", i), string(msg.Payload()))
+
+		if i == N/2-1 {
+			// cumulative acks the first half of messages
+			consumer.CumulativeAck(msg)
+		}
+	}
+
+	consumer.Close()
+
+	// Subscribe again
+	consumer, err = client.Subscribe(ConsumerOptions{
+		Topic:            topicName,
+		SubscriptionName: "sub-1",
+		Type:             Exclusive,
+	})
+	assert.Nil(t, err)
+	defer consumer.Close()
+
+	// We should only receive the 2nd half of messages
+	for i := N / 2; i < N; i++ {
+		msg, err := consumer.Receive(ctx)
+		assert.Nil(t, err)
+		assert.Equal(t, fmt.Sprintf("msg-content-%d", i), string(msg.Payload()))
+
+		consumer.Ack(msg)
+	}
+}
+
 func TestConsumerNack(t *testing.T) {
 	client, err := NewClient(ClientOptions{
 		URL: lookupURL,

pulsar/internal/pulsartracing/consumer_interceptor_test.go

@@ -72,6 +72,14 @@ func (c *mockConsumer) AckID(msgID pulsar.MessageID) error {
 	return nil
 }
 
+func (c *mockConsumer) CumulativeAck(msg pulsar.Message) error {
+	return nil
+}
+
+func (c *mockConsumer) CumulativeAckID(msgID pulsar.MessageID) error {
+	return nil
+}
+
 func (c *mockConsumer) ReconsumeLater(msg pulsar.Message, delay time.Duration) {}
 
 func (c *mockConsumer) Nack(msg pulsar.Message) {}