m-trieu commented on code in PR #31902:
URL: https://github.com/apache/beam/pull/31902#discussion_r1760443991
##########
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.
Review Comment:
done
--
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
