Refactor commit logic out of StreamingDataflowWorker #30312

m-trieu · 2024-02-14T08:30:00Z

StreamingDataflowWorker is a rather large class

For maintainability it would be helpful in the long run to encapsulate different pieces of the worker harness logic into their own class.

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

github-actions · 2024-02-14T09:05:55Z

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

m-trieu · 2024-02-15T01:11:15Z

R: @scwhittle still working on the tests but curious about initial thoughts.
also opted to offload the thread management to executors instead of using raw threads. This may allow us to play around with the thread pool implementation to possibly increase performance (i.e work stealing thread pool).

Thanks!

github-actions · 2024-02-15T01:12:22Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control

m-trieu · 2024-02-21T22:13:39Z

@scwhittle ready for a look! thanks

scwhittle

Didn't get very far through but given round-trip time thought I'd send the comment I did have.

scwhittle · 2024-02-23T17:19:49Z

...e/beam/runners/dataflow/worker/windmill/client/WindmillStreamPoolCloseableStreamFactory.java

+
+  public WindmillStreamPoolCloseableStreamFactory(Supplier<StreamT> streamFactory) {
+    this.streamPool =
+        WindmillStreamPool.create(NUM_COMMIT_STREAMS, COMMIT_STREAM_TIMEOUT, streamFactory);


can we just change StreamPool to have a method returning ClosableStreams? Not sure we need this new class. This class also hard-codes parameters that we might want to change for different stream pools.

done! good idea

scwhittle · 2024-02-29T12:30:21Z

...va/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

+                this::onStreamingCommitComplete,
+                started)
+            : WorkCommitters.createApplianceWorkCommitter(
+                windmillServer::commitWork, numCommitThreads, running::get, started);


the commit threads is always 1 for appliance, should we remove the param?

scwhittle · 2024-02-29T12:35:16Z

...he/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java

