[GitHub] [beam] kileys commented on a change in pull request #14741: [BEAM-12294] Implement close function for BeamFnStatusClient to shutd…

[GitBox <gi...@apache.org>]

kileys commented on a change in pull request #14741:
URL: https://github.com/apache/beam/pull/14741#discussion_r627622716



##########
File path: sdks/java/harness/src/main/java/org/apache/beam/fn/harness/status/BeamFnStatusClient.java
##########
@@ -42,9 +44,12 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class BeamFnStatusClient {
+public class BeamFnStatusClient implements AutoCloseable {
+  private static final Object COMPLETED = new Object();
   private final StreamObserver<WorkerStatusResponse> outboundObserver;
   private final BundleProcessorCache processBundleCache;
+  private final ManagedChannel channel;
+  private final CompletableFuture<Object> inboundObserverCompletion;

Review comment:
       Can we use a CompletableFuture<Void> instead?

##########
File path: sdks/java/harness/src/main/java/org/apache/beam/fn/harness/status/BeamFnStatusClient.java
##########
@@ -53,18 +58,32 @@ public BeamFnStatusClient(
       Function<ApiServiceDescriptor, ManagedChannel> channelFactory,
       BundleProcessorCache processBundleCache,
       PipelineOptions options) {
-    BeamFnWorkerStatusGrpc.BeamFnWorkerStatusStub stub =
-        BeamFnWorkerStatusGrpc.newStub(channelFactory.apply(apiServiceDescriptor));
-    this.outboundObserver = stub.workerStatus(new InboundObserver());
+    this.channel = channelFactory.apply(apiServiceDescriptor);
+    this.outboundObserver =
+        BeamFnWorkerStatusGrpc.newStub(channel).workerStatus(new InboundObserver());
     this.processBundleCache = processBundleCache;
     this.memoryMonitor = MemoryMonitor.fromOptions(options);
+    this.inboundObserverCompletion = new CompletableFuture<>();
     Thread thread = new Thread(memoryMonitor);
     thread.setDaemon(true);
     thread.setPriority(Thread.MIN_PRIORITY);
     thread.setName("MemoryMonitor");
     thread.start();
   }
 
+  @Override
+  public void close() throws Exception {
+    try {
+      inboundObserverCompletion.get();

Review comment:
       Should this get a timeout as well?




-- 
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.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org