scwhittle commented on code in PR #39961:
URL: https://github.com/apache/beam/pull/39961#discussion_r4024304874
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java:
##########
@@ -860,23 +886,45 @@ private void startForNewKey(Work newWork) throws
CoderException {
}
}
- // Returns state bytes read during the bundle execution
- public long getStateBytesRead() {
- return stateBytesRead;
+ public void addBundleOutputMessages(Windmill.OutputMessageBundle
outputBundle) {
+ getOrCreateBundleOutputMessages().add(outputBundle);
+ }
+
+ public void addBundlePubsubMessages(Windmill.PubSubMessageBundle
pubsubBundle) {
+ getOrCreateBundlePubsubMessages().add(pubsubBundle);
+ }
+
+ private List<Windmill.OutputMessageBundle> getOrCreateBundleOutputMessages()
{
+ if (this.bundleOutputMessages == null) {
+ this.bundleOutputMessages = new ArrayList<>();
+ }
+ return this.bundleOutputMessages;
+ }
+
+ private List<Windmill.PubSubMessageBundle> getOrCreateBundlePubsubMessages()
{
+ if (this.bundlePubsubMessages == null) {
+ this.bundlePubsubMessages = new ArrayList<>();
+ }
+ return this.bundlePubsubMessages;
+ }
+
+ private List<Windmill.WorkItemCommitRequest> getOrCreateWorkItemCommits() {
Review Comment:
how about inlining this? it has a single usage and we can create a
single-item list (like Immutable list which has specialized
SingletonImmutableList) if not using multikeybundles and enforce only one key
is present
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WindmillSink.java:
##########
@@ -350,26 +350,43 @@ public long add(WindowedValue<T> data) throws IOException
{
return (long) key.size() + value.size() + metadata.size() + id.size() +
offsetSize;
}
- @Override
- public void close() throws IOException {
+ private void flush(boolean bundleLevel) {
try {
outputBuilder.setDestinationStreamId(destinationName);
for (Windmill.KeyedMessageBundle.Builder keyedOutput :
productionMap.values()) {
outputBuilder.addBundles(keyedOutput.build());
}
if (outputBuilder.getBundlesCount() > 0) {
- context.getOutputBuilder().addOutputMessages(outputBuilder.build());
+ Windmill.OutputMessageBundle bundle = outputBuilder.build();
+ if (bundleLevel) {
+ context.addBundleOutputMessages(bundle);
+ } else {
+ context.getOutputBuilder().addOutputMessages(bundle);
Review Comment:
maybe we should change getOutputBuilder to be getKeyOutputBuilder? could do
separately.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java:
##########
@@ -187,18 +190,31 @@ public long add(WindowedValue<T> data) throws IOException
{
return byteString.size();
}
+ private void flush(boolean bundleLevel) {
+ try {
+ Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build();
+ if (pubsubMessages.getMessagesCount() > 0) {
+ if (bundleLevel) {
+ // If/when we add support for ordering keys, the flush needs to
happen at the key level
+ context.addBundlePubsubMessages(pubsubMessages);
+ } else {
+ context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+ }
+ }
+ } finally {
+ outputBuilder = createOutputBuilder();
Review Comment:
previously we cleared the builder (though see other comment that seems
buggy).
However should we consider just clearing the messages from the bundle so we
don't have to recreate builder, and it can cache message array etc?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PubsubSink.java:
##########
@@ -187,18 +190,31 @@ public long add(WindowedValue<T> data) throws IOException
{
return byteString.size();
}
+ private void flush(boolean bundleLevel) {
+ try {
+ Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build();
+ if (pubsubMessages.getMessagesCount() > 0) {
+ if (bundleLevel) {
+ // If/when we add support for ordering keys, the flush needs to
happen at the key level
+ context.addBundlePubsubMessages(pubsubMessages);
+ } else {
+ context.getOutputBuilder().addPubsubMessages(pubsubMessages);
+ }
+ }
+ } finally {
+ outputBuilder = createOutputBuilder();
+ }
+ }
+
@Override
public void close() throws IOException {
- Windmill.PubSubMessageBundle pubsubMessages = outputBuilder.build();
- if (pubsubMessages.getMessagesCount() > 0) {
- context.getOutputBuilder().addPubsubMessages(pubsubMessages);
- }
- outputBuilder.clear();
Review Comment:
how did this work previously? It seems if we ever closed and then reused
this object, the outputBuilder would not have the topic etc set. Are these
sinks not actually reused in streaming?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java:
##########
@@ -860,23 +886,45 @@ private void startForNewKey(Work newWork) throws
CoderException {
}
}
- // Returns state bytes read during the bundle execution
- public long getStateBytesRead() {
- return stateBytesRead;
+ public void addBundleOutputMessages(Windmill.OutputMessageBundle
outputBundle) {
+ getOrCreateBundleOutputMessages().add(outputBundle);
+ }
+
+ public void addBundlePubsubMessages(Windmill.PubSubMessageBundle
pubsubBundle) {
+ getOrCreateBundlePubsubMessages().add(pubsubBundle);
+ }
+
+ private List<Windmill.OutputMessageBundle> getOrCreateBundleOutputMessages()
{
+ if (this.bundleOutputMessages == null) {
+ this.bundleOutputMessages = new ArrayList<>();
+ }
+ return this.bundleOutputMessages;
+ }
+
+ private List<Windmill.PubSubMessageBundle> getOrCreateBundlePubsubMessages()
{
+ if (this.bundlePubsubMessages == null) {
+ this.bundlePubsubMessages = new ArrayList<>();
+ }
+ return this.bundlePubsubMessages;
+ }
+
+ private List<Windmill.WorkItemCommitRequest> getOrCreateWorkItemCommits() {
Review Comment:
could also inline the pubsub ones since they also just have single usage
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java:
##########
@@ -497,9 +498,9 @@ public void testSetBacklogBytes() {
stepContext.setBacklogBytes(1234.0);
executionContext.finishKey();
- executionContext.flushState();
+ ExecuteWorkResult result = executionContext.flushStateAndReset();
- assertEquals(1234,
executionContext.getOutputBuilder().getSourceBacklogBytes());
+ assertEquals(1234,
result.workItemCommits().get(0).getSourceBacklogBytes());
Review Comment:
assert there is a single commit
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java:
##########
@@ -706,21 +708,20 @@ public void testReadUnboundedReader() throws Exception {
numReadOnThisIteration,
lessThanOrEqualTo(debugOptions.getUnboundedReaderMaxElements()));
// Extract and verify state modifications.
- context.flushState();
- state = context.getOutputBuilder().getSourceStateUpdates().getState();
+ context.finishKey();
+ WindmillComputationKey computationKey = context.getComputationKey();
+ ExecuteWorkResult result = context.flushStateAndReset();
+ Windmill.WorkItemCommitRequest commitRequest =
result.workItemCommits().get(0);
Review Comment:
assert there is a single commit
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
