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



##########
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:
       sure.

##########
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:
       I'm not sure I got what you mean, isn't it already a CompletableFuture.




-- 
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:
[email protected]


Reply via email to