+          () -> {
+            while (true) {
+              try {
+                ready.await();


can we get rid of this countdownlatch?

It seems that we'd just fall through to blocking on the commitQueue which is contained in this class. Alternatively if that doesn't work for some reason, having some explicit start method on the committer interface seems better than passing in this latch that does so out-of-band.

ditto for the SE impl

scwhittle · 2024-02-29T12:40:29Z

...r/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/CloseableStream.java

+    return create(stream, () -> {});
+  }
+
+  public abstract Supplier<StreamT> stream();


seems like this should be StreamT, otherwise we could bind a supplier that creates distinct streams on each call, and we're only closing it once.

...he/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java

scwhittle · 2024-02-29T12:58:07Z

...pache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java

+  private boolean tryAddToCommitStream(Commit commit, CommitWorkStream commitStream) {
+    Preconditions.checkNotNull(commit);
+    // Drop commits for failed work. Such commits will be dropped by Windmill anyway.
+    if (commit.work().isFailed()) {


seems better to do this when getting the initial commit in the streamingCommitLoop and in batchCommitsToStream

otherwise we go onto batchCommitsToStream in a weird state that there was no commit added to the stream or we may count commits in batchCommitsToStream which weren't actually added.

scwhittle · 2024-02-29T13:05:04Z

...va/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

-      activeCommitBytes.addAndGet(-size);
-      return false;
-    }
+  private void onStreamingCommitFailed(Commit commit) {


can we get rid of the separate failed handler? Can we just create a CompletedCommit with appropriate error status and use single onCompletedCommit handler?

...he/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java

scwhittle · 2024-02-29T13:11:49Z

...he/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java

+  private static final Logger LOG = LoggerFactory.getLogger(StreamingApplianceWorkCommitter.class);
+  private static final long TARGET_COMMIT_BUNDLE_BYTES = 32 << 20;
+
+  private final Consumer<CommitWorkRequest> commitWork;


some suffix to make it clearer this is funciton/consumer? makes reading code a little confusing below, there is a function of the same name.

...ain/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java

...he/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java

...va/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

...r/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/CloseableStream.java

scwhittle · 2024-03-04T10:10:45Z

...pache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java

+@ThreadSafe
+final class StreamingEngineWorkCommitter implements WorkCommitter {
+  private static final Logger LOG = LoggerFactory.getLogger(StreamingEngineWorkCommitter.class);
+  private static final int COMMIT_BATCH_SIZE = 5;


I think you might have changed the wrong one, the bytes is still a max since that is the queue limit. I meant what is now COMMIT_BATCH_SIZE to be TARGET_COMMIT_BATCH_SIZE (could use KEYS nstead of SIZE as well)

...pache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java

scwhittle · 2024-03-04T10:15:23Z

...src/main/java/org/apache/beam/runners/dataflow/worker/windmill/state/WindmillStateCache.java

@@ -310,6 +311,10 @@ public void invalidate(ByteString processingKey, long shardingKey) {
      keyIndex.remove(key);
    }

+    public void invalidate(ShardedKey shardedKey) {


scwhittle · 2024-03-04T10:17:44Z

...e/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java

+        fakeWindmillServer.waitForAndGetCommits(commits.size() / 2);
+
+    for (Commit commit : commits) {
+      if (commit.work().isFailed()) {


verify that non-failed ones are in committed

scwhittle · 2024-03-04T10:18:48Z

...e/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java

+      WorkItemCommitRequest request = committed.get(commit.work().getWorkItem().getWorkToken());
+      assertNotNull(request);
+      assertThat(request).isEqualTo(commit.request());
+      assertThat(completeCommits).contains(asCompleteCommit(commit));


add a test that verifies some particular keys failing in windmill, that the non-ok commit status is plumbed correctly

done had to make some changes to FakeWindmillServer

scwhittle · 2024-03-04T10:19:17Z

...eam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java

+
+  @Test
+  public void testParallelism() {
+    workCommitter = createWorkCommitter(ignored -> {});


doesn't seem worth testing IMO

Abacn · 2024-03-06T17:21:56Z

fyi today is release cut. What is the status of this PR? I see the comments are all replied? CC: @scwhittle @m-trieu

scwhittle · 2024-03-06T17:32:16Z

There are still some open comments. This is part of work to support direct path which won't be ready for the cut. So I don't think this should hold up the cut.

m-trieu · 2024-03-07T18:19:53Z

ready for another look @scwhittle thanks!

m-trieu · 2024-03-08T02:46:49Z

presubmit failure is due to flaky test (working on a PR to fix) and not related.

m-trieu · 2024-03-08T08:34:03Z

test flakiness attempted to be fixed here #30572

scwhittle · 2024-03-08T14:58:23Z

...ow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java

+                    .setRemainingBytesForWorkItem(0)
+                    .build());
+            StreamingCommitResponse response =
+                streamingCommitsToOffer.getOrDefault(streamingRequestBuilder.build());


you are specifying hte default here and when creating the queue, seems like you could remove here.

I think instead of request/response pair injection you could set commitstatus for each worktoken. Then you wouldn't need the useInjectableStreamingCommitResponses, you could always lookup and default to OK.

The full request/response isn't adding any additonal testing since both sides are fake. And we need to call onDone with exactly one commit status for the request.

done using map instead of response queue

scwhittle · 2024-03-08T14:59:16Z

...ow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java

@@ -122,6 +131,11 @@ public void setDropStreamingCommits(boolean dropStreamingCommits) {
    this.dropStreamingCommits = dropStreamingCommits;
  }

+  public void setUseInjectableStreamingCommitResponses(


add a comment on what this means

scwhittle · 2024-03-08T15:10:28Z

...e/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java

+  }
+
+  @Test
+  public void testCommit_handlesCompleteCommits_commitStatusOK() {


merge with testCommit_sendsCommitsToStreamingEngine?

scwhittle · 2024-03-08T15:15:05Z

...e/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java

+      commits.add(Commit.create(commitRequest, createComputationState("computationId-" + i), work));
+      fakeWindmillServer
+          .whenCommitWorkStreamCalled()
+          .thenAnswer(request -> toResponse(request, INVALID_TOKEN));


would be better to have some ok and some different commit status to verify

m-trieu · 2024-03-12T08:00:57Z

@scwhittle any more comments here? Thank you!

scwhittle

mostly done

scwhittle · 2024-03-12T09:48:03Z

...pache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java

+      commit.work().setState(Work.State.COMMIT_QUEUED);
+    }
+
+    activeCommitBytes.addAndGet(-commit.getSize());


activeCommitBytes is no longer tracking the bytes until they complete, just until they are added on the stream

You should:

move the decrement into the commitWorkItem ondonefn above. That handles the case it was added successfully, tracking bytes until the commit response is recieved

decrement here only if !isCommitSuccesful to cover the case we're keeping it until next time.

scwhittle · 2024-03-12T09:52:03Z

...e/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java

+    }
+
+    workCommitter.start();
+    commits.forEach(workCommitter::commit);


use a paralleliszed forEach to cover thread-safety a little?

m-trieu · 2024-03-12T19:13:07Z

added drain (fail work and clear queues) when we close the SEWorkCommitter

will need this for direct path. in non direct path mode we would never call stop (except when we call stop on the entire harness in StreamingDataflowWorker) so it will not affect that code path. @scwhittle

scwhittle · 2024-03-12T20:02:52Z

...er/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WeightedBoundedQueue.java

+  }
+
+  public void clear() {
+    queue.clear();


this is not updating limit. To do so properly it seems easiest to poll() until null, but actually perhaps a better alternative is to remove stream() and clear() and just do that loop on poll() until null in StreamingEngineWorkCommitter

That is less API of this class to test and in general if queue is still in use, stream()+clear() may process something that is poll()/take() in between the calls.

scwhittle · 2024-03-12T20:08:26Z

...pache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java

+  }
+
+  private void drainCommitQueue() {
+    commitQueue.stream().forEach(this::failCommit);


see other comment, how about loop on poll instead

scwhittle · 2024-03-12T20:10:24Z

...pache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java

+            (commitStatus) -> {
+              onCommitComplete.accept(CompleteCommit.create(commit, commitStatus));
+              activeCommitBytes.addAndGet(-commit.getSize());
+            });

    if (!isCommitSuccessful) {


// Since the commit was not accepted, revert the changes made above.

scwhittle · 2024-03-13T14:25:45Z

Run Java PreCommit

m-trieu · 2024-03-14T06:06:32Z

test failures should be fixed in #30572 since it removes arbitrary waiting from the tests (replaced with countdownlatch as event triggers)

scwhittle · 2024-03-14T09:18:15Z

The other PR is merged, woudl like to verify this passes now

scwhittle · 2024-03-14T09:18:30Z

Run Java PreCommit

…he StreamingEngineWorkCommitter; add test to make sure drain is working

m-trieu · 2024-03-14T19:03:21Z

wooh! looks good

github-actions bot added runners dataflow labels Feb 14, 2024

m-trieu force-pushed the mt-refactor-commits branch from 5260d57 to 4a70e02 Compare February 15, 2024 00:38

m-trieu force-pushed the mt-refactor-commits branch from 74f8a1a to c9aab72 Compare February 21, 2024 22:11

m-trieu force-pushed the mt-refactor-commits branch from 76fd384 to edf03c0 Compare February 22, 2024 19:55

scwhittle mentioned this pull request Feb 23, 2024

Add DirectCommitWorkStream for direct path #30255

Closed

3 tasks

scwhittle requested changes Feb 23, 2024

View reviewed changes

scwhittle requested changes Feb 29, 2024

View reviewed changes

m-trieu force-pushed the mt-refactor-commits branch from 210b847 to 540acdb Compare March 1, 2024 02:08

scwhittle requested changes Mar 4, 2024

View reviewed changes

m-trieu force-pushed the mt-refactor-commits branch from 540acdb to 7e28f3c Compare March 7, 2024 02:25

scwhittle requested changes Mar 8, 2024

View reviewed changes

scwhittle requested changes Mar 12, 2024

View reviewed changes

scwhittle approved these changes Mar 14, 2024

View reviewed changes

m-trieu added 22 commits March 14, 2024 09:15

add commitworkstream for direct path

bef85dc

have different consumers for failed and complete work

c8aacdd

fix test

97e8781

address cl comments

a95c5bc

address cl comments

f9944c7

fix formatting

31de497

move committing logic to its own file

1386e57

refactor commits

bb89dcb

add closeability to streams

bbfdf50

fix formatting

bfb6496

remove unused variables

5c28ab1

add tests

95e63de

update tests

6c75b74

have windmillStreamPool be able to return a closeableStream

f6846ee

address cl comments

77dd15c

run formatter

4d202de

address cl comments

fdbed9a

remove unused method

e7d1628

remove ResponseQueue for CommitWorkStream in test, modify tests.

0c0915a

address cl comments

2871d85

drain the queue and initial commit by failing the work when we stop t…

975f731

…he StreamingEngineWorkCommitter; add test to make sure drain is working

address cl comments

cfa5039

m-trieu force-pushed the mt-refactor-commits branch from dfb5341 to cfa5039 Compare March 14, 2024 16:16

scwhittle merged commit c298da5 into apache:master Mar 15, 2024
17 checks passed

hjtran pushed a commit to hjtran/beam that referenced this pull request Apr 4, 2024

Refactor commit logic out of StreamingDataflowWorker (apache#30312)

c178724

Refactor commit logic out of StreamingDataflowWorker #30312

Refactor commit logic out of StreamingDataflowWorker #30312

Conversation

m-trieu commented Feb 14, 2024

GitHub Actions Tests Status (on master branch)

github-actions bot commented Feb 14, 2024

m-trieu commented Feb 15, 2024

github-actions bot commented Feb 15, 2024

m-trieu commented Feb 21, 2024

scwhittle left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Abacn commented Mar 6, 2024

scwhittle commented Mar 6, 2024

m-trieu commented Mar 7, 2024

m-trieu commented Mar 8, 2024

m-trieu commented Mar 8, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

m-trieu commented Mar 12, 2024

scwhittle left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

m-trieu commented Mar 12, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

scwhittle commented Mar 13, 2024

m-trieu commented Mar 14, 2024

scwhittle commented Mar 14, 2024

scwhittle commented Mar 14, 2024

m-trieu commented Mar 14, 2024