Add ordering key size validation to validatePubsubMessageSize

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PreparePubsubWriteDoFn.java

@@ -63,6 +63,20 @@ static int validatePubsubMessageSize(PubsubMessage message, int maxPublishBatchS
     }
     int totalSize = payloadSize;
 
+    @Nullable String orderingKey = message.getOrderingKey();
+    if (orderingKey != null) {
+      int orderingKeySize = orderingKey.getBytes(StandardCharsets.UTF_8).length;
+      if (orderingKeySize > PUBSUB_MESSAGE_ATTRIBUTE_MAX_VALUE_BYTES) {
+        throw new SizeLimitExceededException(
+            "Pubsub message ordering key of length "
+                + orderingKeySize
+                + " exceeds maximum of "
+                + PUBSUB_MESSAGE_ATTRIBUTE_MAX_VALUE_BYTES
+                + " bytes. See https://cloud.google.com/pubsub/quotas#resource_limits");
+      }
+      totalSize += orderingKeySize;
+    }
+
     @Nullable Map<String, String> attributes = message.getAttributeMap();
     if (attributes != null) {
       if (attributes.size() > PUBSUB_MESSAGE_MAX_ATTRIBUTES) {

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsub/PreparePubsubWriteDoFnTest.java

@@ -44,6 +44,19 @@ public void testValidatePubsubMessageSizeOnlyPayload() throws SizeLimitExceededE
     assertEquals(data.length, messageSize);
   }
 
+  @Test
+  public void testValidatePubsubMessageSizePayloadAndOrderingKey()
+      throws SizeLimitExceededException {
+    byte[] data = new byte[1024];
+    String orderingKey = "key";
+    PubsubMessage message = new PubsubMessage(data, null, null, orderingKey);
+
+    int messageSize =
+        PreparePubsubWriteDoFn.validatePubsubMessageSize(message, PUBSUB_MESSAGE_MAX_TOTAL_SIZE);
+
+    assertEquals(data.length + orderingKey.getBytes(StandardCharsets.UTF_8).length, messageSize);
+  }
+
   @Test
   public void testValidatePubsubMessageSizePayloadAndAttributes()
       throws SizeLimitExceededException {
@@ -76,6 +89,19 @@ public void testValidatePubsubMessageSizePayloadTooLarge() {
                 message, PUBSUB_MESSAGE_MAX_TOTAL_SIZE));
   }
 
+  @Test
+  public void testValidatePubsubMessageSizePayloadPlusOrderingKeyTooLarge() {
+    byte[] data = new byte[(10 << 20)];
+    String orderingKey = "key";
+    PubsubMessage message = new PubsubMessage(data, null, null, orderingKey);
+
+    assertThrows(
+        SizeLimitExceededException.class,
+        () ->
+            PreparePubsubWriteDoFn.validatePubsubMessageSize(
+                message, PUBSUB_MESSAGE_MAX_TOTAL_SIZE));
+  }
+
   @Test
   public void testValidatePubsubMessageSizePayloadPlusAttributesTooLarge() {
     byte[] data = new byte[(10 << 20)];
@@ -121,6 +147,19 @@ public void testValidatePubsubMessageSizeAttributeValueTooLarge() {
                 message, PUBSUB_MESSAGE_MAX_TOTAL_SIZE));
   }
 
+  @Test
+  public void testValidatePubsubMessageSizeOrderingKeyTooLarge() {
+    byte[] data = new byte[1024];
+    String orderingKey = RandomStringUtils.randomAscii(1025);
+    PubsubMessage message = new PubsubMessage(data, null, null, orderingKey);
+
+    assertThrows(
+        SizeLimitExceededException.class,
+        () ->
+            PreparePubsubWriteDoFn.validatePubsubMessageSize(
+                message, PUBSUB_MESSAGE_MAX_TOTAL_SIZE));
+  }
+
   @Test
   public void testValidatePubsubMessagePayloadTooLarge() {
     byte[] data = new byte[(10 << 20) + 1];