Merge pull request apache#14728 from ihji/cherry-pick-12118

[BEAM-12262] Cherry-pick BEAM-12118

sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/QueueingBeamFnDataClient.java

@@ -83,12 +83,22 @@ boolean offer(ConsumerAndData<?> e, long l, TimeUnit t) throws InterruptedExcept
     // empty in which case it returns null.
     @Nullable
     ConsumerAndData<?> take() throws InterruptedException {
+      // We first poll without blocking to optimize for the case there is data.
+      // If there is no data we end up blocking on take() and thus the extra
+      // poll doesn't matter.
       @Nullable ConsumerAndData<?> result = queue.poll();
       if (result == null) {
         if (closed.get()) {
-          return null;
+          // Poll again to ensure that there is nothing in the queue. Once we observe closed as true
+          // we are guaranteed no additional elements other than the POISON will be added. However
+          // we can't rely on the previous poll result as it could race with additional offers and
+          // close.
+          result = queue.poll();
+        } else {
+          // We are not closed so we perform a blocking take. We are guaranteed that additional
+          // elements will be offered or the POISON will be added by close to unblock this thread.
+          result = queue.take();
         }
-        result = queue.take();
       }
       if (result == POISON) {
         return null;