m-trieu commented on code in PR #31902:
URL: https://github.com/apache/beam/pull/31902#discussion_r1757360784
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java:
##########
@@ -272,42 +263,106 @@ private void startWorkerMetadataConsumer() {
public synchronized void shutdown() {
Preconditions.checkState(started, "StreamingEngineClient never started.");
getWorkerMetadataStream.get().halfClose();
- getWorkBudgetRefresher.stop();
newWorkerMetadataPublisher.shutdownNow();
newWorkerMetadataConsumer.shutdownNow();
channelCachingStubFactory.shutdown();
}
+ private void consumeWindmillWorkerEndpoints(WindmillEndpoints
newWindmillEndpoints) {
+ consumeEndpoints(newWindmillEndpoints).join();
+ }
+
/**
* {@link java.util.function.Consumer<WindmillEndpoints>} used to update
{@link #connections} on
* new backend worker metadata.
*/
- private synchronized void consumeWindmillWorkerEndpoints(WindmillEndpoints
newWindmillEndpoints) {
- isBudgetRefreshPaused.set(true);
+ private synchronized CompletableFuture<Void> consumeEndpoints(
+ WindmillEndpoints newWindmillEndpoints) {
LOG.info("Consuming new windmill endpoints: {}", newWindmillEndpoints);
ImmutableMap<Endpoint, WindmillConnection> newWindmillConnections =
createNewWindmillConnections(newWindmillEndpoints.windmillEndpoints());
-
+ CompletableFuture<Void> closeStaleStreams =
+ closeStaleStreams(newWindmillConnections.values(),
connections.get().windmillStreams());
+ ImmutableMap<WindmillConnection, WindmillStreamSender> newStreams =
+ createAndStartNewStreams(newWindmillConnections.values()).join();
StreamingEngineConnectionState newConnectionsState =
StreamingEngineConnectionState.builder()
.setWindmillConnections(newWindmillConnections)
- .setWindmillStreams(
-
closeStaleStreamsAndCreateNewStreams(newWindmillConnections.values()))
+ .setWindmillStreams(newStreams)
.setGlobalDataStreams(
createNewGlobalDataStreams(newWindmillEndpoints.globalDataEndpoints()))
.build();
-
LOG.info(
"Setting new connections: {}. Previous connections: {}.",
newConnectionsState,
connections.get());
connections.set(newConnectionsState);
- isBudgetRefreshPaused.set(false);
- getWorkBudgetRefresher.requestBudgetRefresh();
+ getWorkBudgetDistributor.distributeBudget(newStreams.values(),
totalGetWorkBudget);
+
+ // Close the streams outside the lock.
+ return closeStaleStreams;
+ }
+
+ /** Close the streams that are no longer valid asynchronously. */
+ private CompletableFuture<Void> closeStaleStreams(
+ Collection<WindmillConnection> newWindmillConnections,
+ ImmutableMap<WindmillConnection, WindmillStreamSender> currentStreams) {
+ return CompletableFuture.allOf(
+ currentStreams.entrySet().stream()
+ .filter(
+ connectionAndStream ->
+
!newWindmillConnections.contains(connectionAndStream.getKey()))
+ .map(
+ entry ->
+ CompletableFuture.runAsync(
+ () -> {
+ LOG.debug("Closing streams to {}",
entry.getKey().backendWorkerToken());
Review Comment:
are debug logs always on? how do we access this
i didn't want to add log.warn/info since it could get noisy
--
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: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
