feat: track batch size using serialized size of PublishRequest (#2113)

sjvanrossum · michaelpri10 · web-flow · commit be78e64f9fdf · 2024-10-11T19:10:18.000-04:00
* feat: track batch size using serialized size of PublishRequest

* fix: compare against batchedBytes instead of messageSize in flush condition

* fix: also count static overhead in flow control

* fix: adjust thresholds in tests

* fix: clean up merge issue

* fix: revert use of topicNameSize in MessageFlowController

* fix: store topicNameSize as initialBatchedBytes in MessagesBatch

---------

Co-authored-by: Mike Prieto &lt;michaelpri10@gmail.com&gt;
diff --git a/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java b/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java
@@ -46,6 +46,7 @@
 import com.google.cloud.pubsub.v1.stub.PublisherStub;
 import com.google.cloud.pubsub.v1.stub.PublisherStubSettings;
 import com.google.common.base.Preconditions;
+import com.google.protobuf.CodedOutputStream;
 import com.google.pubsub.v1.PublishRequest;
 import com.google.pubsub.v1.PublishResponse;
 import com.google.pubsub.v1.PubsubMessage;
@@ -99,6 +100,7 @@ public class Publisher implements PublisherInterface {
   private static final String OPEN_TELEMETRY_TRACER_NAME = "com.google.cloud.pubsub.v1";
 
   private final String topicName;
+  private final int topicNameSize;
 
   private final BatchingSettings batchingSettings;
   private final boolean enableMessageOrdering;
@@ -145,6 +147,8 @@ public static long getApiMaxRequestBytes() {
 
   private Publisher(Builder builder) throws IOException {
     topicName = builder.topicName;
+    topicNameSize =
+        CodedOutputStream.computeStringSize(PublishRequest.TOPIC_FIELD_NUMBER, this.topicName);
 
     this.batchingSettings = builder.batchingSettings;
     FlowControlSettings flowControl = this.batchingSettings.getFlowControlSettings();
@@ -309,7 +313,7 @@ public ApiFuture<String> publish(PubsubMessage message) {
       }
       MessagesBatch messagesBatch = messagesBatches.get(orderingKey);
       if (messagesBatch == null) {
-        messagesBatch = new MessagesBatch(batchingSettings, orderingKey);
+        messagesBatch = new MessagesBatch(batchingSettings, topicNameSize, orderingKey);
         messagesBatches.put(orderingKey, messagesBatch);
       }
 
@@ -636,7 +640,9 @@ private static final class OutstandingPublish {
     OutstandingPublish(PubsubMessageWrapper messageWrapper) {
       this.publishResult = SettableApiFuture.create();
       this.messageWrapper = messageWrapper;
-      this.messageSize = messageWrapper.getPubsubMessage().getSerializedSize();
+      this.messageSize =
+          CodedOutputStream.computeMessageSize(
+              PublishRequest.MESSAGES_FIELD_NUMBER, messageWrapper.getPubsubMessage());
     }
   }
 
@@ -1093,12 +1099,15 @@ void release(long messageSize) {
 
   private class MessagesBatch {
     private List<OutstandingPublish> messages;
+    private int initialBatchedBytes;
     private int batchedBytes;
     private String orderingKey;
     private final BatchingSettings batchingSettings;
 
-    private MessagesBatch(BatchingSettings batchingSettings, String orderingKey) {
+    private MessagesBatch(
+        BatchingSettings batchingSettings, int initialBatchedBytes, String orderingKey) {
       this.batchingSettings = batchingSettings;
+      this.initialBatchedBytes = initialBatchedBytes;
       this.orderingKey = orderingKey;
       reset();
     }
@@ -1111,7 +1120,7 @@ private OutstandingBatch popOutstandingBatch() {
 
     private void reset() {
       messages = new LinkedList<>();
-      batchedBytes = 0;
+      batchedBytes = initialBatchedBytes;
     }
 
     private boolean isEmpty() {
@@ -1150,7 +1159,9 @@ && getBatchedBytes() + outstandingPublish.messageSize >= getMaxBatchBytes()) {
       // immediately.
       // Alternatively if after adding the message we have reached the batch max messages then we
       // have a batch to send.
-      if ((hasBatchingBytes() && outstandingPublish.messageSize >= getMaxBatchBytes())
+      // Note that exceeding {@link Publisher#getApiMaxRequestBytes()} will result in failed
+      // publishes without compression and may yet fail if a request is not sufficiently compressed.
+      if ((hasBatchingBytes() && getBatchedBytes() >= getMaxBatchBytes())
           || getMessagesCount() == batchingSettings.getElementCountThreshold()) {
         batchesToSend.add(popOutstandingBatch());
       }
diff --git a/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java b/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java
@@ -447,7 +447,7 @@ public void testLargeMessagesDoNotReorderBatches() throws Exception {
                 Publisher.Builder.DEFAULT_BATCHING_SETTINGS
                     .toBuilder()
                     .setElementCountThreshold(10L)
-                    .setRequestByteThreshold(20L)
+                    .setRequestByteThreshold(64L)
                     .setDelayThreshold(Duration.ofSeconds(100))
                     .build())
             .setEnableMessageOrdering(true)
@@ -1150,7 +1150,7 @@ public void testPublishFlowControl_throwException() throws Exception {
                             .setLimitExceededBehavior(
                                 FlowController.LimitExceededBehavior.ThrowException)
                             .setMaxOutstandingElementCount(1L)
-                            .setMaxOutstandingRequestBytes(10L)
+                            .setMaxOutstandingRequestBytes(13L)
                             .build())
                     .build())
             .build();
@@ -1192,7 +1192,7 @@ public void testPublishFlowControl_throwExceptionWithOrderingKey() throws Except
                             .setLimitExceededBehavior(
                                 FlowController.LimitExceededBehavior.ThrowException)
                             .setMaxOutstandingElementCount(1L)
-                            .setMaxOutstandingRequestBytes(10L)
+                            .setMaxOutstandingRequestBytes(13L)
                             .build())
                     .build())
             .setEnableMessageOrdering(true)
@@ -1238,7 +1238,7 @@ public void testPublishFlowControl_block() throws Exception {
                         FlowControlSettings.newBuilder()
                             .setLimitExceededBehavior(FlowController.LimitExceededBehavior.Block)
                             .setMaxOutstandingElementCount(2L)
-                            .setMaxOutstandingRequestBytes(10L)
+                            .setMaxOutstandingRequestBytes(13L)
                             .build())
                     .build())
             .build();
