[GitHub] [beam] amaliujia commented on a change in pull request #14480: [BEAM-12118] Modify QueuingBeamFnDataClient to avoid completion latency due to polling

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

amaliujia commented on a change in pull request #14480:
URL: https://github.com/apache/beam/pull/14480#discussion_r610067279



##########
File path: sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClient.java
##########
@@ -71,23 +85,28 @@ public InboundDataClient receive(
     InboundDataClient inboundDataClient =
         this.mainClient.receive(apiServiceDescriptor, inputLocation, queueingConsumer);
     queueingConsumer.inboundDataClient = inboundDataClient;
-    this.inboundDataClients.computeIfAbsent(
-        inboundDataClient, (InboundDataClient idcToStore) -> idcToStore);
+    synchronized (inboundDataClients) {
+      Preconditions.checkState(!isDraining);
+      if (this.inboundDataClients.add(inboundDataClient)) {
+        inboundDataClient.runWhenComplete(() -> completeInbound(inboundDataClient));
+      }
+    }
     return inboundDataClient;
   }
 
-  // Returns true if all the InboundDataClients have finished or cancelled and no values
-  // remain on the queue.
-  private boolean allDone() {
-    for (InboundDataClient inboundDataClient : inboundDataClients.keySet()) {
-      if (!inboundDataClient.isDone()) {
-        return false;
+  private void completeInbound(InboundDataClient client) {
+    Preconditions.checkState(client.isDone());
+    boolean addPoison = false;
+    synchronized (inboundDataClients) {
+      Preconditions.checkState(inboundDataClients.remove(client));
+      finishedClients.add(client);
+      if (inboundDataClients.isEmpty() && isDraining) {
+        addPoison = true;

Review comment:
       I found a hard to understand why adding this boolean flag? To indicate this is draining?  




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