Add correct shutdown mechanics for direct path #31902

m-trieu · 2024-07-16T13:47:58Z

Add correct shutdown mechanics and integrate direct path

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.

github-actions · 2024-07-16T14:34:54Z

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

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

Abacn · 2024-07-22T15:19:50Z

R: @scwhittle

github-actions · 2024-07-22T15:22:00Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

scwhittle

Not finished, mostly looking at stream stuff to help find what is causing stuckness

scwhittle · 2024-07-24T18:53:37Z

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

  /** Send a request to the server. */
  protected final void send(RequestT request) {
    lastSendTimeMs.set(Instant.now().getMillis());
    synchronized (this) {
+      // Check if we should send after we acquire the lock.
+      if (isShutdown()) {
+        LOG.warn("Send called on a shutdown stream.");


if this is possible, don't log as customers don't like warning logs and open issues about them

if this should not be possible, perhaps better to throw an exception so that we notice and fix it.

scwhittle · 2024-07-24T18:54:35Z

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

    }
  }

  /** Starts the underlying stream. */
  protected final void startStream() {
    // Add the stream to the registry after it has been fully constructed.
    streamRegistry.add(this);
-    while (true) {
+    while (!isShutdown.get()) {


I'd just remove this check since you do it first thing below

scwhittle · 2024-07-24T19:00:44Z

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

-          return;
-        }
+    private void tryRestartStream() {
+      if (!isShutdown()) {


this looks racy, we check shutdown above in isStreamDone with synchronization, but then if it is shutdown before here, we end up with an error but won't restart the stream or remove it from the registry.

I would remove this one

scwhittle · 2024-07-24T19:01:21Z

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

-        setLastError(error);
+    private synchronized boolean isStreamDone() {
+      if (isShutdown() || (clientClosed.get() && !hasPendingRequests())) {
+        streamRegistry.remove(AbstractWindmillStream.this);


I think this kind of side effect is confusing in method that just sounds like an accessor

how about maybeTeardownStream()?

Add a comment on return value

scwhittle · 2024-08-07T10:46:47Z

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

@@ -176,7 +180,7 @@ private StreamingDataflowWorker(
      DataflowWorkerHarnessOptions options,
      HotKeyLogger hotKeyLogger,
      Supplier<Instant> clock,
-      StreamingWorkerStatusReporter workerStatusReporter,
+      Function<Supplier<Long>, StreamingWorkerStatusReporter> streamingWorkerStatusReporterFactory,


should this be a functional interface so you can document? it's unclear what the long supplier is

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

scwhittle · 2024-08-07T11:47:14Z

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

+        () -> {
+          try {
+            send(extension);
+          } catch (IllegalStateException e) {


should we handle illegalstateexception internally? or rely on executeSafely to catch it?

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

scwhittle · 2024-08-07T11:56:39Z

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

@@ -182,6 +184,12 @@ protected void onResponse(StreamingCommitResponse response) {
    }
  }

+  @Override
+  protected void shutdownInternal() {
+    pending.values().forEach(pendingRequest -> pendingRequest.onDone.accept(CommitStatus.ABORTED));


I'm worried if pending has somethign inserted between iterating and clear

can you instead use an iterator where you remove as you go so everything removed is guaranteed to be aborted?

scwhittle · 2024-08-07T12:00:56Z

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

@@ -163,7 +165,7 @@ protected void onResponse(StreamingCommitResponse response) {
        continue;
      }
      PendingRequest done = pending.remove(requestId);
-      if (done == null) {
+      if (done == null && !isShutdown()) {


this is going to get nullptr exception below, instead move the isShutdown check to whether or not to log inside this if

scwhittle · 2024-09-12T10:30:27Z

...src/main/java/org/apache/beam/runners/dataflow/options/DataflowStreamingPipelineOptions.java

@@ -302,4 +301,12 @@ public Integer create(PipelineOptions options) {
      return streamingOptions.isEnableStreamingEngine() ? Integer.MAX_VALUE : 1;
    }
  }
+
+  /** EnableStreamingEngine defaults to false unless one of the two experiments is set. */


looks like just 1 experiment

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

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

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

...pache/beam/runners/dataflow/worker/windmill/client/grpc/WindmillStreamShutdownException.java

scwhittle · 2024-09-12T12:58:17Z

...apache/beam/runners/dataflow/worker/windmill/client/grpc/observers/DirectStreamObserver.java

+   * Only send the next value if the phaser is not terminated by the time we acquire the lock since
+   * the phaser can be terminated at any time.
+   */
+  private void tryOnNext(T value) {


not sure these phaser checks are necessary since the outboundObserver itself should stop blocking for onNext if the notifier is terminated.

but if the phaser is terminated, we don't want to call outboundObserver.onNext() right?

Even if you do check here there is still race between checking the phaser and calling onNext regardless. Internally the outboundObserver is already observing the phaser termination via getPhase() (and also blocking respecting phaser termination) so the extra check is just mental overhead I think.

scwhittle · 2024-09-12T13:00:03Z

...apache/beam/runners/dataflow/worker/windmill/client/grpc/observers/DirectStreamObserver.java

          return;
        }
      } catch (TimeoutException e) {
+        if (isReadyNotifier.isTerminated()) {


don't think you need this, awaitAdvanceInterruptibly will return -1 if it's terminated.

scwhittle · 2024-09-12T13:01:10Z

...apache/beam/runners/dataflow/worker/windmill/client/grpc/observers/DirectStreamObserver.java

@@ -39,7 +39,9 @@
 @ThreadSafe
 public final class DirectStreamObserver<T> implements StreamObserver<T> {


maybe this change can be submitted separately? Would be nice to have a test for it showing the previous bug as well.

scwhittle · 2024-09-12T13:05:01Z

...g/apache/beam/runners/dataflow/worker/windmill/work/budget/EvenGetWorkBudgetDistributor.java

@@ -67,17 +70,18 @@ public <T extends GetWorkBudgetSpender> void distributeBudget(
      GetWorkBudgetSpender getWorkBudgetSpender = streamAndDesiredBudget.getKey();
      GetWorkBudget desired = streamAndDesiredBudget.getValue();
      GetWorkBudget remaining = getWorkBudgetSpender.remainingBudget();
-      if (isBelowFiftyPercentOfTarget(remaining, desired)) {
+      if (isBelowFiftyPercentOfTarget(remaining, desired) && isActiveWorkBudgetAware) {


this is unclear to me. It seems like in one case the adjustment is addtiive and another is resetting

Wouldn't we want to only increase by the desired-remaining in either case?

whether its overriding or additive is depending on the internal implementation of GetWorkBudgetSpender.adjustBudget()

based on the current implementation that this is reflecting, we do not account for remaining budget here it is handled internally in adjustBudget.

scwhittle

didn't get through everything but sending comments I had

scwhittle · 2024-09-13T08:41:29Z

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

+            + "<th>Active For</th>"
+            + "<th>State</th>"
+            + "<th>State Active For</th>"
+            + "<th>Produced By</th>"


nit: produced is confusing since we use that for shuffle terminology. How about just "Backend"

scwhittle · 2024-09-13T08:44:58Z

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

-    getWorkBudgetRefresher.requestBudgetRefresh();
+    getWorkBudgetDistributor.distributeBudget(newStreams.values(), totalGetWorkBudget);
+
+    // Close the streams outside the lock.


yeah I don't know if we need to block on closing streams. I think we might want to block on creating the new ones.

one possible thing is that if we are leaking stuff and they never actually close we might not know. But if we have the stream registry I think we'd see that. Or we can add some logging to closeAllStreams that it is taking a long time.

scwhittle · 2024-09-13T09:07:57Z

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

    clientClosed.set(true);
-    requestObserver().onCompleted();
+    requestObserver.onCompleted();


why is it safe to not synchronize? The comment indicated it was for a reason. In general I think that we need to synchronize using the request observer because grpc doesn't expect it to be used simultaneously by multiple threads.

let me try to run a pipeline w synchronization

previously we were calling relying on this to shutdown, so we had removed synchronization but since we arent anymore it could be ok

ran ok in the test pipeline

done.

scwhittle · 2024-09-13T09:11:42Z

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

+    private final Supplier<StreamObserver<RequestT>> requestObserverSupplier;
+
+    @GuardedBy("this")
+    private volatile @Nullable StreamObserver<RequestT> delegateRequestObserver;


You don't need both, synchronization also ensures that memory changes are viewed by other threads. Volatile just ensures one fields's changes are visible to other threads.

This seems like a good overview: https://blogs.oracle.com/javamagazine/post/java-thread-synchronization-volatile-final-atomic-deadlocks

scwhittle · 2024-09-13T09:12:22Z

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

+    }
+
+    @Override
+    public synchronized void onCompleted() {


rm synchronized?

scwhittle · 2024-09-13T11:59:37Z

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

-    }
+    streamRegistry.stream()
+        .sorted(Comparator.comparing(AbstractWindmillStream::backendWorkerToken))
+        .collect(


I missed the multimap bit. I agree the grouping is nice. However isn't the sorting then unnecessary if we're putting it in a map right away?

scwhittle · 2024-09-13T12:01:39Z

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

@@ -264,24 +284,18 @@ protected void startThrottleTimer() {

  @Override
  public void adjustBudget(long itemsDelta, long bytesDelta) {


can we name this setBudget?

scwhittle · 2024-09-13T12:06:43Z

...ain/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java

+      } catch (AppendableInputStream.InvalidInputStreamStateException
+          | VerifyException
+          | CancellationException e) {
+        handleShutdown(request);


Can we make shutdown handling consistent?
we throw an exception for shutdown here but if we don't run the loop due to isShutdown we just return from this function without doing anything.

we do throw a WindmillStreamShutdownException at the bottom if we exit the loop

scwhittle · 2024-09-13T12:07:10Z

...ain/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcGetDataStream.java

+  private void handleShutdown(QueuedRequest request) {
+    if (isShutdown()) {
+      throw new WindmillStreamShutdownException(
+          "Cannot send request=[" + request + "] on closed stream.");


if we keep this, should we pass in the exception above to add as a suppressed exception?

scwhittle · 2024-09-13T12:10:13Z

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

+      }
+    }
+
+    ImmutableList<String> createStreamCancelledErrorMessage() {


add parameter for the limit or just inline above? We're building up a possibly big list just to ignore most of it.

… metric tracking.

… Future proofs this logic for direct path.

… have it implement WorkProvider interface. Move class to windmill/work/provider directory, update visibility for dependent classes and move tests, add GetWorkBudgetOwnerInterface

…ct path

m-trieu · 2024-09-20T06:51:04Z

@scwhittle ready for another look! I will resolve the merge conflicts and rebase

m-trieu changed the title ~~Mt dp streams~~ Add correct shutdown mechanics for direct path Jul 16, 2024

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

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

m-trieu force-pushed the mt-dp-streams branch 3 times, most recently from 01e8586 to ff287f2 Compare August 6, 2024 12:25

scwhittle requested changes Aug 7, 2024

View reviewed changes

m-trieu force-pushed the mt-dp-streams branch 2 times, most recently from bbf2778 to 1e02463 Compare September 12, 2024 05:06

scwhittle requested changes Sep 12, 2024

View reviewed changes

m-trieu force-pushed the mt-dp-streams branch 2 times, most recently from 4820b9f to d03b0bc Compare September 13, 2024 05:19

scwhittle requested changes Sep 13, 2024

View reviewed changes

m-trieu force-pushed the mt-dp-streams branch 2 times, most recently from ade21ba to d95d1b5 Compare September 19, 2024 22:11

m-trieu added 11 commits September 19, 2024 15:12

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

ad30a20

… metric tracking.

add HeartbeatSender interface

c36781d

address PR comments

94e111e

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

3c1737f

… Future proofs this logic for direct path.

don't let FanOutWorkRefreshClient exceptions crash the user worker

be0209e

address PR comments

9058872

address PR comments

3805df6

add WorkProvider interface and have StreamingEngine and Appliance impls

62f7c3c

change StreamingEngineClient to FanOutStreamingEngineWorkProvider and…

f93b351

… have it implement WorkProvider interface. Move class to windmill/work/provider directory, update visibility for dependent classes and move tests, add GetWorkBudgetOwnerInterface

add WorkProvider interface and have StreamingEngine and Appliance impls

f583ef1

change StreamingEngineClient to FanOutStreamingEngineWorkProvider and…

96615c9

… have it implement WorkProvider interface. Move class to windmill/work/provider directory, update visibility for dependent classes and move tests, add GetWorkBudgetOwnerInterface

m-trieu added 9 commits September 19, 2024 15:12

add WorkProvider interface and have StreamingEngine and Appliance impls

f718a48

change StreamingEngineClient to FanOutStreamingEngineWorkProvider and…

552254f

… have it implement WorkProvider interface. Move class to windmill/work/provider directory, update visibility for dependent classes and move tests, add GetWorkBudgetOwnerInterface

add stream management and correct stream shutdown mechanisms for dire…

09af93e

…ct path

integrate direct path components

d5a3f26

address PR comments

1655a49

simplify budgeting logic for GrpcDirectGetWorkStream

5ff1897

fix stuckness on range movements

701bf77

address PR comments, fix broken tests

0db2045

address PR comments

9bad2d7

m-trieu force-pushed the mt-dp-streams branch from d95d1b5 to 9bad2d7 Compare September 20, 2024 03:53

		@@ -39,7 +39,9 @@
		@ThreadSafe
		public final class DirectStreamObserver<T> implements StreamObserver<T> {

		@@ -264,24 +284,18 @@ protected void startThrottleTimer() {

		@Override
		public void adjustBudget(long itemsDelta, long bytesDelta) {

Add correct shutdown mechanics for direct path #31902

Are you sure you want to change the base?

Add correct shutdown mechanics for direct path #31902

Conversation

m-trieu commented Jul 16, 2024 • edited Loading

R: @scwhittle

GitHub Actions Tests Status (on master branch)

github-actions bot commented Jul 16, 2024

Abacn commented Jul 22, 2024

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

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

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

m-trieu commented Sep 20, 2024

m-trieu commented Jul 16, 2024 •

edited

Loading