refactor: add support for client timeouts

object-store/src/main/java/org/hypertrace/config/objectstore/ClientConfig.java

@@ -0,0 +1,11 @@
+package org.hypertrace.config.objectstore;
+
+import java.time.Duration;
+import lombok.Value;
+
+@Value
+public class ClientConfig {
+  Duration timeout;
+  // TODO: Explore unobtrusive ways of setting this default from config
+  public static ClientConfig DEFAULT = new ClientConfig(Duration.ofSeconds(10));
+}

object-store/src/main/java/org/hypertrace/config/objectstore/ContextuallyIdentifiedObjectStore.java

@@ -1,41 +1,50 @@
 package org.hypertrace.config.objectstore;
 
 import com.google.protobuf.Value;
-import io.grpc.Status;
 import java.util.Optional;
 import java.util.function.Function;
 import org.hypertrace.config.service.change.event.api.ConfigChangeEventGenerator;
 import org.hypertrace.config.service.v1.ConfigServiceGrpc.ConfigServiceBlockingStub;
-import org.hypertrace.config.service.v1.ContextSpecificConfig;
-import org.hypertrace.config.service.v1.GetConfigRequest;
-import org.hypertrace.config.service.v1.GetConfigResponse;
 import org.hypertrace.core.grpcutils.context.RequestContext;
 
 public abstract class ContextuallyIdentifiedObjectStore<T> {
   private final ConfigServiceBlockingStub configServiceBlockingStub;
   private final String resourceNamespace;
   private final String resourceName;
   private final Optional<ConfigChangeEventGenerator> configChangeEventGenerator;
+  private final ClientConfig clientConfig;
 
   protected ContextuallyIdentifiedObjectStore(
       ConfigServiceBlockingStub configServiceBlockingStub,
       String resourceNamespace,
       String resourceName,
       ConfigChangeEventGenerator configChangeEventGenerator) {
+    this(
+        configServiceBlockingStub,
+        resourceNamespace,
+        resourceName,
+        configChangeEventGenerator,
+        ClientConfig.DEFAULT);
+  }
+
+  protected ContextuallyIdentifiedObjectStore(
+      ConfigServiceBlockingStub configServiceBlockingStub,
+      String resourceNamespace,
+      String resourceName,
+      ConfigChangeEventGenerator configChangeEventGenerator,
+      ClientConfig clientConfig) {
     this.configServiceBlockingStub = configServiceBlockingStub;
     this.resourceNamespace = resourceNamespace;
     this.resourceName = resourceName;
-    this.configChangeEventGenerator = Optional.of(configChangeEventGenerator);
+    this.configChangeEventGenerator = Optional.ofNullable(configChangeEventGenerator);
+    this.clientConfig = clientConfig;
   }
 
   protected ContextuallyIdentifiedObjectStore(
       ConfigServiceBlockingStub configServiceBlockingStub,
       String resourceNamespace,
       String resourceName) {
-    this.configServiceBlockingStub = configServiceBlockingStub;
-    this.resourceNamespace = resourceNamespace;
-    this.resourceName = resourceName;
-    this.configChangeEventGenerator = Optional.empty();
+    this(configServiceBlockingStub, resourceNamespace, resourceName, null);
   }
 
   protected abstract Optional<T> buildDataFromValue(Value value);
@@ -59,32 +68,9 @@ private IdentifiedObjectStore<T> buildObjectStoreForContext(RequestContext conte
   }
 
   public Optional<ConfigObject<T>> getObject(RequestContext context) {
-    String configContext = this.getConfigContextFromRequestContext(context);
-    try {
-      GetConfigResponse getConfigResponse =
-          context.call(
-              () ->
-                  this.configServiceBlockingStub.getConfig(
-                      GetConfigRequest.newBuilder()
-                          .setResourceName(this.resourceName)
-                          .setResourceNamespace(this.resourceNamespace)
-                          .addContexts(configContext)
-                          .build()));
-
-      ContextSpecificConfig contextSpecificConfig =
-          ContextSpecificConfig.newBuilder()
-              .setContext(configContext)
-              .setConfig(getConfigResponse.getConfig())
-              .setCreationTimestamp(getConfigResponse.getCreationTimestamp())
-              .setUpdateTimestamp(getConfigResponse.getUpdateTimestamp())
-              .build();
-      return ConfigObjectImpl.tryBuild(contextSpecificConfig, this::buildDataFromValue);
-    } catch (Exception exception) {
-      if (Status.fromThrowable(exception).equals(Status.NOT_FOUND)) {
-        return Optional.empty();
-      }
-      throw exception;
-    }
+    return this.buildObjectStoreForContext(context)
+        .getObject(context, this.getConfigContextFromRequestContext(context))
+        .map(Function.identity());
   }
 
   public Optional<T> getData(RequestContext context) {
@@ -107,13 +93,17 @@ private class ContextAwareIdentifiedObjectStoreDelegate extends IdentifiedObject
     private final RequestContext requestContext;
 
     ContextAwareIdentifiedObjectStoreDelegate(RequestContext requestContext) {
-      super(configServiceBlockingStub, resourceNamespace, resourceName);
-      this.requestContext = requestContext;
+      this(requestContext, null);
     }
 
     ContextAwareIdentifiedObjectStoreDelegate(
         RequestContext requestContext, ConfigChangeEventGenerator configChangeEventGenerator) {
-      super(configServiceBlockingStub, resourceNamespace, resourceName, configChangeEventGenerator);
+      super(
+          configServiceBlockingStub,
+          resourceNamespace,
+          resourceName,
+          configChangeEventGenerator,
+          clientConfig);
       this.requestContext = requestContext;
     }
 

object-store/src/main/java/org/hypertrace/config/objectstore/DefaultObjectStore.java

@@ -1,8 +1,10 @@
 package org.hypertrace.config.objectstore;
 
 import com.google.protobuf.Value;
+import io.grpc.Deadline;
 import io.grpc.Status;
 import java.util.Optional;
+import java.util.concurrent.TimeUnit;
 import lombok.extern.slf4j.Slf4j;
 import org.hypertrace.config.service.change.event.api.ConfigChangeEventGenerator;
 import org.hypertrace.config.service.v1.ConfigServiceGrpc.ConfigServiceBlockingStub;
@@ -26,26 +28,39 @@
   private final String resourceNamespace;
   private final String resourceName;
   private final Optional<ConfigChangeEventGenerator> configChangeEventGeneratorOptional;
+  private final ClientConfig clientConfig;
 
   protected DefaultObjectStore(
       ConfigServiceBlockingStub configServiceBlockingStub,
       String resourceNamespace,
       String resourceName,
-      ConfigChangeEventGenerator configChangeEventGenerator) {
+      ConfigChangeEventGenerator configChangeEventGenerator,
+      ClientConfig clientConfig) {
     this.configServiceBlockingStub = configServiceBlockingStub;
     this.resourceNamespace = resourceNamespace;
     this.resourceName = resourceName;
-    this.configChangeEventGeneratorOptional = Optional.of(configChangeEventGenerator);
+    this.configChangeEventGeneratorOptional = Optional.ofNullable(configChangeEventGenerator);
+    this.clientConfig = clientConfig;
+  }
+
+  protected DefaultObjectStore(
+      ConfigServiceBlockingStub configServiceBlockingStub,
+      String resourceNamespace,
+      String resourceName,
+      ConfigChangeEventGenerator configChangeEventGenerator) {
+    this(
+        configServiceBlockingStub,
+        resourceNamespace,
+        resourceName,
+        configChangeEventGenerator,
+        ClientConfig.DEFAULT);
   }
 
   protected DefaultObjectStore(
       ConfigServiceBlockingStub configServiceBlockingStub,
       String resourceNamespace,
       String resourceName) {
-    this.configServiceBlockingStub = configServiceBlockingStub;
-    this.resourceNamespace = resourceNamespace;
-    this.resourceName = resourceName;
-    this.configChangeEventGeneratorOptional = Optional.empty();
+    this(configServiceBlockingStub, resourceNamespace, resourceName, null);
   }
 
   protected abstract Optional<T> buildDataFromValue(Value value);
@@ -65,11 +80,13 @@
       GetConfigResponse getConfigResponse =
           context.call(
               () ->
-                  this.configServiceBlockingStub.getConfig(
-                      GetConfigRequest.newBuilder()
-                          .setResourceName(this.resourceName)
-                          .setResourceNamespace(this.resourceNamespace)
-                          .build()));
+                  this.configServiceBlockingStub
+                      .withDeadline(getDeadline())
+                      .getConfig(
+                          GetConfigRequest.newBuilder()
+                              .setResourceName(this.resourceName)
+                              .setResourceNamespace(this.resourceNamespace)
+                              .build()));
 
       return ConfigObjectImpl.tryBuild(
           getConfigResponse.getConfig(),
@@ -90,6 +107,7 @@
           context.call(
               () ->
                   this.configServiceBlockingStub
+                      .withDeadline(getDeadline())
                       .getConfig(
                           GetConfigRequest.newBuilder()
                               .setResourceName(this.resourceName)
@@ -110,12 +128,14 @@
     UpsertConfigResponse response =
         context.call(
             () ->
-                this.configServiceBlockingStub.upsertConfig(
-                    UpsertConfigRequest.newBuilder()
-                        .setResourceName(this.resourceName)
-                        .setResourceNamespace(this.resourceNamespace)
-                        .setConfig(this.buildValueFromData(data))
-                        .build()));
+                this.configServiceBlockingStub
+                    .withDeadline(getDeadline())
+                    .upsertConfig(
+                        UpsertConfigRequest.newBuilder()
+                            .setResourceName(this.resourceName)
+                            .setResourceNamespace(this.resourceNamespace)
+                            .setConfig(this.buildValueFromData(data))
+                            .build()));
 
     ConfigObject<T> upsertedObject =
         ConfigObjectImpl.tryBuild(response, this::buildDataFromValue)
@@ -153,11 +173,13 @@
           context
               .call(
                   () ->
-                      this.configServiceBlockingStub.deleteConfig(
-                          DeleteConfigRequest.newBuilder()
-                              .setResourceName(this.resourceName)
-                              .setResourceNamespace(this.resourceNamespace)
-                              .build()))
+                      this.configServiceBlockingStub
+                          .withDeadline(getDeadline())
+                          .deleteConfig(
+                              DeleteConfigRequest.newBuilder()
+                                  .setResourceName(this.resourceName)
+                                  .setResourceNamespace(this.resourceNamespace)
+                                  .build()))
               .getDeletedConfig();
       ConfigObject<T> object =
           ConfigObjectImpl.tryBuild(deletedConfig, this::buildDataFromValue)
@@ -176,4 +198,8 @@
       throw exception;
     }
   }
+
+  protected Deadline getDeadline() {
+    return Deadline.after(this.clientConfig.getTimeout().toMillis(), TimeUnit.MILLISECONDS);
+  }
 }