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..b38a738e515 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; @@ -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 getOnlineFeatures(List featureRefs, List rows, Str return getOnlineFeatures(featureRefs, rows); } - protected FeastClient(ManagedChannel channel, Optional credentials) { + protected FeastClient(ManagedChannel channel, Optional credentials, Deadline deadline) { this.channel = channel; TracingClientInterceptor tracingInterceptor = TracingClientInterceptor.newBuilder().withTracer(GlobalTracer.get()).build(); @@ -213,6 +214,8 @@ protected FeastClient(ManagedChannel channel, Optional 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 index 1dfb9989c95..dfcd63ac49e 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; @@ -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