From b2fa5ddb519cac076df61c26c598a44cbc34d9b6 Mon Sep 17 00:00:00 2001
From: Mike Prieto <mikeprieto@google.com>
Date: Fri, 9 Feb 2024 12:46:44 -0500
Subject: [PATCH] fix: Only synchronize on the outstandingReceipts object in
 the MessageDispatcher

---
 .../cloud/pubsub/v1/MessageDispatcher.java    | 53 ++++++++++---------
 1 file changed, 29 insertions(+), 24 deletions(-)

diff --git a/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java b/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java
index 38abb9572..1810badd2 100644
--- a/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java
+++ b/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java
@@ -382,7 +382,7 @@ private void notifyReceiptComplete() {
     }
   }
 
-  synchronized void processReceivedMessages(List<ReceivedMessage> messages) {
+  void processReceivedMessages(List<ReceivedMessage> messages) {
     Instant totalExpiration = now().plus(maxAckExtensionPeriod);
     List<OutstandingMessage> outstandingBatch = new ArrayList<>(messages.size());
     for (ReceivedMessage message : messages) {
@@ -398,7 +398,9 @@ synchronized void processReceivedMessages(List<ReceivedMessage> messages) {
       if (this.exactlyOnceDeliveryEnabled.get()) {
         // For exactly once deliveries we don't add to outstanding batch because we first
         // process the receipt modack. If that is successful then we process the message.
-        outstandingReceipts.put(message.getAckId(), new ReceiptCompleteData(outstandingMessage));
+        synchronized (outstandingReceipts) {
+          outstandingReceipts.put(message.getAckId(), new ReceiptCompleteData(outstandingMessage));
+        }
       } else if (pendingMessages.putIfAbsent(message.getAckId(), ackHandler) != null) {
         // putIfAbsent puts ackHandler if ackID isn't previously mapped, then return the
         // previously-mapped element.
@@ -417,34 +419,37 @@ synchronized void processReceivedMessages(List<ReceivedMessage> messages) {
     processBatch(outstandingBatch);
   }
 
-  synchronized void notifyAckSuccess(AckRequestData ackRequestData) {
-
-    if (outstandingReceipts.containsKey(ackRequestData.getAckId())) {
-      outstandingReceipts.get(ackRequestData.getAckId()).notifyReceiptComplete();
-      List<OutstandingMessage> outstandingBatch = new ArrayList<>();
-
-      for (Iterator<Entry<String, ReceiptCompleteData>> it =
-              outstandingReceipts.entrySet().iterator();
-          it.hasNext(); ) {
-        Map.Entry<String, ReceiptCompleteData> receipt = it.next();
-        // If receipt is complete then add to outstandingBatch to process the batch
-        if (receipt.getValue().isReceiptComplete()) {
-          it.remove();
-          if (pendingMessages.putIfAbsent(
-                  receipt.getKey(), receipt.getValue().getOutstandingMessage().ackHandler)
-              == null) {
-            outstandingBatch.add(receipt.getValue().getOutstandingMessage());
+  void notifyAckSuccess(AckRequestData ackRequestData) {
+    synchronized (outstandingReceipts) {
+      if (outstandingReceipts.containsKey(ackRequestData.getAckId())) {
+        outstandingReceipts.get(ackRequestData.getAckId()).notifyReceiptComplete();
+        List<OutstandingMessage> outstandingBatch = new ArrayList<>();
+
+        for (Iterator<Entry<String, ReceiptCompleteData>> it =
+                outstandingReceipts.entrySet().iterator();
+            it.hasNext(); ) {
+          Map.Entry<String, ReceiptCompleteData> receipt = it.next();
+          // If receipt is complete then add to outstandingBatch to process the batch
+          if (receipt.getValue().isReceiptComplete()) {
+            it.remove();
+            if (pendingMessages.putIfAbsent(
+                    receipt.getKey(), receipt.getValue().getOutstandingMessage().ackHandler)
+                == null) {
+              outstandingBatch.add(receipt.getValue().getOutstandingMessage());
+            }
+          } else {
+            break;
           }
-        } else {
-          break;
         }
+        processBatch(outstandingBatch);
       }
-      processBatch(outstandingBatch);
     }
   }
 
-  synchronized void notifyAckFailed(AckRequestData ackRequestData) {
-    outstandingReceipts.remove(ackRequestData.getAckId());
+  void notifyAckFailed(AckRequestData ackRequestData) {
+    synchronized (outstandingReceipts) {
+      outstandingReceipts.remove(ackRequestData.getAckId());
+    }
   }
 
   private void processBatch(List<OutstandingMessage> batch) {