Merge pull request #104 from ExpediaGroup/support-grpc-deadlines-java…

…-client Add deadline to Java gRPC client
ExpediaGroup · May 15, 2024 · 3a6e0c9 · 3a6e0c9
2 parents 4996dca + 7e745ff
commit 3a6e0c9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 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
@@ -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;
@@ -57,9 +58,9 @@ public class FeastClient implements AutoCloseable {
    * @param port port number of Feast serving GRPC server
    * @return {@link FeastClient}
    */
-  public static FeastClient create(String host, int port) {
+  public static FeastClient create(String host, int port, Deadline deadline) {
     // configure client with no security config.
-    return FeastClient.createSecure(host, port, SecurityConfig.newBuilder().build());
+    return FeastClient.createSecure(host, port, SecurityConfig.newBuilder().build(), deadline);
   }
 
   /**
@@ -71,7 +72,7 @@ public static FeastClient create(String host, int port) {
    *     SecurityConfig} for options.
    * @return {@link FeastClient}
    */
-  public static FeastClient createSecure(String host, int port, SecurityConfig securityConfig) {
+  public static FeastClient createSecure(String host, int port, SecurityConfig securityConfig, Deadline deadline) {
     // Configure client TLS
     ManagedChannel channel = null;
     if (securityConfig.isTLSEnabled()) {
@@ -98,7 +99,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(), deadline);
   }
 
   /**
@@ -201,7 +202,7 @@ public List<Row> getOnlineFeatures(List<String> featureRefs, List<Row> rows, Str
     return getOnlineFeatures(featureRefs, rows);
   }
 
-  protected FeastClient(ManagedChannel channel, Optional<CallCredentials> credentials) {
+  protected FeastClient(ManagedChannel channel, Optional<CallCredentials> credentials, Deadline deadline) {
     this.channel = channel;
     TracingClientInterceptor tracingInterceptor =
         TracingClientInterceptor.newBuilder().withTracer(GlobalTracer.get()).build();
@@ -213,6 +214,8 @@ protected FeastClient(ManagedChannel channel, Optional<CallCredentials> credenti
       servingStub = servingStub.withCallCredentials(credentials.get());
     }
 
+    servingStub = servingStub.withDeadline(deadline);
+
     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
@@ -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;
@@ -86,7 +87,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(), Deadline.after(200, TimeUnit.MILLISECONDS));
   }
 
   @Test