break nested classes out of StreamingDataflowWorker #28537

m-trieu · 2023-09-19T18:16:59Z

StreamingDataflowWorker file is quite large. Opting to break out the inner static classes to a top level and introduce /streaming directory under the worker/ directory.

There will be future changes to these classes to support direct path, and this change will make the files easier to work with.

includes some formatting changes required by spotless

Slowly switching to java.util.Optional instead of the google Optional.

r: @kennknowles

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.

kennknowles · 2023-09-19T18:19:21Z

Can you separate the independent bits into separate commits? (can still be one PR)

m-trieu · 2023-09-19T19:09:01Z

@kennknowles I split the commits, will do this ahead next time. Thanks!

github-actions · 2023-09-19T19:34:55Z

Assigning reviewers. If you would like to opt out of this review, comment assign to next reviewer:

R: @AnandInguva added as fallback since no labels match configuration

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

m-trieu · 2023-09-19T19:47:11Z

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

+    return Transport.getJsonFactory().fromString(input, MapTask.class);
+  }
+
+  public static void main(String[] args) throws Exception {


@kennknowles out of curiosity, do you happen to know where this is referenced? I was looking for it in a gradle file, but couldn't find it and not sure if its even referenced in those.

Thanks!

It is referenced in dataflow internal launching of the Dataflow v1 Java worker (not in github repo).
With beam portability, the runner specifics like this are behind the FnApi but for the v1 harness the dataflow details leaked some. They were initially separate repositories but we kept having jar compatibility issues so we moved it to the external beam repository and single jar.

m-trieu · 2023-09-20T22:49:54Z

r: @scwhittle

github-actions · 2023-09-20T22:51:37Z

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

scwhittle · 2023-09-21T10:05:04Z

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

+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Bounded set of queues, with a maximum total weight. */
+public class WeightedBoundedQueue<V> {


would be nice to add some tests for this now that it's visible

scwhittle · 2023-09-21T10:07:53Z

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

+  }
+
+  /** Returns the current weight of the queue. */
+  public int weight() {


IMO this method name and size() aren't particularly obvious, what do you think?

maybe
queuedElementsWeight()
queuedElementsCount()

scwhittle · 2023-09-21T10:09:37Z

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

+            ? ImmutableMap.copyOf(transformUserNameToStateFamily)
+            : ImmutableMap.of();
+    this.computationStateCache = computationStateCache;
+    Preconditions.checkNotNull(mapTask.getStageName());


ditto (also if kept put the preconditions at the top?)

moved, mapTask is not nullable, but unfortunately systemName and stageName fields in map task are

scwhittle · 2023-09-21T10:11:01Z

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

+    this.mapTask = mapTask;
+    this.executor = executor;
+    this.transformUserNameToStateFamily =
+        transformUserNameToStateFamily != null


add @ Nullable annotation or remove

scwhittle · 2023-09-21T10:12:23Z

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

+    return executionStateQueue;
+  }
+
+  /** Mark the given shardedKey and work as active. */


and schedules execution of work if there is no active work for the shardedKey already processing.

scwhittle · 2023-09-21T11:23:13Z

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

+  static final int MAX_COMMIT_QUEUE_BYTES = 500 << 20; // 500MB
+  static final int NUM_COMMIT_STREAMS = 1;
+  static final int GET_WORK_STREAM_TIMEOUT_MINUTES = 3;
+  static final Duration COMMIT_STREAM_TIMEOUT = Duration.standardMinutes(1);


should the other constants be moved to the top here?
(if not, what is the motivation for moving some of them?)

i think either the google java linter style or the spotless apply linter organizes constants in public, protected package private, then private order.

this was applied by spotless apply or google java formatter

scwhittle · 2023-09-21T11:23:50Z

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

-  }
-
+  private static final Random clientIdGenerator = new Random();
+  final WindmillStateCache stateCache;


can this be private final? and stay in same location?

this is accessed in the StreamingDataflowWorkerTest.testMergeWindowsCaching, i can mark it as @VisibleForTesting?

I plan to do some breakup/cleanup of StreamingDataflowWorker in a later CL. If we opt to use dependency injection and inject the cache for tests, we can make all of this private.

I was thinking we can eventually have either StreamingApplianceDataflowWorker or StreamingEngineDataflowWorker (depending on the StreamingDataflowWorkerOptions). StreamingDataflowWorker can instantiate either one depending on the passed in options, and we can have different tests for each.

We can inject all the dependencies into the constructors and have some static factory class/methods to create the shapes we want (for test or not etc).

Future code changes won't ever touch the explicit code paths/classes for appliance.

