[
https://issues.apache.org/jira/browse/BEAM-3218?focusedWorklogId=133714&page=com.atlassian.jira.plugin.system.issuetabpanels:worklog-tabpanel#worklog-133714
]
ASF GitHub Bot logged work on BEAM-3218:
----------------------------------------
Author: ASF GitHub Bot
Created on: 10/Aug/18 17:41
Start Date: 10/Aug/18 17:41
Worklog Time Spent: 10m
Work Description: amotamed closed pull request #6183: [BEAM-3218] Added
Quota checks for PubsubMessage in PubsubBoundedWriter
URL: https://github.com/apache/beam/pull/6183
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git
a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/StringUtils.java
b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/StringUtils.java
index 3969be64b03..862082da9d0 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/StringUtils.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/StringUtils.java
@@ -19,7 +19,9 @@
import static com.google.common.base.Preconditions.checkNotNull;
+import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
+import java.util.Collection;
import java.util.List;
/** Utilities for working with JSON and other human-readable string formats. */
@@ -141,4 +143,16 @@ public static int getLevenshteinDistance(final String s,
final String t) {
return v1[t.length()];
}
+
+ public static int getTotalSizeInBytes(Collection<String> list) {
+ int out = 0;
+ for (String s : list) {
+ out += getStringByteSize(s);
+ }
+ return out;
+ }
+
+ public static int getStringByteSize(String s) {
+ return s.getBytes(StandardCharsets.UTF_8).length;
+ }
}
diff --git
a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/StringUtilsTest.java
b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/StringUtilsTest.java
index 9e9686ca201..c675314d67b 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/StringUtilsTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/StringUtilsTest.java
@@ -20,6 +20,7 @@
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
+import java.util.Arrays;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -54,4 +55,14 @@ public void testLevenshteinDistance() {
assertEquals(1, StringUtils.getLevenshteinDistance("abc", "ab1c")); //
insertion
assertEquals(1, StringUtils.getLevenshteinDistance("abc", "a1c")); //
modification
}
+
+ @Test
+ public void testTotalSizeInBytes() {
+ assertEquals(0, StringUtils.getTotalSizeInBytes(Arrays.asList("")));
+ assertEquals(1, StringUtils.getTotalSizeInBytes(Arrays.asList("a")));
+ assertEquals(2, StringUtils.getTotalSizeInBytes(Arrays.asList("ab")));
+ assertEquals(2, StringUtils.getTotalSizeInBytes(Arrays.asList("a", "b")));
+ assertEquals(3, StringUtils.getTotalSizeInBytes(Arrays.asList("abc")));
+ assertEquals(3, StringUtils.getTotalSizeInBytes(Arrays.asList("a", "b",
"c")));
+ }
}
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 5f4027adce5..e860b208edf 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
@@ -851,9 +851,11 @@ public void populateDisplayData(DisplayData.Builder
builder) {
*/
public class PubsubBoundedWriter extends DoFn<T, Void> {
- private static final int MAX_PUBLISH_BATCH_SIZE = 100;
+ private static final int MAX_PUBLISH_BATCH_SIZE = 1000;
+ private static final int MAX_PUBLISH_BYTE_SIZE = 1024 * 1024 * 10;
private transient List<OutgoingMessage> output;
private transient PubsubClient pubsubClient;
+ private int currentByteSize;
@StartBundle
public void startBundle(StartBundleContext c) throws IOException {
@@ -862,6 +864,7 @@ public void startBundle(StartBundleContext c) throws
IOException {
this.pubsubClient =
FACTORY.newClient(
getTimestampAttribute(), null,
c.getPipelineOptions().as(PubsubOptions.class));
+ this.currentByteSize = 0;
}
@ProcessElement
@@ -870,12 +873,17 @@ public void processElement(ProcessContext c) throws
IOException {
PubsubMessage message = getFormatFn().apply(c.element());
payload = message.getPayload();
Map<String, String> attributes = message.getAttributeMap();
- // NOTE: The record id is always null.
- output.add(new OutgoingMessage(payload, attributes,
c.timestamp().getMillis(), null));
- if (output.size() >= MAX_PUBLISH_BATCH_SIZE) {
+ boolean shouldPublish =
+ output.size() == MAX_PUBLISH_BATCH_SIZE
+ || currentByteSize + message.getMessageByteSize() >
MAX_PUBLISH_BYTE_SIZE;
+ if (shouldPublish) {
publish();
}
+
+ // NOTE: The record id is always null.
+ output.add(new OutgoingMessage(payload, attributes,
c.timestamp().getMillis(), null));
+ currentByteSize += message.getMessageByteSize();
}
@FinishBundle
@@ -886,6 +894,7 @@ public void finishBundle() throws IOException {
output = null;
pubsubClient.close();
pubsubClient = null;
+ currentByteSize = 0;
}
private void publish() throws IOException {
@@ -895,6 +904,7 @@ private void publish() throws IOException {
PubsubClient.topicPathFromName(topic.project, topic.topic),
output);
checkState(n == output.size());
output.clear();
+ currentByteSize = 0;
}
@Override
diff --git
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubMessage.java
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubMessage.java
index 2b83dc9a15a..84b5a19b874 100644
---
a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubMessage.java
+++
b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubMessage.java
@@ -18,6 +18,7 @@
package org.apache.beam.sdk.io.gcp.pubsub;
import static com.google.common.base.Preconditions.checkNotNull;
+import static org.apache.beam.sdk.util.StringUtils.getTotalSizeInBytes;
import java.util.Map;
import javax.annotation.Nullable;
@@ -52,4 +53,12 @@ public String getAttribute(String attribute) {
public Map<String, String> getAttributeMap() {
return attributes;
}
+
+ public int getMessageByteSize() {
+ return message.length + getAttributeMapByteSize();
+ }
+
+ int getAttributeMapByteSize() {
+ return getTotalSizeInBytes(attributes.values()) +
getTotalSizeInBytes(attributes.keySet());
+ }
}
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
Issue Time Tracking
-------------------
Worklog Id: (was: 133714)
Time Spent: 1h 10m (was: 1h)
> Change PubsubBoundedWriter to check total byte size
> ---------------------------------------------------
>
> Key: BEAM-3218
> URL: https://issues.apache.org/jira/browse/BEAM-3218
> Project: Beam
> Issue Type: Bug
> Components: io-java-gcp
> Reporter: Thang Nguyen
> Assignee: Reuven Lax
> Priority: Minor
> Time Spent: 1h 10m
> Remaining Estimate: 0h
>
> The PubsubBoundedWriter [does not
> check|https://github.com/apache/beam/blob/master/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubIO.java#L895-L897]
> the total byte size of outgoing messages when publishing.
> AC:
> * Add a check to ensure the total bytes of the outgoing messages is less than
> pubsub's allowed max
--
This message was sent by Atlassian JIRA
(v7.6.3#76005)
