feat: Receipt modack

google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/MessageDispatcher.java

@@ -28,6 +28,7 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.pubsub.v1.PubsubMessage;
 import com.google.pubsub.v1.ReceivedMessage;
+import io.opencensus.trace.Link;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -89,7 +90,9 @@ class MessageDispatcher {
   private final LinkedBlockingQueue<AckRequestData> pendingAcks = new LinkedBlockingQueue<>();
   private final LinkedBlockingQueue<AckRequestData> pendingNacks = new LinkedBlockingQueue<>();
   private final LinkedBlockingQueue<AckRequestData> pendingReceipts = new LinkedBlockingQueue<>();
-
+  private final ConcurrentMap<String, OutstandingMessage> outstandingReceipts = new ConcurrentHashMap<>();
+  private final List<OutstandingMessage> exactlyOnceOutstandingBatch = new ArrayList<>();
+  private final LinkedBlockingQueue<OutstandingMessage> exactlyOncePendingBatch = new LinkedBlockingQueue<>();
   private final AtomicInteger messageDeadlineSeconds = new AtomicInteger();
   private final AtomicBoolean extendDeadline = new AtomicBoolean(true);
   private final Lock jobLock;
@@ -261,6 +264,16 @@ public void run() {
                               TimeUnit.SECONDS);
                     }
                     processOutstandingOperations();