Also for my own knowledge, why don't we use Guice or Dagger dependency injection (if we want to avoid the runtime overhead)? @scwhittle

@kennknowles
I'm not very aware of Java ecosystem and didn't originally write these files. I'm not sure if either of those frameworks are used in Beam but I don't know if that was a conscious choice.

scwhittle · 2023-09-21T11:25:10Z

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

+  // is enabled.
+  private ScheduledExecutorService globalConfigRefreshTimer;
+  // Periodic sender of debug information to the debug capture service.
+  private DebugCapture.Manager debugCaptureManager = null;


can remove = null, mark it nullable

scwhittle · 2023-09-21T11:27:37Z

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

-      workIdBuilder.append('-');
-      workIdBuilder.append(Long.toHexString(workItem.getWorkToken()));
-      DataflowWorkerLoggingMDC.setWorkId(workIdBuilder.toString());
+      String workIdBuilder =


no longer a builder. Does Java optimize repeated string additions better now? IIRC we maybe saw cpu or objects on a profile

I believe StringBuilder is very useful when concatenating in a loop, for simple concatenation the compiler will optimize for us. https://dzone.com/articles/string-concatenation-performacne-improvement-in-ja

scwhittle · 2023-09-21T11:28:47Z

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

@@ -1393,7 +1367,7 @@ private Commit batchCommitsToStream(CommitWorkStream commitStream) {
      Commit commit;
      try {
        if (commits < 5) {
-          commit = commitQueue.poll(10 - 2 * commits, TimeUnit.MILLISECONDS);
+          commit = commitQueue.poll(10 - 2L * commits, TimeUnit.MILLISECONDS);


why is L required?

to avoid integer multiplication implicitly cast to long (complaint from the static checker)

scwhittle · 2023-09-22T09:52:27Z

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

-  }
-
+  private static final Random clientIdGenerator = new Random();
+  final WindmillStateCache stateCache;


@kennknowles
I'm not very aware of Java ecosystem and didn't originally write these files. I'm not sure if either of those frameworks are used in Beam but I don't know if that was a conscious choice.

scwhittle · 2023-09-22T10:02:27Z

.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/WeightBoundedQueueTest.java

+    assertEquals(MAX_WEIGHT, queue.queuedElementsWeight());
+    assertEquals(1, queue.size());
+
+    // Have another thread poll the queue, pulling off the only value inside and freeing up the


can just poll directly from this thread instead of via a new thread, easier to verify result etc then.

scwhittle · 2023-09-22T10:02:59Z

.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/WeightBoundedQueueTest.java

+        WeightedBoundedQueue.create(MAX_WEIGHT, i -> Math.min(MAX_WEIGHT, i));
+
+    // Insert value that takes all the capacity into the queue.
+    Thread thread1 = new Thread(() -> queue.put(MAX_WEIGHT));


can just put directly from the test thread

scwhittle · 2023-09-22T10:22:43Z

.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/WeightBoundedQueueTest.java

+    // Try to insert another value into the queue. This will block since there is no capacity in the
+    // queue.
+    Thread thread2 = new Thread(() -> queue.put(MAX_WEIGHT));
+    thread2.start();


sleep 100ms? just to give some time for thread to get into blocking put (test will pass either way so shouldn't add flakiness)

scwhittle · 2023-09-22T10:23:49Z

.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/WeightBoundedQueueTest.java

+    WeightedBoundedQueue<Integer> queue =
+        WeightedBoundedQueue.create(MAX_WEIGHT, i -> Math.min(MAX_WEIGHT, i));
+
+    assertNull(queue.poll());


add test of poll with timeout method

scwhittle · 2023-09-22T10:24:27Z

.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/WeightBoundedQueueTest.java

+                throw new RuntimeException(e);
+              }
+            });
+    takeThread.start();


add a sleep to give chance for it to start blocking

scwhittle · 2023-09-22T10:24:59Z

.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/WeightBoundedQueueTest.java

