From e53b502259534a45c8d5714d54f0e89cb60f2967 Mon Sep 17 00:00:00 2001
From: Jose Acevedo <sharp.acevedo@gmail.com>
Date: Wed, 22 May 2024 16:23:21 -0500
Subject: [PATCH] Added deadline to gRPC Java Client

Signed-off-by: Jose Acevedo <sharp.acevedo@gmail.com>
---
 .../src/main/java/dev/feast/FeastClient.java  | 42 ++++++++++++++++++-
 .../test/java/dev/feast/FeastClientTest.java  |  4 +-
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/java/serving-client/src/main/java/dev/feast/FeastClient.java b/java/serving-client/src/main/java/dev/feast/FeastClient.java
index c10a76ecf81..b2767e3c909 100644
--- a/java/serving-client/src/main/java/dev/feast/FeastClient.java
+++ b/java/serving-client/src/main/java/dev/feast/FeastClient.java
@@ -26,6 +26,7 @@
 import feast.proto.serving.ServingServiceGrpc.ServingServiceBlockingStub;
 import feast.proto.types.ValueProto;
 import io.grpc.CallCredentials;
+import io.grpc.Deadline;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
@@ -63,7 +64,20 @@ public static FeastClient create(String host, int port) {
   }
 
   /**
-   * Create a authenticated client that can access Feast serving with authentication enabled.
+   * Create a client to access Feast Serving.
+   *
+   * @param host hostname or ip address of Feast serving GRPC server
+   * @param port port number of Feast serving GRPC server
+   * @param deadline GRPC deadline of Feast serving GRPC server {@link Deadline}
+   * @return {@link FeastClient}
+   */
+  public static FeastClient create(String host, int port, Deadline deadline) {
+    // configure client with no security config.
+    return FeastClient.createSecure(host, port, SecurityConfig.newBuilder().build(), deadline);
+  }
+
+  /**
+   * Create an authenticated client that can access Feast serving with authentication enabled.
    *
    * @param host hostname or ip address of Feast serving GRPC server
    * @param port port number of Feast serving GRPC server
@@ -72,6 +86,21 @@ public static FeastClient create(String host, int port) {
    * @return {@link FeastClient}
    */
   public static FeastClient createSecure(String host, int port, SecurityConfig securityConfig) {
+    return createSecure(host, port, securityConfig, null);
+  }
+
+  /**
+   * Create an authenticated client that can access Feast serving with authentication enabled.
+   *
+   * @param host hostname or ip address of Feast serving GRPC server
+   * @param port port number of Feast serving GRPC server
+   * @param securityConfig security options to configure the Feast client. See {@link
+   *     SecurityConfig} for options.
+   * @param deadline GRPC deadline of Feast serving GRPC server {@link Deadline}
+   * @return {@link FeastClient}
+   */
+  public static FeastClient createSecure(
+      String host, int port, SecurityConfig securityConfig, Deadline deadline) {
     // Configure client TLS
     ManagedChannel channel = null;
     if (securityConfig.isTLSEnabled()) {
@@ -98,7 +127,7 @@ public static FeastClient createSecure(String host, int port, SecurityConfig sec
       channel = ManagedChannelBuilder.forAddress(host, port).usePlaintext().build();
     }
 
-    return new FeastClient(channel, securityConfig.getCredentials());
+    return new FeastClient(channel, securityConfig.getCredentials(), Optional.ofNullable(deadline));
   }
 
   /**
@@ -202,6 +231,11 @@ public List<Row> getOnlineFeatures(List<String> featureRefs, List<Row> rows, Str
   }
 
   protected FeastClient(ManagedChannel channel, Optional<CallCredentials> credentials) {
+    this(channel, credentials, Optional.empty());
+  }
+
+  protected FeastClient(
+      ManagedChannel channel, Optional<CallCredentials> credentials, Optional<Deadline> deadline) {
     this.channel = channel;
     TracingClientInterceptor tracingInterceptor =
         TracingClientInterceptor.newBuilder().withTracer(GlobalTracer.get()).build();
@@ -213,6 +247,10 @@ protected FeastClient(ManagedChannel channel, Optional<CallCredentials> credenti
       servingStub = servingStub.withCallCredentials(credentials.get());
     }
 
+    if (deadline.isPresent()) {
+      servingStub = servingStub.withDeadline(deadline.get());
+    }
+
     this.stub = servingStub;
   }
 
diff --git a/java/serving-client/src/test/java/dev/feast/FeastClientTest.java b/java/serving-client/src/test/java/dev/feast/FeastClientTest.java
index 1dfb9989c95..9846cff9f79 100644
--- a/java/serving-client/src/test/java/dev/feast/FeastClientTest.java
+++ b/java/serving-client/src/test/java/dev/feast/FeastClientTest.java
@@ -38,6 +38,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.Before;
 import org.junit.Rule;
@@ -45,6 +46,7 @@
 
 public class FeastClientTest {
   private final String AUTH_TOKEN = "test token";
+  private final Deadline DEADLINE = Deadline.after(2, TimeUnit.SECONDS);
 
   @Rule public GrpcCleanupRule grpcRule;
   private AtomicBoolean isAuthenticated;
@@ -86,7 +88,7 @@ public void setup() throws Exception {
     ManagedChannel channel =
         this.grpcRule.register(
             InProcessChannelBuilder.forName(serverName).directExecutor().build());
-    this.client = new FeastClient(channel, Optional.empty());
+    this.client = new FeastClient(channel, Optional.empty(), Optional.of(DEADLINE));
   }
 
   @Test