feat/add-async-channel-ping

Change-Id: I3e390b6f7a6c9beaee52d74f37ef557629af7759

Author: nicholsl

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java

@@ -96,43 +96,12 @@ public void primeChannel(ManagedChannel managedChannel) {
   }
 
   private void primeChannelUnsafe(ManagedChannel managedChannel) throws IOException {
-    sendPrimeRequests(managedChannel);
+    sendPrimeRequestsBlocking(managedChannel);
   }
 
-  private void sendPrimeRequests(ManagedChannel managedChannel) {
+  private void sendPrimeRequestsBlocking(ManagedChannel managedChannel) {
     try {
-      ClientCall<PingAndWarmRequest, PingAndWarmResponse> clientCall =
-          managedChannel.newCall(
-              BigtableGrpc.getPingAndWarmMethod(),
-              CallOptions.DEFAULT
-                  .withCallCredentials(callCredentials)
-                  .withDeadline(Deadline.after(1, TimeUnit.MINUTES)));
-
-      SettableApiFuture<PingAndWarmResponse> future = SettableApiFuture.create();
-      clientCall.start(
-          new ClientCall.Listener<PingAndWarmResponse>() {
-            PingAndWarmResponse response;
-
-            @Override
-            public void onMessage(PingAndWarmResponse message) {
-              response = message;
-            }
-
-            @Override
-            public void onClose(Status status, Metadata trailers) {
-              if (status.isOk()) {
-                future.set(response);
-              } else {
-                future.setException(status.asException());
-              }
-            }
-          },
-          createMetadata(headers, request));
-      clientCall.sendMessage(request);
-      clientCall.halfClose();
-      clientCall.request(Integer.MAX_VALUE);
-
-      future.get(1, TimeUnit.MINUTES);
+      sendPrimeRequestsAsync(managedChannel).get(1, TimeUnit.MINUTES);
     } catch (Throwable e) {
       // TODO: Not sure if we should swallow the error here. We are pre-emptively swapping
       // channels if the new
@@ -141,6 +110,59 @@ public void onClose(Status status, Metadata trailers) {
     }
   }
 
+  /**
+   * Asynchronously sends a PingAndWarm request.
+   *
+   * @param managedChannel The channel to send the request on.
+   * @return A ListenableFuture that will be completed with the PingAndWarmResponse or an exception.
+   */
+  public SettableApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(
+      ManagedChannel managedChannel) {
+    ClientCall<PingAndWarmRequest, PingAndWarmResponse> clientCall =
+        managedChannel.newCall(
+            BigtableGrpc.getPingAndWarmMethod(),
+            CallOptions.DEFAULT
+                .withCallCredentials(callCredentials)
+                .withDeadline(Deadline.after(1, TimeUnit.MINUTES)));
+
+    SettableApiFuture<PingAndWarmResponse> future = SettableApiFuture.create();
+    clientCall.start(
+        new ClientCall.Listener<PingAndWarmResponse>() {
+          private PingAndWarmResponse response;
+
+          @Override
+          public void onMessage(PingAndWarmResponse message) {
+            response = message;
+          }
+
+          @Override
+          public void onClose(Status status, Metadata trailers) {
+            if (status.isOk()) {
+              future.set(response);
+            } else {
+              // Propagate the gRPC error to the future.
+              future.setException(status.asException(trailers));
+            }
+          }
+        },
+        createMetadata(headers, request));
+
+    try {
+      // Send the request message.
+      clientCall.sendMessage(request);
+      // Signal that no more messages will be sent.
+      clientCall.halfClose();
+      // Request the response from the server.
+      clientCall.request(Integer.MAX_VALUE);
+    } catch (Throwable t) {
+      // If sending fails, cancel the call and notify the future.
+      clientCall.cancel("Failed to send priming request", t);
+      future.setException(t);
+    }
+
+    return future;
+  }
+
   private static Metadata createMetadata(Map<String, String> headers, PingAndWarmRequest request) {
     Metadata metadata = new Metadata();
 

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimerTest.java

@@ -18,6 +18,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import com.google.api.core.ApiFunction;
+import com.google.api.core.SettableApiFuture;
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.OAuth2Credentials;
 import com.google.bigtable.v2.BigtableGrpc.BigtableImplBase;
@@ -39,6 +40,8 @@
 import java.io.IOException;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import java.util.logging.Handler;
 import java.util.logging.LogRecord;
 import java.util.logging.Logger;
@@ -166,6 +169,38 @@ public void testHeadersAreSent() {
     }
   }
 
+  // New test for the async success path
+  @Test
+  public void testAsyncSuccess() throws Exception {
+    SettableApiFuture<PingAndWarmResponse> future = primer.sendPrimeRequestsAsync(channel);
+
+    PingAndWarmResponse response = future.get(1, TimeUnit.SECONDS);
+    assertThat(response).isNotNull();
+    assertThat(future.isDone()).isTrue();
+  }
+
+  // New test for the async failure path
+  @Test
+  public void testAsyncFailure() {
+    // Configure the server to return a gRPC error
+    fakeService.pingAndWarmCallback =
+        new ApiFunction<PingAndWarmRequest, PingAndWarmResponse>() {
+          @Override
+          public PingAndWarmResponse apply(PingAndWarmRequest pingAndWarmRequest) {
+            throw new StatusRuntimeException(Status.UNAVAILABLE);
+          }
+        };
+
+    SettableApiFuture<PingAndWarmResponse> future = primer.sendPrimeRequestsAsync(channel);
+
+    try {
+      future.get(1, TimeUnit.SECONDS);
+    } catch (Exception e) {
+      // Assert that the future completes with an ExecutionException
+      assertThat(e).isInstanceOf(ExecutionException.class);
+    }
+  }
+
   private static class MetadataInterceptor implements ServerInterceptor {
     ConcurrentLinkedQueue<Metadata> metadataList = new ConcurrentLinkedQueue<>();