+  }
+
+  @Test
+  public void testTake() throws InterruptedException {


add a variant where you poll with very large timeout and ensure it gets the element

done used polling at 1 minute, inserting element at 30 seconds, so poll should return the element.

Thanks, I just meant something like testPoll_withTimeout with a large timeout we never expected to get hit. You could remove testPoll_withLargeTimeout since I don't think it adds much beyond withTimeout and theses tests don't seem to run with much parallelism and we don't want to slow the build down too much.

sg removed!

scwhittle · 2023-09-22T10:26:31Z

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

+   * Marks the work for the given shardedKey as complete. Schedules queued work for the key if any.
+   */
+  public void completeWork(ShardedKey shardedKey, long workToken) {
+    Work nextWork;


ping on this and some other uncompleted comments

scwhittle · 2023-09-25T09:20:18Z

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

@@ -125,23 +126,28 @@ public boolean activateWork(ShardedKey shardedKey, Work work) {
  /**
   * Marks the work for the given shardedKey as complete. Schedules queued work for the key if any.
   */
-  public void completeWork(ShardedKey shardedKey, long workToken) {
-    Work nextWork;
+  public void completeWorkAndScheduleNextWork(ShardedKey shardedKey, long workToken) {


nit: maybe CompleteWorkAndScheduleNextWorkForKey

scwhittle · 2023-09-25T09:20:42Z

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

    synchronized (activeWork) {
-      Queue<Work> queue = activeWork.get(shardedKey);
-      if (queue == null) {
+      Queue<Work> workQueue = activeWork.get(shardedKey);


nullable annotation or optional

made optional and moved into ActiveWorkState + tests

scwhittle · 2023-09-25T09:27:33Z

.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/WeightBoundedQueueTest.java

+  }
+
+  @Test
+  public void testTake() throws InterruptedException {


Thanks, I just meant something like testPoll_withTimeout with a large timeout we never expected to get hit. You could remove testPoll_withLargeTimeout since I don't think it adds much beyond withTimeout and theses tests don't seem to run with much parallelism and we don't want to slow the build down too much.

scwhittle · 2023-09-25T09:28:03Z

.../src/test/java/org/apache/beam/runners/dataflow/worker/streaming/WeightBoundedQueueTest.java

+            });
+    takeThread.start();
+
+    Thread putThread =


remove putThread, can just be inlined to test.

scwhittle · 2023-09-25T09:31:31Z

...-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/counters/NameContext.java

@@ -31,7 +31,10 @@ public abstract class NameContext {
   * systemName} and a {@code userName}.
   */
  public static NameContext create(
-      String stageName, String originalName, String systemName, String userName) {
+      String stageName,


Ping, I realized this is just method and not members generated by autovalue below. But can stageName be non-null below?

scwhittle · 2023-09-25T09:34:29Z

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

+
+  @Override
+  public void close() throws Exception {
+    ExecutionState executionState;


scwhittle · 2023-09-25T09:34:37Z

.../main/java/org/apache/beam/runners/dataflow/worker/streaming/KeyCommitTooLargeException.java

+
+import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
+
+public class KeyCommitTooLargeException extends Exception {


scwhittle · 2023-09-25T09:35:06Z

...-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ShardedKey.java

+    if (keyToDisplay.size() > 100) {
+      keyToDisplay = keyToDisplay.substring(0, 100);
+    }
+    return String.format("%016x-%s", shardingKey(), TextFormat.escapeBytes(keyToDisplay));


scwhittle · 2023-09-25T09:35:19Z

...w-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/StageInfo.java

+    return counterUpdates;
+  }
+
+  public String getStageName() {


scwhittle · 2023-09-26T13:01:40Z

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

+                            "Active key %s without work, expected token %d",
+                            shardedKey, workToken)));
+
+    if (completedWork.getWorkItem().getWorkToken() != workToken) {


we should be comparing the workToken and cacheToken

discussed offline, will follow up in next cl

scwhittle · 2023-09-26T13:03:06Z

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

+
+    // Ensure we don't already have this work token queued.
+    for (Work queuedWork : workQueue) {
+      if (queuedWork.getWorkItem().getWorkToken() == work.getWorkItem().getWorkToken()) {


if the cache token is different then it isn't a pure duplicate it could be a retry, in that case it might be better to take the more recent observed item (guessing it is more likely the newer one) or keep both.

sgtm, I don't think this is currently done.

do we want to queue the retry then? What happens if we process both, would Windmill make sure only 1 is committed w/ the exactly once processing?

in more detail if we take the most recent observed item, then would we have to remove the older Work(sameWorkToken, existingCacheToken) from the queue, and add the more recent or passed in Work(sameWorkToken, newCacheToken) to the queue?

discussed offline, will follow up in next cl

scwhittle · 2023-09-26T13:12:31Z

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

+                    .build());
+  }
+
+  synchronized CommitsPendingCountAndActiveWorkStatus getPendingCommitsAndPrintActiveWorkAt(


what about just changing this to printTo(PrintWriter writer) and doing the printing here?
Or just returning a String that we print?

I don't think we need to expose this for testing (don't see test anyway yet) and the current logic is a bit spread out with MAX_PRINTABLE_COMMIT_PENDING_KEYS used here when building and in printing. The header for the table is also separated from the population of the table here. Just having all of that in a single method instead of introducing CommitsPendingCountAndActiveWorkStatus seems cleaner.

scwhittle · 2023-09-26T13:13:33Z

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

+          execute(work);
+          return true;
+        }
+        // This will never happen, the switch is exhaustive.


move inside default case

scwhittle · 2023-09-26T13:16:16Z

...taflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java

@@ -120,6 +120,11 @@ public Collection<Windmill.LatencyAttribution> getLatencyAttributions() {
    return list;
  }

+  boolean isStuckAt(Instant stuckCommitDeadline) {


isStuckCommittingAt

scwhittle · 2023-09-26T13:19:58Z

...ker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java

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


This seems odd to me because I don't think this is a legal state for the class. If somethig is active it is in the queue, if it becomes active and there is nothing it should clean itself up.

It seems like it would be better to test through the interface of the class instead of injecting internal state members, maybe verifying internal state but not directly manipulating it from the test.

For example here you coudl do the above, add an item, activate it, complete it, verify that the key doesn't exist in the internal map.

done im also going to remove all of the tests that we can't "simulate" with the external API. changes the test map to readonly as well.

scwhittle · 2023-09-26T13:21:53Z

...ker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java

+    ShardedKey shardedKey = shardedKey("someKey", 1L);
+    Deque<Work> workQueue = new ArrayDeque<>();
+    workQueue.addLast(workInQueue);
+    activeWork.put(shardedKey, workQueue);


ditto, this seems like it should just be combined with the above QUEUED test which sets up the right state, instead of hard-coding the state here which might not match the actual implementation.

If there was a bug and queueing the work returned QUEUED but didn't actually put it in the internal state, then the above test would pass and this test would pass.

Ditto for the below

scwhittle

Thanks, minor test improvement but otherwise feel free to ask someone else for review/merge.

scwhittle · 2023-09-27T08:29:13Z

...ker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkStateTest.java

+            shardedKey, workToBeCompleted.getWorkItem().getWorkToken());
+
+    assertTrue(nextWorkOpt.isPresent());
+    assertSame(nextWork, nextWorkOpt.get());


would be good to then complete this returned work and verify that works

Use ActiveWorkState class instead of an activeWorkMap in ComputationState

m-trieu · 2023-09-27T18:49:20Z

Thanks for reviewing @scwhittle @kennknowles !
@AnandInguva I squashed the commits, and this is ready for merge.

Thank you!

github-actions bot added runners dataflow labels Sep 19, 2023

m-trieu force-pushed the mtrieu-refactor branch from fc2d91f to 6a4c788 Compare September 19, 2023 19:08

github-actions bot added the Next Action: Reviewers label Sep 19, 2023

m-trieu commented Sep 19, 2023

View reviewed changes

scwhittle requested changes Sep 21, 2023

View reviewed changes

m-trieu force-pushed the mtrieu-refactor branch 2 times, most recently from b8229e0 to 766b440 Compare September 21, 2023 22:35

scwhittle requested changes Sep 22, 2023

View reviewed changes

m-trieu force-pushed the mtrieu-refactor branch from 766b440 to 6e2aba8 Compare September 25, 2023 04:57

scwhittle requested changes Sep 25, 2023

View reviewed changes

scwhittle requested changes Sep 26, 2023

View reviewed changes

m-trieu force-pushed the mtrieu-refactor branch from 175900c to ce0ada3 Compare September 26, 2023 18:57

scwhittle approved these changes Sep 27, 2023

View reviewed changes

Break out nested classes from StreamingDataflowWorker.

89b946a

Use ActiveWorkState class instead of an activeWorkMap in ComputationState

m-trieu force-pushed the mtrieu-refactor branch from ce0ada3 to 89b946a Compare September 27, 2023 18:45

AnandInguva merged commit 170310e into apache:master Sep 28, 2023
17 checks passed

m-trieu mentioned this pull request Oct 20, 2023

check for cachetoken representing a retry before activating and completing work #29082

Merged

3 tasks

m-trieu deleted the mtrieu-refactor branch June 10, 2024 15:41


		import org.apache.beam.runners.dataflow.worker.windmill.Windmill;

		public class KeyCommitTooLargeException extends Exception {

break nested classes out of StreamingDataflowWorker #28537

break nested classes out of StreamingDataflowWorker #28537

Conversation

m-trieu commented Sep 19, 2023

r: @kennknowles

GitHub Actions Tests Status (on master branch)

kennknowles commented Sep 19, 2023

m-trieu commented Sep 19, 2023

github-actions bot commented Sep 19, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

m-trieu commented Sep 20, 2023

github-actions bot commented Sep 20, 2023

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

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

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