feat: introduce closeAsync to Batcher

gax/src/main/java/com/google/api/gax/batching/Batcher.java

@@ -31,6 +31,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.BetaApi;
+import com.google.api.core.InternalExtensionOnly;
 
 /**
  * Represents a batching context where individual elements will be accumulated and flushed in a
@@ -45,6 +46,7 @@
  * @param <ElementResultT> The type of the result for each individual element.
  */
 @BetaApi("The surface for batching is not stable yet and may change in the future.")
+@InternalExtensionOnly
 public interface Batcher<ElementT, ElementResultT> extends AutoCloseable {
 
   /**
@@ -79,4 +81,10 @@ public interface Batcher<ElementT, ElementResultT> extends AutoCloseable {
    */
   @Override
   void close() throws InterruptedException;
+
+  /**
+   * Closes this Batcher by preventing new elements from being added and send outstanding elements.
+   * The returned future will be resolved when the last element completes
+   */
+  ApiFuture<Void> closeAsync();
 }

gax/src/main/java/com/google/api/gax/batching/BatcherImpl.java

@@ -52,6 +52,7 @@
 import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -89,7 +90,7 @@ public class BatcherImpl<ElementT, ElementResultT, RequestT, ResponseT>
   private final Object flushLock = new Object();
   private final Object elementLock = new Object();
   private final Future<?> scheduledFuture;
-  private volatile boolean isClosed = false;
+  private SettableApiFuture<Void> closeFuture;
   private final BatcherStats batcherStats = new BatcherStats();
   private final FlowController flowController;
 
@@ -172,7 +173,10 @@ public BatcherImpl(
   /** {@inheritDoc} */
   @Override
   public ApiFuture<ElementResultT> add(ElementT element) {
-    Preconditions.checkState(!isClosed, "Cannot add elements on a closed batcher");
+    // Note: there is no need to synchronize over closeFuture. The write & read of the variable
+    // will only be done from a single calling thread.
+    Preconditions.checkState(closeFuture == null, "Cannot add elements on a closed batcher");
+
     // This is not the optimal way of throttling. It does not send out partial batches, which
     // means that the Batcher might not use up all the resources allowed by FlowController.
     // The more efficient implementation should look like:
@@ -257,9 +261,20 @@ public void onFailure(Throwable throwable) {
   }
 
   private void onBatchCompletion() {
-    if (numOfOutstandingBatches.decrementAndGet() == 0) {
-      synchronized (flushLock) {
+    boolean shouldClose = false;
+
+    synchronized (flushLock) {
+      if (numOfOutstandingBatches.decrementAndGet() == 0) {
         flushLock.notifyAll();
+        shouldClose = closeFuture != null;
+      }
+    }
+    if (shouldClose) {
+      BatchingException overallError = batcherStats.asException();
+      if (overallError != null) {
+        closeFuture.setException(overallError);
+      } else {
+        closeFuture.set(null);
       }
     }
   }
@@ -279,17 +294,57 @@ private void awaitAllOutstandingBatches() throws InterruptedException {
   /** {@inheritDoc} */
   @Override
   public void close() throws InterruptedException {
-    if (isClosed) {
-      return;
+    try {
+      closeAsync().get();
+    } catch (ExecutionException e) {
+      // Original stacktrace of a batching exception is not useful, so rethrow the error with
+      // the caller stacktrace
+      if (e.getCause() instanceof BatchingException) {
+        BatchingException cause = (BatchingException) e.getCause();
+        throw new BatchingException(cause.getMessage());
+      } else {
+        throw new IllegalStateException("unexpected error closing the batcher", e.getCause());
+      }
+    }
+  }
+
+  @Override
+  public ApiFuture<Void> closeAsync() {
+    if (closeFuture != null) {
+      return closeFuture;
+    }
+
+    // Send any buffered elements
+    // Must come before numOfOutstandingBatches check below
+    sendOutstanding();
+
+    boolean closeImmediately;
+
+    synchronized (flushLock) {
+      // prevent admission of new elements
+      closeFuture = SettableApiFuture.create();
+      // check if we can close immediately
+      closeImmediately = numOfOutstandingBatches.get() == 0;
     }
-    flush();
+
+    // Clean up accounting
     scheduledFuture.cancel(true);
-    isClosed = true;
     currentBatcherReference.closed = true;
     currentBatcherReference.clear();
-    BatchingException exception = batcherStats.asException();
-    if (exception != null) {
-      throw exception;
+
+    // notify futures
+    if (closeImmediately) {
+      finishClose();
+    }
+    return closeFuture;
+  }
+
+  private void finishClose() {
+    BatchingException batchingException = batcherStats.asException();
+    if (batchingException != null) {
+      closeFuture.setException(batchingException);
+    } else {
+      closeFuture.set(null);
     }
   }
 

gax/src/test/java/com/google/api/gax/batching/BatcherImplTest.java

@@ -92,7 +92,11 @@ public class BatcherImplTest {
   @After
   public void tearDown() throws InterruptedException {
     if (underTest != null) {
-      underTest.close();
+      try {
+        underTest.close();
+      } catch (BatchingException ignored) {
+
+      }
     }
   }