From ed1b8b27f9e2fac6a7fa950784f6dc2521ce331f Mon Sep 17 00:00:00 2001
From: 1fanwang <1fannnw@gmail.com>
Date: Tue, 25 Aug 2026 16:01:48 -0400
Subject: [PATCH] [PubsubIO] Account for message metadata in bounded write
 batch size

Signed-off-by: 1fanwang <1fannnw@gmail.com>
---
 .../beam/sdk/io/gcp/pubsub/PubsubIO.java      |  7 +---
 .../beam/sdk/io/gcp/pubsub/PubsubIOTest.java  | 41 +++++++++++++++++++
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
index fa2399abd91c..2f522c653963 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java
@@ -1912,14 +1912,11 @@ public void processElement(@Element PubsubMessage message, @Timestamp Instant ti
         // TODO(sjvanrossum): https://github.com/apache/beam/issues/31800
         // - Size validation makes no distinction between JSON and Protobuf encoding
         // - Accounting for HTTP to gRPC transcoding is non-trivial
-        PreparePubsubWriteDoFn.validatePubsubMessage(message, maxPublishBatchByteSize);
+        final int messageSize =
+            PreparePubsubWriteDoFn.validatePubsubMessage(message, maxPublishBatchByteSize);
         // NOTE: The record id is always null since it will be assigned by Pub/Sub.
         final OutgoingMessage msg =
             OutgoingMessage.of(message, timestamp.getMillis(), null, message.getTopic());
-        // TODO(sjvanrossum): https://github.com/apache/beam/issues/31800
-        // - Size validation makes no distinction between JSON and Protobuf encoding
-        // - Accounting for HTTP to gRPC transcoding is non-trivial
-        final int messageSize = msg.getMessage().getData().size();
 
         final PubsubTopic pubsubTopic;
         ValueProvider<PubsubTopic> topicProvider = getTopicProvider();
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
index 3d9c65aa1376..7a88a330d0fe 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIOTest.java
@@ -67,6 +67,7 @@
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.DoFnTester;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.SimpleFunction;
 import org.apache.beam.sdk.transforms.display.DisplayData;
@@ -82,6 +83,7 @@
 import org.apache.beam.sdk.values.ValueInSingleWindow;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -921,6 +923,45 @@ public void testBigMessageBounded() throws IOException {
     }
   }
 
+  @Test
+  public void testBoundedWriteBatchSizeIncludesAttributesAndOrderingKey() throws Exception {
+    PubsubClient mockClient = Mockito.mock(PubsubClient.class);
+    PubsubClient.PubsubClientFactory mockFactory =
+        Mockito.mock(PubsubClient.PubsubClientFactory.class);
+    Mockito.when(
+            mockFactory.newClient(
+                Mockito.isNull(), Mockito.isNull(), Mockito.any(), Mockito.isNull()))
+        .thenReturn(mockClient);
+
+    List<Integer> publishedBatchSizes = Lists.newArrayList();
+    Mockito.when(mockClient.publish(Mockito.any(), Mockito.anyList()))
+        .thenAnswer(
+            invocation -> {
+              List<OutgoingMessage> messages = invocation.getArgument(1);
+              publishedBatchSizes.add(messages.size());
+              return messages.size();
+            });
+
+    PubsubIO.Write<PubsubMessage> write =
+        PubsubIO.writeMessages()
+            .to("projects/project/topics/topic")
+            .withClientFactory(mockFactory)
+            .withOrderingKey()
+            .withMaxBatchBytesSize(40);
+    PubsubMessage message =
+        new PubsubMessage(
+                "0123456789".getBytes(StandardCharsets.UTF_8), ImmutableMap.of("k", "value"))
+            .withOrderingKey("order");
+
+    try (DoFnTester<PubsubMessage, Void> tester =
+        DoFnTester.of(write.new PubsubBoundedWriter(100, 40))) {
+      tester.setCloningBehavior(DoFnTester.CloningBehavior.DO_NOT_CLONE);
+      tester.processBundle(ImmutableList.of(message, message));
+    }
+
+    assertEquals(ImmutableList.of(1, 1), publishedBatchSizes);
+  }
+
   @Test
   public void testReadValidate() throws IOException {
     PubsubOptions options = TestPipeline.testingPipelineOptions().as(PubsubOptions.class);