feat(pubsub): batch receipt modacks (#10234)

hongalex · web-flow · commit 4c2cd10fb73b · 2024-05-20T14:10:46.000-07:00
* feat(pubsub): batch receipt modacks

* add comment on why we use success ack result

* update streaming pull retry test nack logic

* remove extra gofunc wrapper
diff --git a/pubsub/iterator.go b/pubsub/iterator.go
@@ -62,20 +62,22 @@ var (
 )
 
 type messageIterator struct {
-	ctx        context.Context
-	cancel     func() // the function that will cancel ctx; called in stop
-	po         *pullOptions
-	ps         *pullStream
-	subc       *vkit.SubscriberClient
-	subName    string
-	kaTick     <-chan time.Time // keep-alive (deadline extensions)
-	ackTicker  *time.Ticker     // message acks
-	nackTicker *time.Ticker     // message nacks
-	pingTicker *time.Ticker     //  sends to the stream to keep it open
-	failed     chan struct{}    // closed on stream error
-	drained    chan struct{}    // closed when stopped && no more pending messages
-	wg         sync.WaitGroup
-
+	ctx           context.Context
+	cancel        func() // the function that will cancel ctx; called in stop
+	po            *pullOptions
+	ps            *pullStream
+	subc          *vkit.SubscriberClient
+	subName       string
+	kaTick        <-chan time.Time // keep-alive (deadline extensions)
+	ackTicker     *time.Ticker     // message acks
+	nackTicker    *time.Ticker     // message nacks
+	pingTicker    *time.Ticker     //  sends to the stream to keep it open
+	receiptTicker *time.Ticker     // sends receipt modacks
+	failed        chan struct{}    // closed on stream error
+	drained       chan struct{}    // closed when stopped && no more pending messages
+	wg            sync.WaitGroup
+
+	// This mutex guards the structs related to lease extension.
 	mu          sync.Mutex
 	ackTimeDist *distribution.D // dist uses seconds
 
@@ -91,7 +93,9 @@ type messageIterator struct {
 	// ack IDs whose ack deadline is to be modified
 	// ModAcks don't have AckResults but allows reuse of the SendModAck function.
 	pendingModAcks map[string]*AckResult
-	err            error // error from stream failure
+	// ack IDs whose receipt need to be acknowledged with a modack.
+	pendingReceipts map[string]*AckResult
+	err             error // error from stream failure
 
 	eoMu                      sync.RWMutex
 	enableExactlyOnceDelivery bool
@@ -127,6 +131,7 @@ func newMessageIterator(subc *vkit.SubscriberClient, subName string, po *pullOpt
 	ackTicker := time.NewTicker(100 * time.Millisecond)
 	nackTicker := time.NewTicker(100 * time.Millisecond)
 	pingTicker := time.NewTicker(30 * time.Second)
+	receiptTicker := time.NewTicker(100 * time.Millisecond)
 	cctx, cancel := context.WithCancel(context.Background())
 	cctx = withSubscriptionKey(cctx, subName)
 	it := &messageIterator{
@@ -140,13 +145,15 @@ func newMessageIterator(subc *vkit.SubscriberClient, subName string, po *pullOpt
 		ackTicker:          ackTicker,
 		nackTicker:         nackTicker,
 		pingTicker:         pingTicker,
+		receiptTicker:      receiptTicker,
 		failed:             make(chan struct{}),
 		drained:            make(chan struct{}),
 		ackTimeDist:        distribution.New(int(maxDurationPerLeaseExtension/time.Second) + 1),
 		keepAliveDeadlines: map[string]time.Time{},
 		pendingAcks:        map[string]*AckResult{},
 		pendingNacks:       map[string]*AckResult{},
 		pendingModAcks:     map[string]*AckResult{},
+		pendingReceipts:    map[string]*AckResult{},
 	}
 	it.wg.Add(1)
 	go it.sender()
@@ -307,11 +314,15 @@ func (it *messageIterator) receive(maxToPull int32) ([]*Message, error) {
 	it.mu.Unlock()
 
 	if len(ackIDs) > 0 {
-		// When exactly once delivery is not enabled, modacks are fire and forget.
 		if !exactlyOnceDelivery {
-			go func() {
-				it.sendModAck(ackIDs, deadline, false)
-			}()
+			// When exactly once delivery is not enabled, modacks are fire and forget.
+			// Add pending receipt modacks to queue to batch with other modacks.
+			it.mu.Lock()
+			for id := range ackIDs {
+				// Use a SuccessAckResult (dummy) since we don't propagate modacks back to the user.
+				it.pendingReceipts[id] = newSuccessAckResult()
+			}
+			it.mu.Unlock()
 			return msgs, nil
 		}
 
@@ -402,6 +413,7 @@ func (it *messageIterator) sender() {
 	defer it.ackTicker.Stop()
 	defer it.nackTicker.Stop()
 	defer it.pingTicker.Stop()
+	defer it.receiptTicker.Stop()
 	defer func() {
 		if it.ps != nil {
 			it.ps.CloseSend()
@@ -414,6 +426,7 @@ func (it *messageIterator) sender() {
 		sendNacks := false
 		sendModAcks := false
 		sendPing := false
+		sendReceipt := false
 
 		dl := it.ackDeadline()
 
@@ -456,9 +469,12 @@ func (it *messageIterator) sender() {
 			it.mu.Lock()
 			// Ping only if we are processing messages via streaming.
 			sendPing = !it.po.synchronous
+		case <-it.receiptTicker.C:
+			it.mu.Lock()
+			sendReceipt = (len(it.pendingReceipts) > 0)
 		}
 		// Lock is held here.
-		var acks, nacks, modAcks map[string]*AckResult
+		var acks, nacks, modAcks, receipts map[string]*AckResult
 		if sendAcks {
 			acks = it.pendingAcks
 			it.pendingAcks = map[string]*AckResult{}
@@ -471,6 +487,10 @@ func (it *messageIterator) sender() {
 			modAcks = it.pendingModAcks
 			it.pendingModAcks = map[string]*AckResult{}
 		}
+		if sendReceipt {
+			receipts = it.pendingReceipts
+			it.pendingReceipts = map[string]*AckResult{}
+		}
 		it.mu.Unlock()
 		// Make Ack and ModAck RPCs.
 		if sendAcks {
@@ -486,6 +506,9 @@ func (it *messageIterator) sender() {
 		if sendPing {
 			it.pingStream()
 		}
+		if sendReceipt {
+			it.sendModAck(receipts, dl, true)
+		}
 	}
 }
 
diff --git a/pubsub/streaming_pull_test.go b/pubsub/streaming_pull_test.go
@@ -245,15 +245,17 @@ func TestStreamingPullRetry(t *testing.T) {
 	server.wait()
 	for i := 0; i < len(testMessages); i++ {
 		id := testMessages[i].AckId
+		server.mu.Lock()
 		if i%2 == 0 {
 			if !server.Acked[id] {
 				t.Errorf("msg %q should have been acked but wasn't", id)
 			}
 		} else {
-			if dl, ok := server.Deadlines[id]; !ok || dl != 0 {
-				t.Errorf("msg %q should have been nacked but wasn't", id)
+			if server.Acked[id] {
+				t.Errorf("msg %q should have not been acked", id)
 			}
 		}
+		server.mu.Unlock()
 	}
 }
 

Original file line number	Diff line number	Diff line change
`@@ -245,15 +245,17 @@ func TestStreamingPullRetry(t *testing.T) {`
`245`	`245`	`server.wait()`
`246`	`246`	`for i := 0; i < len(testMessages); i++ {`
`247`	`247`	`id := testMessages[i].AckId`
	`248`	`+ server.mu.Lock()`
`248`	`249`	`if i%2 == 0 {`
`249`	`250`	`if !server.Acked[id] {`
`250`	`251`	`t.Errorf("msg %q should have been acked but wasn't", id)`
`251`	`252`	`}`
`252`	`253`	`} else {`
`253`		`- if dl, ok := server.Deadlines[id]; !ok \|\| dl != 0 {`
`254`		`- t.Errorf("msg %q should have been nacked but wasn't", id)`
	`254`	`+ if server.Acked[id] {`
	`255`	`+ t.Errorf("msg %q should have not been acked", id)`
`255`	`256`	`}`
`256`	`257`	`}`
	`258`	`+ server.mu.Unlock()`
`257`	`259`	`}`
`258`	`260`	`}`
`259`	`261`