Fix/Refactor GetData interfaces for Direct Path integration #31784

m-trieu · 2024-07-05T09:48:12Z

GetData/RefreshWork changes as part of breaking up #31504

Remove MetricTrackingWindmillServerStub and replace w/ GetDataClient and WorkRefreshClient
GetDataClient has 3 implementations Appliance, StreamingEngine, and Direct
WorkRefreshClient also has 3 implementations Appliance, StreamingEngine, and FanOut

R: @scwhittle

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.

… metric tracking.

github-actions · 2024-07-05T10:35:29Z

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

scwhittle

Thanks for separating this out! It's a lot easier to review than mixed in with the rest of the direct path changes.

Nothing too major. Since I'm out next week feel free to ask someone else to finish up reviewing.

scwhittle · 2024-07-05T21:31:24Z

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

+
+  /** Returns a new {@link Work} instance with the same state and a different failure handler. */
+  public Work withFailureHandler(Runnable onFailed) {
+    return new Work(


assert there was not already a failure handler? maybe just have a setfailurehandler method?

dropped this to add in a different PR

scwhittle · 2024-07-05T21:38:19Z

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

@@ -85,7 +88,7 @@ public static StreamingEngineWorkCommitter create(
  @Override
  @SuppressWarnings("FutureReturnValueIgnored")
  public void start() {
-    if (!commitSenders.isShutdown()) {
+    if (isRunning.compareAndSet(false, true) && !commitSenders.isShutdown()) {


do we call start multiple times?
seems simpler if we just fail on it called multiple times and remove if it is possible

we do not
was just opting for this route to not use locks and to future proof incase changes are made in the future where someone intentionally or not calls start() multiple times. In this case start() and stop() logic are only executed once, no matter how many times called.

I don't think we should throw an exception since we will have to handle it in a caller or else crash, but if this is best can update thie code.

I think for some things like this which we expect to just be started once and live for ever, it can be simpler to just prevent some behavior instead of trying to support it. If start can be called multiple times, one caller could stop while the other doesn't but hte end state is very different
start1/start2/stop1 -> stopped
start1/stop1/start2 -> started
So it seems safer to just prevent it since we only really expect these methods to be called by the main startup path and then in teardown for tests.

ack updating

scwhittle · 2024-07-05T21:39:37Z

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

@@ -104,17 +116,17 @@ public long currentActiveCommitBytes() {

  @Override
  public void stop() {
-    if (!commitSenders.isTerminated()) {
+    if (isRunning.compareAndSet(true, false) && !commitSenders.isTerminated()) {


ditto, can we just enforce stop called once and after start?

scwhittle · 2024-07-05T21:45:05Z

.../org/apache/beam/runners/dataflow/worker/windmill/client/getdata/ApplianceGetDataClient.java

+      return;
+    }
+
+    for (Map.Entry<HeartbeatSender, Heartbeat> heartbeatToSend : heartbeats.entrySet()) {


should jsut be a single sender in this case, could assert and simplify code

scwhittle · 2024-07-05T21:47:40Z

...che/beam/runners/dataflow/worker/windmill/client/getdata/ThrottlingGetDataMetricTracker.java

+    return () ->
+        getDataMetrics
+            .activeHeartbeats()
+            .getAndUpdate(existing -> Math.max(existing - numHeartbeats, 0));


think you can get rid of the max now that we're no longer setting to 0

scwhittle · 2024-07-05T21:59:14Z

...ava/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/DirectHeartbeatSender.java

+
+  @Override
+  public int hashCode() {
+    return getDataStream.hashCode();


could use Objects.hash(DirectHeartbeatSender.class, getDataStream);

scwhittle · 2024-07-05T22:01:07Z

...ava/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/DirectHeartbeatSender.java

+ *     stream instance.
+ */
+@Internal
+public final class DirectHeartbeatSender implements HeartbeatSender {


maybe FixedStreamHeartbeatSender/ SingleStreamHeartbeatSender ?

Since the code woudl work if the stream was direct or to the dispatcher

scwhittle · 2024-07-05T22:04:47Z

...r/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/Heartbeat.java

+
+/** Heartbeat requests and the work that was used to generate the heartbeat requests. */
+@AutoValue
+public abstract class Heartbeat {


Seems like this shoudl be Heartbeats or HeartbeatsBatch , something to indicate it is multiple.

Can we change the colleciton of work and requests to be immutable?
Since just created one place you could use some ImmutalbeMap builder there and change create to take those immutable maps.

Or could add some builder class for this that has immutable map builder inside of it.

...er/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java

scwhittle · 2024-07-05T22:15:16Z

.../java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java

+      // Get heartbeat requests for computation's current active work, aggregated by GetDataStream
+      // to correctly fan-out the heartbeat requests.
+      Table<HeartbeatSender, RefreshableWork, HeartbeatRequest> heartbeats =
+          HeartbeatRequests.getRefreshableKeyHeartbeats(


this seems to be the only use of this and it doesn't really use the table, just iterates over the cells. Can this method instead be changed to add to the heartbeatsender -> heartbeats builder? (see other comment)

Seems like we're giong to have a lot of unneeded allocations to build the rows and tables instead of just building into what we want in the end.

… Future proofs this logic for direct path.

github-actions · 2024-07-12T15:06:11Z

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

R: @damondouglas 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).

scwhittle · 2024-07-18T01:11:32Z

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

@@ -147,16 +150,25 @@ private static LatencyAttribution.Builder createLatencyAttributionWithActiveLate
    return latencyAttribution;
  }

+  public RefreshableWork refreshableView() {


can we remove this? can it just auto-cast? or in the case where perhaps it isn't for streams just do it via a cast there?

scwhittle · 2024-07-18T01:16:40Z

...va/worker/src/test/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSourcesTest.java

        Instant::now,
        Collections.emptyList());
  }

+  private static GetDataClient createMockGetDataClient() {


why does mock(GetDataClient) not work?
ditto for other tests

if it doesn't for some reason, can we have a shared factory for this for tests?

scwhittle · 2024-07-23T17:30:42Z

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

+
+  ShardedKey getShardedKey();
+
+  boolean isRefreshable(Instant refreshDeadline);


the meaning of refreshDeadline isn't particularly clear. We could add a comment but I'm wondering if it makes more sense to have this determination outside of refreshablework class and instead when we are generating refreshablework.

Maybe should we just expose the start time or last refresh time and just compare when generating refreshablework?

scwhittle · 2024-07-23T17:41:07Z

...er/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java

+  void shutdown();
+
+  /** Reflects that {@link #shutdown()} was explicitly called. */
+  boolean isShutdown();


should this be removed? shutdown indicates that no further interaction with the stream should happen after shutdown. So if external callers are checking shutdown state before doing something it seems racy.

scwhittle · 2024-07-23T17:47:05Z

...er/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java

+  /** Reflects that {@link #shutdown()} was explicitly called. */
+  boolean isShutdown();
+
+  Type streamType();


why is this needed? I don't see any callers and
generally it is a code-smell because it means that if we add a type we have to update scattered usages of the type where if instead whatever callers are switching on the type of was part of the interface everything is grouped.

https://code-smells.com/object-orientation-abusers/switch-statement

scwhittle · 2024-07-23T17:49:40Z

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

@@ -94,7 +97,16 @@ public void start() {

  @Override
  public void commit(Commit commit) {
-    commitQueue.put(commit);
+    if (commit.work().isFailed() || !isRunning.get()) {


the log should only be if !isRunning

the commit work could be failed due to heartbeat response and the log woudl be inaccurate

scwhittle · 2024-07-23T17:53:27Z

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

@@ -85,7 +88,7 @@ public static StreamingEngineWorkCommitter create(
  @Override
  @SuppressWarnings("FutureReturnValueIgnored")
  public void start() {
-    if (!commitSenders.isShutdown()) {
+    if (isRunning.compareAndSet(false, true) && !commitSenders.isShutdown()) {


I think for some things like this which we expect to just be started once and live for ever, it can be simpler to just prevent some behavior instead of trying to support it. If start can be called multiple times, one caller could stop while the other doesn't but hte end state is very different
start1/start2/stop1 -> stopped
start1/stop1/start2 -> started
So it seems safer to just prevent it since we only really expect these methods to be called by the main startup path and then in teardown for tests.

...che/beam/runners/dataflow/worker/windmill/client/getdata/ThrottlingGetDataMetricTracker.java

...va/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java

.../java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java

...main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/HeartbeatSender.java

...rg/apache/beam/runners/dataflow/worker/windmill/work/refresh/FixedStreamHeartbeatSender.java

scwhittle

A few more comments but nothing major!

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

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

scwhittle · 2024-07-29T18:57:07Z

...er/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java

-      return Type.GET_WORKER_METADATA;
+  interface GetWorkerMetadataStream extends WindmillStream {}
+
+  class WindmillStreamShutdownException extends RuntimeException {


could this be moved into AbstractWindmillStream?

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

.../src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/Heartbeats.java

...rg/apache/beam/runners/dataflow/worker/windmill/work/refresh/FixedStreamHeartbeatSender.java

arunpandianp · 2024-08-05T07:32:03Z

.../java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java

+      }
+    }
+
+    sendHeartbeatSafely(firstHeartbeat);


firstHeartbeat here will be null, if heartbeatsBySender is empty

I see exceptions like

java.lang.NullPointerException: Cannot invoke "java.util.Map$Entry.getValue()" because "heartbeat" is null at org.apache.beam.runners.dataflow.worker.windmill.work.refresh.ActiveWorkRefresher.sendHeartbeatSafely(ActiveWorkRefresher.java:182) at org.apache.beam.runners.dataflow.worker.windmill.work.refresh.ActiveWorkRefresher.refreshActiveWork(ActiveWorkRefresher.java:149) at org.apache.beam.runners.dataflow.worker.windmill.work.refresh.ActiveWorkRefresher.lambda$start$0(ActiveWorkRefresher.java:94) at java.base/java.util.concurrent.Executors$RunnableAdapter.call(Executors.java:572) at java.base/java.util.concurrent.FutureTask.runAndReset(FutureTask.java:358) at java.base/java.util.concurrent.ScheduledThreadPoolExecutor$ScheduledFutureTask.run(ScheduledThreadPoolExecutor.java:305) at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1144) at java.base/java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:642) at java.base/java.lang.Thread.run(Thread.java:1583)

Oops we lost the empty case in refactoring, sending PR.

Sgtm! I'll approve

m-trieu added 2 commits July 5, 2024 02:46

add getdata clients and shareable logic for GetData w/ throttling and…

01cfef4

… metric tracking.

add HeartbeatSender interface

4ba59e3

github-actions bot added runners dataflow labels Jul 5, 2024

scwhittle requested changes Jul 5, 2024

View reviewed changes

m-trieu added 2 commits July 12, 2024 00:38

address PR comments

f6cceaf

move side input state API into Work instance w/ keyed state fetching.…

4693870

… Future proofs this logic for direct path.

github-actions bot added the Next Action: Reviewers label Jul 12, 2024

don't let FanOutWorkRefreshClient exceptions crash the user worker

3e7f124

m-trieu force-pushed the mt-get-data branch from 7974c98 to 3e7f124 Compare July 13, 2024 03:29

This was referenced Jul 13, 2024

Add WorkProvider interfaces and implementations #31883

Merged

Integrate direct path #31902

Open

scwhittle requested changes Jul 18, 2024

View reviewed changes

scwhittle requested changes Jul 23, 2024

View reviewed changes

m-trieu force-pushed the mt-get-data branch 7 times, most recently from 34d960b to 15e982d Compare July 27, 2024 03:44

address PR comments

d5acf0f

m-trieu force-pushed the mt-get-data branch from 15e982d to d5acf0f Compare July 27, 2024 03:45

scwhittle requested changes Jul 29, 2024

View reviewed changes

m-trieu and others added 2 commits July 30, 2024 16:12

address PR comments

908a0e0

Merge branch 'master' into mt-get-data

492ac5a

scwhittle approved these changes Jul 30, 2024

View reviewed changes

scwhittle merged commit 570f2f8 into apache:master Jul 30, 2024
17 checks passed

arunpandianp reviewed Aug 5, 2024

View reviewed changes


		ShardedKey getShardedKey();

		boolean isRefreshable(Instant refreshDeadline);

Fix/Refactor GetData interfaces for Direct Path integration #31784

Fix/Refactor GetData interfaces for Direct Path integration #31784

Conversation

m-trieu commented Jul 5, 2024

GitHub Actions Tests Status (on master branch)

github-actions bot commented Jul 5, 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

m-trieu Jul 12, 2024 • edited Loading

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

github-actions bot commented Jul 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

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 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

m-trieu Jul 12, 2024 •

edited

Loading