+                    // List<OutstandingMessage> outstandingBatch = new ArrayList<>();
+                    // for(OutstandingMessage message: exactlyOnceOutstandingBatch){
+                    //   if(!exactlyOncePendingBatch.isEmpty() &&
+                    //       message.receivedMessage.getAckId() == exactlyOncePendingBatch.peek().receivedMessage.getAckId()){
+                    //         outstandingBatch.add(message);
+                    //         exactlyOncePendingBatch.poll();
+                    //         exactlyOnceOutstandingBatch.remove(message);
+                    //   }
+                    // }
+                    // processBatch(outstandingBatch);
                   } catch (Throwable t) {
                     // Catch everything so that one run failing doesn't prevent subsequent runs.
                     logger.log(Level.WARNING, "failed to run periodic job", t);
@@ -348,6 +361,7 @@ private OutstandingMessage(ReceivedMessage receivedMessage, AckHandler ackHandle
       this.receivedMessage = receivedMessage;
       this.ackHandler = ackHandler;
     }
+
   }
 
   void processReceivedMessages(List<ReceivedMessage> messages) {
@@ -361,7 +375,7 @@ void processReceivedMessages(List<ReceivedMessage> messages) {
       AckRequestData ackRequestData = builder.build();
       AckHandler ackHandler =
           new AckHandler(ackRequestData, message.getMessage().getSerializedSize(), totalExpiration);
-      if (pendingMessages.putIfAbsent(message.getAckId(), ackHandler) != null) {
+      if (!this.exactlyOnceDeliveryEnabled.get() && pendingMessages.putIfAbsent(message.getAckId(), ackHandler) != null) {
         // putIfAbsent puts ackHandler if ackID isn't previously mapped, then return the
         // previously-mapped element.
         // If the previous element is not null, we already have the message and the new one is
@@ -372,13 +386,48 @@ void processReceivedMessages(List<ReceivedMessage> messages) {
         // totally expire so that pubsub service sends us the message again.
         continue;
       }
-      outstandingBatch.add(new OutstandingMessage(message, ackHandler));
+      OutstandingMessage outstandingMessage = new OutstandingMessage(message, ackHandler);
+      if (this.exactlyOnceDeliveryEnabled.get()) {
+        exactlyOncePendingBatch.add(outstandingMessage);
+        outstandingReceipts.put(message.getAckId(), outstandingMessage);
+      } else {
+        outstandingBatch.add(outstandingMessage);
+      }
       pendingReceipts.add(ackRequestData);
     }
-
     processBatch(outstandingBatch);
   }
 
+  void notifyAckSuccess(AckRequestData ackRequestData) {
+
+    if(outstandingReceipts.containsKey(ackRequestData.getAckId())) {
+      OutstandingMessage outstandingMessage = outstandingReceipts.get(ackRequestData.getAckId());
+      pendingMessages.putIfAbsent(outstandingMessage.receivedMessage.getAckId(),
+          outstandingMessage.ackHandler);
+
+      exactlyOnceOutstandingBatch.add(outstandingMessage);
+      List<OutstandingMessage> outstandingBatch = new ArrayList<>();
+      for (OutstandingMessage message : exactlyOnceOutstandingBatch) {
+        if (!exactlyOncePendingBatch.isEmpty() &&
+            message.receivedMessage.getAckId()
+                == exactlyOncePendingBatch.peek().receivedMessage.getAckId()) {
+          outstandingBatch.add(message);
+          exactlyOncePendingBatch.poll();
+          exactlyOnceOutstandingBatch.remove(message);
+        }
+      }
+      processBatch(outstandingBatch);
+    }
+  }
+
+  void notifyAckFailed(AckRequestData ackRequestData) {
+    if(outstandingReceipts.containsKey(ackRequestData.getAckId())){
+      outstandingReceipts.remove(ackRequestData);
+      OutstandingMessage outstandingMessage = outstandingReceipts.get(ackRequestData.getAckId());
+      exactlyOncePendingBatch.remove(outstandingMessage);
+    }
+  }
+
   private void processBatch(List<OutstandingMessage> batch) {
     messagesWaiter.incrementPendingCount(batch.size());
     for (OutstandingMessage message : batch) {
@@ -519,6 +568,7 @@ void extendDeadlines() {
 
   @InternalApi
   void processOutstandingOperations() {
+
     List<ModackRequestData> modackRequestData = new ArrayList<ModackRequestData>();
 
     // Nacks are modacks with an expiration of 0

google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/StreamingSubscriberConnection.java

@@ -519,6 +519,7 @@ public void onSuccess(Empty empty) {
         for (AckRequestData ackRequestData : ackRequestDataList) {
           // This will check if a response is needed, and if it has already been set
           ackRequestData.setResponse(AckResponse.SUCCESSFUL, setResponseOnSuccess);
+          messageDispatcher.notifyAckSuccess(ackRequestData);
           // Remove from our pending operations
           pendingRequests.remove(ackRequestData);
         }
@@ -541,6 +542,7 @@ public void onFailure(Throwable t) {
           Map<String, String> metadataMap = getMetadataMapFromThrowable(t);
           ackRequestDataList.forEach(
               ackRequestData -> {
+                messageDispatcher.notifyAckFailed(ackRequestData);
                 String ackId = ackRequestData.getAckId();
                 if (metadataMap.containsKey(ackId)) {
                   // An error occured

google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/MessageDispatcherTest.java

@@ -139,6 +139,16 @@ public void testReceiptMessageReceiver() {
         .receiveMessage(eq(TEST_MESSAGE.getMessage()), any(AckReplyConsumer.class));
   }
 
+  // @Test
+  // public void testBatchSizeofExactlyOnceDelivered() {
+  //   MessageReceiver mockMessageReceiver = mock(MessageReceiver.class);
+  //   MessageDispatcher messageDispatcher = getMessageDispatcher(mockMessageReceiver);
+  //   List<ReceivedMessage> receivedMessageList = new ArrayList<ReceivedMessage>();
+  //   Collections.addAll(receivedMessageList, TEST_MESSAGE, TEST_MESSAGE, TEST_MESSAGE, TEST_MESSAGE, TEST_MESSAGE);
+  //   messageDispatcher.processReceivedMessages(receivedMessageList);
+  //
+  // }
+
   @Test
   public void testReceiptMessageReceiverWithAckResponse() {
     MessageReceiverWithAckResponse mockMessageReceiverWithAckResponse =