scwhittle commented on code in PR #31784:
URL: https://github.com/apache/beam/pull/31784#discussion_r1695742159
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/getdata/StreamPoolGetDataClient.java:
##########
@@ -18,33 +18,29 @@
package org.apache.beam.runners.dataflow.worker.windmill.client.getdata;
import java.io.PrintWriter;
-import java.util.Map;
import javax.annotation.concurrent.ThreadSafe;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import
org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
import
org.apache.beam.runners.dataflow.worker.windmill.Windmill.KeyedGetDataRequest;
import org.apache.beam.runners.dataflow.worker.windmill.client.CloseableStream;
import
org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.GetDataStream;
import
org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStreamPool;
-import
org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender;
-import
org.apache.beam.runners.dataflow.worker.windmill.work.refresh.Heartbeats;
import org.apache.beam.sdk.annotations.Internal;
-import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
/**
* StreamingEngine implementation of {@link GetDataClient}.
*
- * @implNote Uses {@link WindmillStreamPool} to send/receive requests.
Depending on options, may use
- * a dedicated stream pool for heartbeats.
+ * @implNote Uses {@link WindmillStreamPool} to send requests. Depending on
options, may use a
+ * dedicated stream pool for heartbeats.
Review Comment:
update comment, no longer sends heartbeats
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java:
##########
@@ -206,14 +206,17 @@ synchronized ImmutableListMultimap<ShardedKey,
RefreshableWork> getReadOnlyActiv
.collect(
flatteningToImmutableListMultimap(
Entry::getKey,
- e ->
e.getValue().stream().map(ExecutableWork::work).map(Work::refreshableView)));
+ e ->
+ e.getValue().stream()
+ .map(ExecutableWork::work)
Review Comment:
nit: combine to single map?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java:
##########
@@ -333,14 +331,8 @@ private StreamingDataflowWorker(
ID_GENERATOR,
stageInfoMap);
- this.heartbeatSender =
- options.isEnableStreamingEngine()
- ? new StreamPoolHeartbeatSender(
- options.getUseSeparateWindmillHeartbeatStreams()
- ? WindmillStreamPool.create(
- 1, GET_DATA_STREAM_TIMEOUT,
windmillServer::getDataStream)
- : getDataStreamPool)
- : new ApplianceHeartbeatSender(windmillServer::getData);
+ // Register standard file systems.
Review Comment:
nit: can we move this to the top? it affects some executors used within
filesystems and seems safer to do it before we start other stuff.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java:
##########
@@ -115,32 +132,72 @@ private void invalidateStuckCommits() {
}
}
+ /** Create {@link Heartbeats} and group them by {@link HeartbeatSender}. */
private void refreshActiveWork() {
Instant refreshDeadline =
clock.get().minus(Duration.millis(activeWorkRefreshPeriodMillis));
+ Map<HeartbeatSender, Heartbeats> heartbeatsBySender =
+ aggregateHeartbeatsBySender(refreshDeadline);
+ if (heartbeatsBySender.isEmpty()) {
+ return;
+ }
+
+ if (heartbeatsBySender.size() == 1) {
+ // If there is a single HeartbeatSender, just use the calling thread to
send heartbeats.
+ Map.Entry<HeartbeatSender, Heartbeats> heartbeat =
+ Iterables.getOnlyElement(heartbeatsBySender.entrySet());
+ sendHeartbeat(heartbeat);
+ } else {
+ // If there are multiple HeartbeatSenders, send out the heartbeats in
parallel using the
+ // fanOutActiveWorkRefreshExecutor.
+ List<CompletableFuture<Void>> fanOutRefreshActiveWork = new
ArrayList<>();
+ for (Map.Entry<HeartbeatSender, Heartbeats> heartbeat :
heartbeatsBySender.entrySet()) {
+ fanOutRefreshActiveWork.add(
+ CompletableFuture.runAsync(
+ () -> sendHeartbeat(heartbeat),
fanOutActiveWorkRefreshExecutor));
+ }
+
+ // Don't block until we kick off all the refresh active work RPCs.
Review Comment:
It would be equivalent to just loop over fanOutRefreshActiveWork and join
each individually (if we're not surfacing exceptions from them it behaves the
same and seems likely less code).
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java:
##########
@@ -115,32 +132,72 @@ private void invalidateStuckCommits() {
}
}
+ /** Create {@link Heartbeats} and group them by {@link HeartbeatSender}. */
private void refreshActiveWork() {
Instant refreshDeadline =
clock.get().minus(Duration.millis(activeWorkRefreshPeriodMillis));
+ Map<HeartbeatSender, Heartbeats> heartbeatsBySender =
+ aggregateHeartbeatsBySender(refreshDeadline);
+ if (heartbeatsBySender.isEmpty()) {
+ return;
+ }
+
+ if (heartbeatsBySender.size() == 1) {
+ // If there is a single HeartbeatSender, just use the calling thread to
send heartbeats.
+ Map.Entry<HeartbeatSender, Heartbeats> heartbeat =
+ Iterables.getOnlyElement(heartbeatsBySender.entrySet());
+ sendHeartbeat(heartbeat);
+ } else {
+ // If there are multiple HeartbeatSenders, send out the heartbeats in
parallel using the
+ // fanOutActiveWorkRefreshExecutor.
+ List<CompletableFuture<Void>> fanOutRefreshActiveWork = new
ArrayList<>();
+ for (Map.Entry<HeartbeatSender, Heartbeats> heartbeat :
heartbeatsBySender.entrySet()) {
+ fanOutRefreshActiveWork.add(
+ CompletableFuture.runAsync(
+ () -> sendHeartbeat(heartbeat),
fanOutActiveWorkRefreshExecutor));
+ }
+
+ // Don't block until we kick off all the refresh active work RPCs.
+ @SuppressWarnings("rawtypes")
+ CompletableFuture<Void> parallelFanOutRefreshActiveWork =
+ CompletableFuture.allOf(fanOutRefreshActiveWork.toArray(new
CompletableFuture[0]));
+ parallelFanOutRefreshActiveWork.join();
+ }
+ }
+
+ /** Aggregate the heartbeats across computations by HeartbeatSender for
correct fan out. */
+ private Map<HeartbeatSender, Heartbeats> aggregateHeartbeatsBySender(Instant
refreshDeadline) {
Map<HeartbeatSender, Heartbeats.Builder> heartbeatsBySender = new
HashMap<>();
// Aggregate the heartbeats across computations by HeartbeatSender for
correct fan out.
for (ComputationState computationState : computations.get()) {
for (RefreshableWork work :
computationState.getRefreshableWork(refreshDeadline)) {
heartbeatsBySender
.computeIfAbsent(work.heartbeatSender(), ignored ->
Heartbeats.builder())
- .addWork(work)
- .addHeartbeatRequest(computationState.getComputationId(),
createHeartbeatRequest(work));
+ .add(computationState.getComputationId(), work, sampler);
}
}
- heartbeatSender.accept(
- heartbeatsBySender.entrySet().stream()
- .collect(toImmutableMap(Map.Entry::getKey, e ->
e.getValue().build())));
+ return heartbeatsBySender.entrySet().stream()
+ .collect(toImmutableMap(Map.Entry::getKey, e -> e.getValue().build()));
+ }
+
+ private void sendHeartbeat(Map.Entry<HeartbeatSender, Heartbeats> heartbeat)
{
Review Comment:
could add noException or something to method name to help see above that it
is safe to call inline or on executor without handling.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/FixedStreamHeartbeatSender.java:
##########
@@ -49,15 +50,24 @@ public static FixedStreamHeartbeatSender
create(GetDataStream getDataStream) {
@Override
public void sendHeartbeats(Heartbeats heartbeats) {
- if (getDataStream.isShutdown()) {
+ String threadName = Thread.currentThread().getName();
+ try {
+ String backendWorkerToken = getDataStream.backendWorkerToken();
+ if (!backendWorkerToken.isEmpty()) {
+ // Decorate the thread name w/ the backendWorkerToken for debugging.
Resets the thread's
+ // name after sending the heartbeats succeeds or fails.
+ Thread.currentThread().setName(threadName + "-" + backendWorkerToken);
Review Comment:
how about we avoid the getName/setName if no backend worker token.
Think you could have threadName start empty, then intialize it here if token
was nonempty, and then below set it back only if threadname is nonempty.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/Heartbeats.java:
##########
@@ -45,21 +46,32 @@ public abstract static class Builder {
abstract ImmutableList.Builder<RefreshableWork> workBuilder();
- public final Builder addWork(RefreshableWork work) {
+ public final Builder add(
+ String computationId, RefreshableWork work,
DataflowExecutionStateSampler sampler) {
workBuilder().add(work);
+ addHeartbeatRequest(computationId, createHeartbeatRequest(work,
sampler));
return this;
}
+ private Windmill.HeartbeatRequest createHeartbeatRequest(
+ RefreshableWork work, DataflowExecutionStateSampler sampler) {
+ return Windmill.HeartbeatRequest.newBuilder()
+ .setShardingKey(work.getShardedKey().shardingKey())
+ .setWorkToken(work.id().workToken())
+ .setCacheToken(work.id().cacheToken())
+
.addAllLatencyAttribution(work.getHeartbeatLatencyAttributions(sampler))
+ .build();
+ }
+
abstract Builder setHeartbeatRequests(
Review Comment:
can this be removed so that add is used?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/Heartbeats.java:
##########
@@ -45,21 +46,32 @@ public abstract static class Builder {
abstract ImmutableList.Builder<RefreshableWork> workBuilder();
- public final Builder addWork(RefreshableWork work) {
+ public final Builder add(
+ String computationId, RefreshableWork work,
DataflowExecutionStateSampler sampler) {
workBuilder().add(work);
+ addHeartbeatRequest(computationId, createHeartbeatRequest(work,
sampler));
return this;
}
+ private Windmill.HeartbeatRequest createHeartbeatRequest(
+ RefreshableWork work, DataflowExecutionStateSampler sampler) {
+ return Windmill.HeartbeatRequest.newBuilder()
+ .setShardingKey(work.getShardedKey().shardingKey())
+ .setWorkToken(work.id().workToken())
+ .setCacheToken(work.id().cacheToken())
+
.addAllLatencyAttribution(work.getHeartbeatLatencyAttributions(sampler))
+ .build();
+ }
+
abstract Builder setHeartbeatRequests(
ImmutableListMultimap<String, Windmill.HeartbeatRequest> value);
abstract ImmutableListMultimap.Builder<String, Windmill.HeartbeatRequest>
heartbeatRequestsBuilder();
Review Comment:
can this be private?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/Heartbeats.java:
##########
@@ -45,21 +46,32 @@ public abstract static class Builder {
abstract ImmutableList.Builder<RefreshableWork> workBuilder();
Review Comment:
can this be private?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/Heartbeats.java:
##########
@@ -45,21 +46,32 @@ public abstract static class Builder {
abstract ImmutableList.Builder<RefreshableWork> workBuilder();
- public final Builder addWork(RefreshableWork work) {
+ public final Builder add(
+ String computationId, RefreshableWork work,
DataflowExecutionStateSampler sampler) {
workBuilder().add(work);
+ addHeartbeatRequest(computationId, createHeartbeatRequest(work,
sampler));
return this;
}
+ private Windmill.HeartbeatRequest createHeartbeatRequest(
+ RefreshableWork work, DataflowExecutionStateSampler sampler) {
+ return Windmill.HeartbeatRequest.newBuilder()
+ .setShardingKey(work.getShardedKey().shardingKey())
+ .setWorkToken(work.id().workToken())
+ .setCacheToken(work.id().cacheToken())
+
.addAllLatencyAttribution(work.getHeartbeatLatencyAttributions(sampler))
+ .build();
+ }
+
abstract Builder setHeartbeatRequests(
ImmutableListMultimap<String, Windmill.HeartbeatRequest> value);
abstract ImmutableListMultimap.Builder<String, Windmill.HeartbeatRequest>
heartbeatRequestsBuilder();
- public final Builder addHeartbeatRequest(
+ private void addHeartbeatRequest(
Review Comment:
just inline above and remove?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java:
##########
@@ -214,11 +216,13 @@ private boolean tryAddToCommitBatch(Commit commit,
CommitWorkStream.RequestBatch
return isCommitAccepted;
}
- // Helper to batch additional commits into the commit batch as long as they
fit.
- // Returns a commit that was removed from the queue but not consumed or null.
- private Commit expandBatch(CommitWorkStream.RequestBatcher batcher) {
+ /**
+ * Helper to batch additional commits into the commit batch as long as they
fit. Returns a commit
+ * that was removed from the queue but not consumed or null.
+ */
+ private @Nullable Commit expandBatch(CommitWorkStream.RequestBatcher
batcher) {
int commits = 1;
- while (true) {
+ while (isRunning.get()) {
Review Comment:
nit: I'd just remove this one, it adds overhead to streams with many keys
and we're already checking between each batch
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/getdata/ThrottlingGetDataMetricTracker.java:
##########
@@ -26,121 +26,84 @@
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
/**
- * Wraps GetData calls that tracks metrics for the number of in-flight
requests and throttles
- * requests when memory pressure is high.
+ * Wraps GetData calls to track metrics for the number of in-flight requests
and throttles requests
+ * when memory pressure is high.
*/
@Internal
@ThreadSafe
public final class ThrottlingGetDataMetricTracker {
+ private static final String GET_STATE_DATA_RESOURCE_CONTEXT = "GetStateData";
+ private static final String GET_SIDE_INPUT_RESOURCE_CONTEXT =
"GetSideInputData";
+
private final MemoryMonitor gcThrashingMonitor;
- private final GetDataMetrics getDataMetrics;
+ private final AtomicInteger activeStateReads;
+ private final AtomicInteger activeSideInputs;
+ private final AtomicInteger activeHeartbeats;
public ThrottlingGetDataMetricTracker(MemoryMonitor gcThrashingMonitor) {
this.gcThrashingMonitor = gcThrashingMonitor;
- this.getDataMetrics = GetDataMetrics.create();
+ this.activeStateReads = new AtomicInteger();
+ this.activeSideInputs = new AtomicInteger();
+ this.activeHeartbeats = new AtomicInteger();
+ }
+
+ /**
+ * Tracks a state data fetch. If there is memory pressure, may throttle
requests. Returns an
+ * {@link AutoCloseable} that will decrement the metric after the call is
finished.
+ */
+ AutoCloseable trackStateDataFetchWithThrottling() {
+ gcThrashingMonitor.waitForResources(GET_STATE_DATA_RESOURCE_CONTEXT);
+ activeStateReads.getAndIncrement();
+ return activeStateReads::getAndDecrement;
}
/**
- * Tracks a GetData call. If there is memory pressure, may throttle
requests. Returns an {@link
- * AutoCloseable} that will decrement the metric after the call is finished.
+ * Tracks a side input fetch. If there is memory pressure, may throttle
requests. Returns an
+ * {@link AutoCloseable} that will decrement the metric after the call is
finished.
*/
- public AutoCloseable trackSingleCallWithThrottling(Type callType) {
- gcThrashingMonitor.waitForResources(callType.debugName);
- AtomicInteger getDataMetricTracker = getDataMetrics.getMetricFor(callType);
- getDataMetricTracker.getAndIncrement();
- return getDataMetricTracker::getAndDecrement;
+ AutoCloseable trackSideInputFetchWithThrottling() {
+ gcThrashingMonitor.waitForResources(GET_SIDE_INPUT_RESOURCE_CONTEXT);
+ activeSideInputs.getAndIncrement();
+ return activeSideInputs::getAndDecrement;
}
/**
* Tracks heartbeat request metrics. Returns an {@link AutoCloseable} that
will decrement the
* metric after the call is finished.
*/
public AutoCloseable trackHeartbeats(int numHeartbeats) {
- getDataMetrics
- .activeHeartbeats()
- .getAndUpdate(currentActiveHeartbeats -> currentActiveHeartbeats +
numHeartbeats);
- return () ->
- getDataMetrics.activeHeartbeats().getAndUpdate(existing -> existing -
numHeartbeats);
+ activeHeartbeats.getAndUpdate(
Review Comment:
prefer getAndAdd(numHeartbeats)/getAndAdd(-numHeartbeats)
can likely use native instructions instead of cas and easier to read
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java:
##########
@@ -69,40 +71,44 @@ private StreamingEngineWorkCommitter(
new ThreadFactoryBuilder()
.setDaemon(true)
.setPriority(Thread.MAX_PRIORITY)
- .setNameFormat("CommitThread-%d")
+ .setNameFormat(
+ backendWorkerToken.isEmpty()
+ ? "CommitThread-%d"
+ : "CommitThread-" + backendWorkerToken + "-%d")
.build());
this.activeCommitBytes = new AtomicLong();
this.onCommitComplete = onCommitComplete;
this.numCommitSenders = numCommitSenders;
this.isRunning = new AtomicBoolean(false);
}
- public static StreamingEngineWorkCommitter create(
- Supplier<CloseableStream<CommitWorkStream>> commitWorkStreamFactory,
- int numCommitSenders,
- Consumer<CompleteCommit> onCommitComplete) {
- return new StreamingEngineWorkCommitter(
- commitWorkStreamFactory, numCommitSenders, onCommitComplete);
+ public static Builder builder() {
+ return new AutoBuilder_StreamingEngineWorkCommitter_Builder()
+ .setBackendWorkerToken("")
Review Comment:
nit: use the NO worker token constant?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java:
##########
@@ -64,14 +72,23 @@ public ActiveWorkRefresher(
Supplier<Collection<ComputationState>> computations,
DataflowExecutionStateSampler sampler,
ScheduledExecutorService activeWorkRefreshExecutor,
- Consumer<Map<HeartbeatSender, Heartbeats>> heartbeatSender) {
+ HeartbeatTracker heartbeatTracker) {
this.clock = clock;
this.activeWorkRefreshPeriodMillis = activeWorkRefreshPeriodMillis;
this.stuckCommitDurationMillis = stuckCommitDurationMillis;
this.computations = computations;
this.sampler = sampler;
this.activeWorkRefreshExecutor = activeWorkRefreshExecutor;
- this.heartbeatSender = heartbeatSender;
+ this.heartbeatTracker = heartbeatTracker;
+ this.fanOutActiveWorkRefreshExecutor =
+ Executors.newCachedThreadPool(
+ new ThreadFactoryBuilder()
+ // Work refresh runs as a background process, don't let
failures crash
+ // the worker.
+ .setUncaughtExceptionHandler(
Review Comment:
Let's remove this, sendHeartbeats is not throwing exceptions now
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/refresh/ActiveWorkRefresher.java:
##########
@@ -115,32 +132,72 @@ private void invalidateStuckCommits() {
}
}
+ /** Create {@link Heartbeats} and group them by {@link HeartbeatSender}. */
private void refreshActiveWork() {
Instant refreshDeadline =
clock.get().minus(Duration.millis(activeWorkRefreshPeriodMillis));
+ Map<HeartbeatSender, Heartbeats> heartbeatsBySender =
+ aggregateHeartbeatsBySender(refreshDeadline);
+ if (heartbeatsBySender.isEmpty()) {
+ return;
+ }
+
+ if (heartbeatsBySender.size() == 1) {
+ // If there is a single HeartbeatSender, just use the calling thread to
send heartbeats.
+ Map.Entry<HeartbeatSender, Heartbeats> heartbeat =
+ Iterables.getOnlyElement(heartbeatsBySender.entrySet());
+ sendHeartbeat(heartbeat);
+ } else {
+ // If there are multiple HeartbeatSenders, send out the heartbeats in
parallel using the
+ // fanOutActiveWorkRefreshExecutor.
+ List<CompletableFuture<Void>> fanOutRefreshActiveWork = new
ArrayList<>();
+ for (Map.Entry<HeartbeatSender, Heartbeats> heartbeat :
heartbeatsBySender.entrySet()) {
+ fanOutRefreshActiveWork.add(
+ CompletableFuture.runAsync(
+ () -> sendHeartbeat(heartbeat),
fanOutActiveWorkRefreshExecutor));
+ }
+
+ // Don't block until we kick off all the refresh active work RPCs.
Review Comment:
Could also combine inlining single element to reduce threads used in fanout
case by 1 doing something like:
// Send one heartbeat using this thread and start rest async.
@Nullable Map.Entry<HeartbeatSender, Heartbeats> inlineHeartbeat = null;
List<CompletableFuture<Void>> fanOutRefreshActiveWork = new ArrayList<>();
for (Map.Entry<HeartbeatSender, Heartbeats> heartbeat :
heartbeatsBySender.entrySet()) {
if (inlineHeartbeat == null) {
inlineHeartbeat = heartbeat;
} else {
// runAsync
}
}
sendHeartbeat(inlineHeartbeat);
// loop and block on async
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java:
##########
@@ -140,10 +113,11 @@ default void close() {
/** Interface for streaming GetWorkerMetadata requests to Windmill. */
@ThreadSafe
- interface GetWorkerMetadataStream extends WindmillStream {
- @Override
- default Type streamType() {
- return Type.GET_WORKER_METADATA;
+ interface GetWorkerMetadataStream extends WindmillStream {}
+
+ class WindmillStreamShutdownException extends RuntimeException {
Review Comment:
could this be moved into AbstractWindmillStream?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
