extract semaphore logic out of WeightBoundedQueue to allow for sharing the weigher #32905

m-trieu · 2024-10-22T22:24:25Z

Factor out semaphore logic out of WeightBoundedQueue to allow for sharing the weigher

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-10-22T23:06:37Z

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-10-24T00:07:23Z

@scwhittle ready to for review, the failure to unrelated thanks!

m-trieu · 2024-10-25T05:25:55Z

Run Java Precommit

scwhittle · 2024-10-30T10:28:34Z

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

@@ -118,19 +118,17 @@ public final class StreamingDataflowWorker {
   */
  public static final int MAX_SINK_BYTES = 10_000_000;

+  public static final String STREAMING_ENGINE_USE_JOB_SETTINGS_FOR_HEARTBEAT_POOL =


revert this file

scwhittle · 2024-10-30T10:30:31Z

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

+    return new WeightedSemaphore<>(maxWeight, new Semaphore(maxWeight, true), weigherFn);
+  }
+
+  void acquire(V value) {


public acquire and release? seems like basic functionality

scwhittle · 2024-10-30T10:30:55Z

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

+    return new WeightedSemaphore<>(maxWeight, new Semaphore(maxWeight, true), weigherFn);
+  }
+
+  void acquire(V value) {


name acquireUninterruptibly?

scwhittle · 2024-10-30T10:33:28Z

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

@@ -61,11 +62,10 @@ public final class StreamingEngineWorkCommitter implements WorkCommitter {
      Supplier<CloseableStream<CommitWorkStream>> commitWorkStreamFactory,
      int numCommitSenders,
      Consumer<CompleteCommit> onCommitComplete,
-      String backendWorkerToken) {
+      String backendWorkerToken,
+      WeightedSemaphore<Commit> weigher) {


rename weigher throughout to commitByteSemaphore or something capturing it is the limiter not just a weighing function?

scwhittle · 2024-10-30T10:35:42Z

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

+    Thread putThread =
+        new Thread(
+            () -> {
+              try {


remove the sleep? the element is already added by queue1 above so the sleep doesn't seem necessary

scwhittle · 2024-10-30T10:36:59Z

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

    }
    return result;
  }

  /** Returns and removes the next value, or blocks until one is available. */
  public @Nullable V take() throws InterruptedException {
    V result = queue.take();
-    limit.release(weigher.apply(result));
+    weigher.release(result);
    return result;
  }

  /** Returns the current weight of the queue. */


can this be removed and just look at the weighted semaphore? It is confusing since it doesn't count just bytes in this queue.

scwhittle · 2024-10-30T10:39:28Z

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

-  private final int maxWeight;
-  private final Semaphore limit;
-  private final Function<V, Integer> weigher;
+  private final WeightedSemaphore<V> weigher;


name weightedSemaphore? weigher just sounds like it is the weigh function

scwhittle · 2024-10-30T10:40:55Z

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

@@ -86,16 +86,19 @@ public final class StreamingEngineWorkCommitter implements WorkCommitter {
  public static Builder builder() {
    return new AutoBuilder_StreamingEngineWorkCommitter_Builder()
        .setBackendWorkerToken(NO_BACKEND_WORKER_TOKEN)
+        .setWeigher(
+            WeightedSemaphore.create(


maybe require this to be set by caller instead of hiding it? It would be nice for the callers to know about it because then they could display it on their status page

scwhittle · 2024-10-30T10:42:34Z

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

+            });
+    putThread.start();
+
+    // Should only see the first value in the queue, since the queue is at capacity. putThread


could put the sleep here if you want to give a little more time for putThread to run and become blocked (doesn't matter if it doesn't or not so not racy but better test if it does block).

…to be injected and shared

m-trieu · 2024-11-02T08:00:24Z

back to you @scwhittle thank you

github-actions · 2024-11-02T08:36:56Z

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

R: @Abacn 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

minor cleanups

scwhittle · 2024-11-05T09:44:58Z

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

-  public static <V> WeightedBoundedQueue<V> create(int maxWeight, Function<V, Integer> weigherFn) {
-    return new WeightedBoundedQueue<>(
-        new LinkedBlockingQueue<>(), maxWeight, new Semaphore(maxWeight, true), weigherFn);
+  public static <V> WeightedBoundedQueue<V> create(WeightedSemaphore<V> weigher) {


nit: weightedSemaphore

scwhittle · 2024-11-05T09:49:35Z

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

+  }
+
+  public void acquireUninterruptibly(V value) {
+    limit.acquireUninterruptibly(weigher.apply(value));


How about capping the calculated weight to not be more than the max so it will be the single item available to process?

Otherwise it seems that we coudl block forever. this is currently done in the specific weigher logic but seems safer to do here and in release.

scwhittle · 2024-11-05T09:50:25Z

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

+
+  public static WeightedSemaphore<Commit> maxCommitByteSemaphore() {
+    return WeightedSemaphore.create(
+        MAX_QUEUED_COMMITS_BYTES, commit -> Math.min(MAX_QUEUED_COMMITS_BYTES, commit.getSize()));


see other comment, remove the min here.

scwhittle · 2024-11-05T09:50:35Z

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

-            MAX_COMMIT_QUEUE_BYTES, commit -> Math.min(MAX_COMMIT_QUEUE_BYTES, commit.getSize()));
+            WeightedSemaphore.create(
+                MAX_COMMIT_QUEUE_BYTES,
+                commit -> Math.min(MAX_COMMIT_QUEUE_BYTES, commit.getSize())));


remove min.

scwhittle · 2024-11-05T09:52:46Z

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

+    Thread.sleep(100);
+    assertEquals(MAX_WEIGHT, weigher.currentWeight());
+    assertEquals(1, queue1.size());
+    assertEquals(MAX_WEIGHT, weigher.currentWeight());


rm, doing this twice.

scwhittle · 2024-11-05T09:53:29Z

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

+    putThread.join();
+
+    assertEquals(MAX_WEIGHT, weigher.currentWeight());
+    assertEquals(1, queue2.size());


poll from queue2 and verify the weight goes to zero.

m-trieu · 2024-11-05T19:21:09Z

back to you thanks! @scwhittle

scwhittle · 2024-11-05T20:31:36Z

Run Java PreCommit

github-actions bot added runners dataflow labels Oct 22, 2024

m-trieu force-pushed the mt-shared-weigher branch from 94041c4 to 881d43f Compare October 22, 2024 22:25

scwhittle requested changes Oct 30, 2024

View reviewed changes

m-trieu added 2 commits October 31, 2024 18:01

extract semaphore logic out of WeightBoundedQueue to allow a weigher …

ec3b662

…to be injected and shared

address PR comments

27a9758

m-trieu force-pushed the mt-shared-weigher branch from 881d43f to 27a9758 Compare November 2, 2024 07:53

github-actions bot added the Next Action: Reviewers label Nov 2, 2024

scwhittle requested changes Nov 5, 2024

View reviewed changes

address PR comments

f614330

scwhittle approved these changes Nov 5, 2024

View reviewed changes

extract semaphore logic out of WeightBoundedQueue to allow for sharing the weigher #32905

Are you sure you want to change the base?

extract semaphore logic out of WeightBoundedQueue to allow for sharing the weigher #32905

Conversation

m-trieu commented Oct 22, 2024

GitHub Actions Tests Status (on master branch)

github-actions bot commented Oct 22, 2024

m-trieu commented Oct 24, 2024

m-trieu commented Oct 25, 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

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 Nov 2, 2024

github-actions bot commented Nov 2, 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

m-trieu commented Nov 5, 2024

scwhittle commented Nov 5, 2024