m-trieu commented on code in PR #32775:
URL: https://github.com/apache/beam/pull/32775#discussion_r1803918211
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java:
##########
@@ -219,191 +201,195 @@ static FanOutStreamingEngineWorkerHarness forTesting(
@SuppressWarnings("ReturnValueIgnored")
@Override
public synchronized void start() {
- Preconditions.checkState(!started, "StreamingEngineClient cannot start
twice.");
- // Starts the stream, this value is memoized.
- getWorkerMetadataStream.get();
- startWorkerMetadataConsumer();
- getWorkBudgetRefresher.start();
+ Preconditions.checkState(!started, "FanOutStreamingEngineWorkerHarness
cannot start twice.");
+ getWorkerMetadataStream =
+ streamFactory.createGetWorkerMetadataStream(
+ dispatcherClient.getWindmillMetadataServiceStubBlocking(),
+ getWorkerMetadataThrottleTimer,
+ this::consumeWorkerMetadata);
started = true;
}
public ImmutableSet<HostAndPort> currentWindmillEndpoints() {
- return connections.get().windmillConnections().keySet().stream()
+ return backends.get().windmillStreams().keySet().stream()
.map(Endpoint::directEndpoint)
.filter(Optional::isPresent)
.map(Optional::get)
- .filter(
- windmillServiceAddress ->
- windmillServiceAddress.getKind() !=
WindmillServiceAddress.Kind.IPV6)
- .map(
- windmillServiceAddress ->
- windmillServiceAddress.getKind() ==
WindmillServiceAddress.Kind.GCP_SERVICE_ADDRESS
- ? windmillServiceAddress.gcpServiceAddress()
- :
windmillServiceAddress.authenticatedGcpServiceAddress().gcpServiceAddress())
+ .map(WindmillServiceAddress::getServiceAddress)
.collect(toImmutableSet());
}
/**
- * Fetches {@link GetDataStream} mapped to globalDataKey if one exists, or
defaults to {@link
- * GetDataStream} pointing to dispatcher.
+ * Fetches {@link GetDataStream} mapped to globalDataKey if or throws {@link
+ * NoSuchElementException} if one is not found.
*/
private GetDataStream getGlobalDataStream(String globalDataKey) {
- return
Optional.ofNullable(connections.get().globalDataStreams().get(globalDataKey))
- .map(Supplier::get)
- .orElseGet(
- () ->
- streamFactory.createGetDataStream(
- dispatcherClient.getWindmillServiceStub(), new
ThrottleTimer()));
- }
-
- @SuppressWarnings("FutureReturnValueIgnored")
- private void startWorkerMetadataConsumer() {
- newWorkerMetadataConsumer.submit(
- () -> {
- while (true) {
- Optional.ofNullable(newWindmillEndpoints.poll())
- .ifPresent(this::consumeWindmillWorkerEndpoints);
- }
- });
+ return
Optional.ofNullable(backends.get().globalDataStreams().get(globalDataKey))
+ .map(GlobalDataStreamSender::get)
+ .orElseThrow(
+ () -> new NoSuchElementException("No endpoint for global data tag:
" + globalDataKey));
}
@VisibleForTesting
@Override
public synchronized void shutdown() {
- Preconditions.checkState(started, "StreamingEngineClient never started.");
- getWorkerMetadataStream.get().halfClose();
- getWorkBudgetRefresher.stop();
- newWorkerMetadataPublisher.shutdownNow();
- newWorkerMetadataConsumer.shutdownNow();
+ Preconditions.checkState(started, "FanOutStreamingEngineWorkerHarness
never started.");
+ Preconditions.checkNotNull(getWorkerMetadataStream).halfClose();
+ workerMetadataConsumer.shutdownNow();
Review Comment:
Done
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java:
##########
@@ -219,191 +201,195 @@ static FanOutStreamingEngineWorkerHarness forTesting(
@SuppressWarnings("ReturnValueIgnored")
@Override
public synchronized void start() {
- Preconditions.checkState(!started, "StreamingEngineClient cannot start
twice.");
- // Starts the stream, this value is memoized.
- getWorkerMetadataStream.get();
- startWorkerMetadataConsumer();
- getWorkBudgetRefresher.start();
+ Preconditions.checkState(!started, "FanOutStreamingEngineWorkerHarness
cannot start twice.");
+ getWorkerMetadataStream =
+ streamFactory.createGetWorkerMetadataStream(
+ dispatcherClient.getWindmillMetadataServiceStubBlocking(),
+ getWorkerMetadataThrottleTimer,
+ this::consumeWorkerMetadata);
started = true;
}
public ImmutableSet<HostAndPort> currentWindmillEndpoints() {
- return connections.get().windmillConnections().keySet().stream()
+ return backends.get().windmillStreams().keySet().stream()
.map(Endpoint::directEndpoint)
.filter(Optional::isPresent)
.map(Optional::get)
- .filter(
- windmillServiceAddress ->
- windmillServiceAddress.getKind() !=
WindmillServiceAddress.Kind.IPV6)
- .map(
- windmillServiceAddress ->
- windmillServiceAddress.getKind() ==
WindmillServiceAddress.Kind.GCP_SERVICE_ADDRESS
- ? windmillServiceAddress.gcpServiceAddress()
- :
windmillServiceAddress.authenticatedGcpServiceAddress().gcpServiceAddress())
+ .map(WindmillServiceAddress::getServiceAddress)
.collect(toImmutableSet());
}
/**
- * Fetches {@link GetDataStream} mapped to globalDataKey if one exists, or
defaults to {@link
- * GetDataStream} pointing to dispatcher.
+ * Fetches {@link GetDataStream} mapped to globalDataKey if or throws {@link
+ * NoSuchElementException} if one is not found.
*/
private GetDataStream getGlobalDataStream(String globalDataKey) {
- return
Optional.ofNullable(connections.get().globalDataStreams().get(globalDataKey))
- .map(Supplier::get)
- .orElseGet(
- () ->
- streamFactory.createGetDataStream(
- dispatcherClient.getWindmillServiceStub(), new
ThrottleTimer()));
- }
-
- @SuppressWarnings("FutureReturnValueIgnored")
- private void startWorkerMetadataConsumer() {
- newWorkerMetadataConsumer.submit(
- () -> {
- while (true) {
- Optional.ofNullable(newWindmillEndpoints.poll())
- .ifPresent(this::consumeWindmillWorkerEndpoints);
- }
- });
+ return
Optional.ofNullable(backends.get().globalDataStreams().get(globalDataKey))
+ .map(GlobalDataStreamSender::get)
+ .orElseThrow(
+ () -> new NoSuchElementException("No endpoint for global data tag:
" + globalDataKey));
}
@VisibleForTesting
@Override
public synchronized void shutdown() {
- Preconditions.checkState(started, "StreamingEngineClient never started.");
- getWorkerMetadataStream.get().halfClose();
- getWorkBudgetRefresher.stop();
- newWorkerMetadataPublisher.shutdownNow();
- newWorkerMetadataConsumer.shutdownNow();
+ Preconditions.checkState(started, "FanOutStreamingEngineWorkerHarness
never started.");
+ Preconditions.checkNotNull(getWorkerMetadataStream).halfClose();
+ workerMetadataConsumer.shutdownNow();
channelCachingStubFactory.shutdown();
}
- /**
- * {@link java.util.function.Consumer<WindmillEndpoints>} used to update
{@link #connections} on
- * new backend worker metadata.
- */
+ private void consumeWorkerMetadata(WindmillEndpoints windmillEndpoints) {
+ synchronized (metadataLock) {
+ // Only process versions greater than what we currently have to prevent
double processing of
+ // metadata. workerMetadataConsumer is single-threaded so we maintain
ordering.
+ if (windmillEndpoints.version() > pendingMetadataVersion) {
+ pendingMetadataVersion = windmillEndpoints.version();
+ workerMetadataConsumer.execute(() ->
consumeWindmillWorkerEndpoints(windmillEndpoints));
+ }
+ }
+ }
+
private synchronized void consumeWindmillWorkerEndpoints(WindmillEndpoints
newWindmillEndpoints) {
- isBudgetRefreshPaused.set(true);
- LOG.info("Consuming new windmill endpoints: {}", newWindmillEndpoints);
- ImmutableMap<Endpoint, WindmillConnection> newWindmillConnections =
- createNewWindmillConnections(newWindmillEndpoints.windmillEndpoints());
+ // Since this is run on a single threaded executor, multiple versions of
the metadata maybe
+ // queued up while a previous version of the windmillEndpoints were being
consumed. Only consume
+ // the endpoints if they are the most current version.
+ synchronized (metadataLock) {
+ if (newWindmillEndpoints.version() < pendingMetadataVersion) {
+ return;
+ }
+ }
- StreamingEngineConnectionState newConnectionsState =
- StreamingEngineConnectionState.builder()
- .setWindmillConnections(newWindmillConnections)
- .setWindmillStreams(
-
closeStaleStreamsAndCreateNewStreams(newWindmillConnections.values()))
+ LOG.debug(
+ "Consuming new endpoints: {}. previous metadata version: {}, current
metadata version: {}",
+ newWindmillEndpoints,
+ activeMetadataVersion,
+ newWindmillEndpoints.version());
+ closeStaleStreams(newWindmillEndpoints);
+ ImmutableMap<Endpoint, WindmillStreamSender> newStreams =
+
createAndStartNewStreams(newWindmillEndpoints.windmillEndpoints()).join();
+ StreamingEngineBackends newBackends =
+ StreamingEngineBackends.builder()
+ .setWindmillStreams(newStreams)
.setGlobalDataStreams(
createNewGlobalDataStreams(newWindmillEndpoints.globalDataEndpoints()))
.build();
+ backends.set(newBackends);
+ getWorkBudgetDistributor.distributeBudget(newStreams.values(),
totalGetWorkBudget);
+ activeMetadataVersion = newWindmillEndpoints.version();
+ }
+
+ /** Close the streams that are no longer valid asynchronously. */
+ @SuppressWarnings("FutureReturnValueIgnored")
+ private void closeStaleStreams(WindmillEndpoints newWindmillEndpoints) {
+ StreamingEngineBackends currentBackends = backends.get();
+ ImmutableMap<Endpoint, WindmillStreamSender> currentWindmillStreams =
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
