Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Dataflow Streaming] Use concurrent queue for commit queue #33690

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -17,25 +17,29 @@
*/
package org.apache.beam.runners.dataflow.worker.streaming;

import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.ConcurrentLinkedQueue;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;

/** Queue bounded by a {@link WeightedSemaphore}. */
public final class WeightedBoundedQueue<V> {
public final class WeightedBoundedQueue<V extends @NonNull Object> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we might want to make this less generally named or at least comment on spinning so we know that before it gets used somewhere else.


private final LinkedBlockingQueue<V> queue;
private final ConcurrentLinkedQueue<V> queue;
private final WeightedSemaphore<V> weightedSemaphore;

private WeightedBoundedQueue(
LinkedBlockingQueue<V> linkedBlockingQueue, WeightedSemaphore<V> weightedSemaphore) {
this.queue = linkedBlockingQueue;
ConcurrentLinkedQueue<V> queue, WeightedSemaphore<V> weightedSemaphore) {
this.queue = queue;
this.weightedSemaphore = weightedSemaphore;
}

public static <V> WeightedBoundedQueue<V> create(WeightedSemaphore<V> weightedSemaphore) {
return new WeightedBoundedQueue<>(new LinkedBlockingQueue<>(), weightedSemaphore);
public static <V extends @NonNull Object> WeightedBoundedQueue<V> create(
WeightedSemaphore<V> weightedSemaphore) {
return new WeightedBoundedQueue<>(new ConcurrentLinkedQueue<>(), weightedSemaphore);
}

/**
Expand All @@ -60,14 +64,29 @@ public void put(V value) {
* Retrieves and removes the head of this queue, waiting up to the specified wait time if
* necessary for an element to become available.
*
* @param timeout how long to wait before giving up, in units of {@code unit}
* @param unit a {@code TimeUnit} determining how to interpret the {@code timeout} parameter
* @param timeout how long to wait before giving up
* @return the head of this queue, or {@code null} if the specified waiting time elapses before an
* element is available
* @throws InterruptedException if interrupted while waiting
*/
public @Nullable V poll(long timeout, TimeUnit unit) throws InterruptedException {
@Nullable V result = queue.poll(timeout, unit);
public @Nullable V poll(Duration timeout) throws InterruptedException {
@Nullable V result;
Instant deadline = Instant.now().plus(timeout);
int spin = 0;
while (true) {
if (++spin > 1000) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how much cpu does this burn if the queue is empty?

we could also look into other specialized queues like:
https://github.com/JCTools/JCTools/blob/master/jctools-core/src/main/java/org/jctools/queues/MpscBlockingConsumerArrayQueue.java
(random googling, would need some validation).

Thread.sleep(1);
result = queue.poll();
if (result != null || Instant.now().isAfter(deadline)) {
break;
}
spin = 0;
}
result = queue.poll();
if (result != null) {
break;
}
}
if (result != null) {
weightedSemaphore.release(result);
}
Expand All @@ -76,7 +95,7 @@ public void put(V value) {

/** Returns and removes the next value, or blocks until one is available. */
public V take() throws InterruptedException {
V result = queue.take();
V result = Preconditions.checkNotNull(poll(Duration.ofDays(1000)));
weightedSemaphore.release(result);
return result;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
package org.apache.beam.runners.dataflow.worker.windmill.client.commits;

import com.google.auto.value.AutoBuilder;
import java.time.Duration;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
Expand Down Expand Up @@ -227,7 +228,7 @@ private boolean tryAddToCommitBatch(Commit commit, CommitWorkStream.RequestBatch
Commit commit;
try {
if (commits < TARGET_COMMIT_BATCH_KEYS) {
commit = commitQueue.poll(10 - 2L * commits, TimeUnit.MILLISECONDS);
commit = commitQueue.poll(Duration.ofMillis(10 - 2L * commits));
} else {
commit = commitQueue.poll();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNull;

import java.util.concurrent.TimeUnit;
import java.time.Duration;
import java.util.concurrent.atomic.AtomicInteger;
import javax.annotation.Nullable;
import org.junit.Rule;
Expand Down Expand Up @@ -119,7 +119,7 @@ public void testPoll_withTimeout() throws InterruptedException {
() -> {
int polled;
try {
polled = queue.poll(pollWaitTimeMillis, TimeUnit.MILLISECONDS);
polled = queue.poll(Duration.ofMillis(pollWaitTimeMillis));
pollResult.set(polled);
} catch (InterruptedException e) {
throw new RuntimeException(e);
Expand Down Expand Up @@ -152,7 +152,7 @@ public void testPoll_withTimeout_timesOut() throws InterruptedException {
() -> {
@Nullable Integer polled;
try {
polled = queue.poll(pollWaitTimeMillis, TimeUnit.MILLISECONDS);
polled = queue.poll(Duration.ofMillis(pollWaitTimeMillis));
if (polled != null) {
pollResult.set(polled);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
import java.util.Map;
import org.apache.beam.sdk.expansion.ExternalTransformRegistrar;
import org.apache.beam.sdk.io.aws2.common.ClientConfiguration;
import org.apache.beam.sdk.io.aws2.kinesis.KinesisIO;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

separate change?

import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ExternalTransformBuilder;
import org.apache.beam.sdk.transforms.PTransform;
Expand Down
Loading