refactor: use parallelizable scheduler for explore + entity (#93)

Author: aaron-steinfeld

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/dao/EntityDaoModule.java

@@ -15,10 +15,10 @@
 import org.hypertrace.core.graphql.common.utils.BiConverter;
 import org.hypertrace.core.graphql.common.utils.Converter;
 import org.hypertrace.core.graphql.common.utils.TriConverter;
-import org.hypertrace.core.graphql.context.GraphQlRequestContextBuilder;
 import org.hypertrace.core.graphql.rx.BoundedIoScheduler;
 import org.hypertrace.core.graphql.spi.config.GraphQlServiceConfig;
 import org.hypertrace.core.graphql.utils.grpc.GrpcChannelRegistry;
+import org.hypertrace.core.graphql.utils.grpc.GrpcContextBuilder;
 import org.hypertrace.gateway.service.v1.baseline.BaselineEntity;
 import org.hypertrace.gateway.service.v1.common.AggregatedMetricValue;
 import org.hypertrace.gateway.service.v1.common.Expression;
@@ -43,8 +43,8 @@ protected void configure() {
     bind(BaselineDao.class).to(GatewayBaselineDao.class);
     requireBinding(CallCredentials.class);
     requireBinding(GraphQlServiceConfig.class);
-    requireBinding(GraphQlRequestContextBuilder.class);
     requireBinding(GrpcChannelRegistry.class);
+    requireBinding(GrpcContextBuilder.class);
     requireBinding(Key.get(Scheduler.class, BoundedIoScheduler.class));
 
     requireBinding(

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/dao/GatewayBaselineDao.java

@@ -5,6 +5,7 @@
 import static org.hypertrace.core.graphql.common.utils.CollectorUtils.flatten;
 
 import io.grpc.CallCredentials;
+import io.reactivex.rxjava3.core.Scheduler;
 import io.reactivex.rxjava3.core.Single;
 import java.util.Collection;
 import java.util.Set;
@@ -13,9 +14,10 @@
 import javax.inject.Singleton;
 import org.hypertrace.core.graphql.common.utils.Converter;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.hypertrace.core.graphql.rx.BoundedIoScheduler;
 import org.hypertrace.core.graphql.spi.config.GraphQlServiceConfig;
-import org.hypertrace.core.graphql.utils.grpc.GraphQlGrpcContextBuilder;
 import org.hypertrace.core.graphql.utils.grpc.GrpcChannelRegistry;
+import org.hypertrace.core.graphql.utils.grpc.GrpcContextBuilder;
 import org.hypertrace.gateway.service.GatewayServiceGrpc;
 import org.hypertrace.gateway.service.v1.baseline.BaselineEntitiesRequest;
 import org.hypertrace.gateway.service.v1.baseline.BaselineEntitiesResponse;
@@ -34,9 +36,9 @@
 
 @Singleton
 class GatewayBaselineDao implements BaselineDao {
-  private static final int DEFAULT_DEADLINE_SEC = 10;
   private final GatewayServiceGrpc.GatewayServiceFutureStub gatewayServiceStub;
-  private final GraphQlGrpcContextBuilder grpcContextBuilder;
+  private final GrpcContextBuilder grpcContextBuilder;
+  private final Scheduler boundedIoScheduler;
   private final Converter<Collection<MetricAggregationRequest>, Set<Expression>>
       aggregationConverter;
   private final Converter<Collection<MetricSeriesRequest>, Set<TimeAggregation>> seriesConverter;
@@ -47,10 +49,12 @@ class GatewayBaselineDao implements BaselineDao {
       GrpcChannelRegistry grpcChannelRegistry,
       GraphQlServiceConfig serviceConfig,
       CallCredentials credentials,
-      GraphQlGrpcContextBuilder grpcContextBuilder,
+      GrpcContextBuilder grpcContextBuilder,
+      @BoundedIoScheduler Scheduler boundedIoScheduler,
       Converter<Collection<MetricAggregationRequest>, Set<Expression>> aggregationConverter,
       Converter<Collection<MetricSeriesRequest>, Set<TimeAggregation>> seriesConverter) {
     this.grpcContextBuilder = grpcContextBuilder;
+    this.boundedIoScheduler = boundedIoScheduler;
     this.gatewayServiceStub =
         GatewayServiceGrpc.newFutureStub(
                 grpcChannelRegistry.forAddress(
@@ -68,6 +72,7 @@ public Single<BaselineEntitiesResponse> getBaselines(
       EntitiesResponse entitiesResponse,
       EntityRequest request) {
     return this.buildRequest(entitiesRequest, entitiesResponse, request)
+        .subscribeOn(this.boundedIoScheduler)
         .flatMap(baselineEntitiesRequest -> makeRequest(context, baselineEntitiesRequest));
   }
 
@@ -126,7 +131,7 @@ private Single<BaselineEntitiesResponse> makeRequest(
     return Single.fromFuture(
         this.grpcContextBuilder
             .build(context)
-            .callInContext(
+            .call(
                 () ->
                     this.gatewayServiceStub
                         .withDeadlineAfter(

hypertrace-graphql-entity-schema/src/main/java/org/hypertrace/graphql/entity/dao/GatewayServiceEntityDao.java

@@ -3,13 +3,15 @@
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 import io.grpc.CallCredentials;
+import io.reactivex.rxjava3.core.Scheduler;
 import io.reactivex.rxjava3.core.Single;
 import javax.inject.Inject;
 import javax.inject.Singleton;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.hypertrace.core.graphql.rx.BoundedIoScheduler;
 import org.hypertrace.core.graphql.spi.config.GraphQlServiceConfig;
-import org.hypertrace.core.graphql.utils.grpc.GraphQlGrpcContextBuilder;
 import org.hypertrace.core.graphql.utils.grpc.GrpcChannelRegistry;
+import org.hypertrace.core.graphql.utils.grpc.GrpcContextBuilder;
 import org.hypertrace.gateway.service.GatewayServiceGrpc;
 import org.hypertrace.gateway.service.GatewayServiceGrpc.GatewayServiceFutureStub;
 import org.hypertrace.gateway.service.v1.entity.EntitiesRequest;
@@ -21,26 +23,29 @@
 @Singleton
 class GatewayServiceEntityDao implements EntityDao {
   private final GatewayServiceFutureStub gatewayServiceStub;
-  private final GraphQlGrpcContextBuilder grpcContextBuilder;
+  private final GrpcContextBuilder grpcContextBuilder;
   private final GatewayServiceEntityRequestBuilder requestBuilder;
   private final GatewayServiceEntityConverter entityConverter;
   private final BaselineDao baselineDao;
   private final GraphQlServiceConfig serviceConfig;
+  private final Scheduler boundedIoScheduler;
 
   @Inject
   GatewayServiceEntityDao(
       GraphQlServiceConfig serviceConfig,
       CallCredentials credentials,
-      GraphQlGrpcContextBuilder grpcContextBuilder,
+      GrpcContextBuilder grpcContextBuilder,
       GrpcChannelRegistry grpcChannelRegistry,
       GatewayServiceEntityRequestBuilder requestBuilder,
       GatewayServiceEntityConverter entityConverter,
-      BaselineDao baselineDao) {
+      BaselineDao baselineDao,
+      @BoundedIoScheduler Scheduler boundedIoScheduler) {
     this.grpcContextBuilder = grpcContextBuilder;
     this.requestBuilder = requestBuilder;
     this.entityConverter = entityConverter;
     this.baselineDao = baselineDao;
     this.serviceConfig = serviceConfig;
+    this.boundedIoScheduler = boundedIoScheduler;
 
     this.gatewayServiceStub =
         GatewayServiceGrpc.newFutureStub(
@@ -53,6 +58,7 @@ class GatewayServiceEntityDao implements EntityDao {
   public Single<EntityResultSet> getEntities(EntityRequest request) {
     return this.requestBuilder
         .buildRequest(request)
+        .subscribeOn(this.boundedIoScheduler)
         .flatMap(
             serverRequest ->
                 this.makeRequest(request.resultSetRequest().context(), serverRequest)
@@ -75,7 +81,7 @@ private Single<EntitiesResponse> makeRequest(
     return Single.fromFuture(
         this.grpcContextBuilder
             .build(context)
-            .callInContext(
+            .call(
                 () ->
                     this.gatewayServiceStub
                         .withDeadlineAfter(

hypertrace-graphql-explorer-schema/build.gradle.kts

@@ -25,6 +25,7 @@ dependencies {
   implementation("org.hypertrace.core.graphql:hypertrace-core-graphql-attribute-store")
   implementation("org.hypertrace.core.graphql:hypertrace-core-graphql-deserialization")
   implementation("org.hypertrace.core.graphql:hypertrace-core-graphql-schema-utils")
+  implementation("org.hypertrace.core.graphql:hypertrace-core-graphql-rx-utils")
 
   testImplementation("org.junit.jupiter:junit-jupiter")
   testImplementation("org.mockito:mockito-core")

hypertrace-graphql-explorer-schema/src/main/java/org/hypertrace/graphql/explorer/dao/ExplorerDaoModule.java

@@ -4,6 +4,7 @@
 import com.google.inject.Key;
 import com.google.inject.TypeLiteral;
 import io.grpc.CallCredentials;
+import io.reactivex.rxjava3.core.Scheduler;
 import java.util.Collection;
 import java.util.List;
 import java.util.Set;
@@ -15,9 +16,10 @@
 import org.hypertrace.core.graphql.common.schema.attributes.MetricAggregationType;
 import org.hypertrace.core.graphql.common.schema.results.arguments.filter.FilterArgument;
 import org.hypertrace.core.graphql.common.utils.Converter;
-import org.hypertrace.core.graphql.context.GraphQlRequestContextBuilder;
+import org.hypertrace.core.graphql.rx.BoundedIoScheduler;
 import org.hypertrace.core.graphql.spi.config.GraphQlServiceConfig;
 import org.hypertrace.core.graphql.utils.grpc.GrpcChannelRegistry;
+import org.hypertrace.core.graphql.utils.grpc.GrpcContextBuilder;
 import org.hypertrace.gateway.service.v1.common.Expression;
 import org.hypertrace.gateway.service.v1.common.Filter;
 import org.hypertrace.gateway.service.v1.common.OrderByExpression;
@@ -34,8 +36,9 @@ protected void configure() {
     bind(ExplorerDao.class).to(GatewayServiceExplorerDao.class);
     requireBinding(CallCredentials.class);
     requireBinding(GraphQlServiceConfig.class);
-    requireBinding(GraphQlRequestContextBuilder.class);
     requireBinding(GrpcChannelRegistry.class);
+    requireBinding(GrpcContextBuilder.class);
+    requireBinding(Key.get(Scheduler.class, BoundedIoScheduler.class));
     requireBinding(Key.get(new TypeLiteral<Converter<Value, Object>>() {}));
     requireBinding(Key.get(new TypeLiteral<Converter<AttributeModelType, AttributeType>>() {}));
     requireBinding(

hypertrace-graphql-explorer-schema/src/main/java/org/hypertrace/graphql/explorer/dao/GatewayServiceExplorerDao.java

@@ -3,13 +3,15 @@
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 import io.grpc.CallCredentials;
+import io.reactivex.rxjava3.core.Scheduler;
 import io.reactivex.rxjava3.core.Single;
 import javax.inject.Inject;
 import javax.inject.Singleton;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.hypertrace.core.graphql.rx.BoundedIoScheduler;
 import org.hypertrace.core.graphql.spi.config.GraphQlServiceConfig;
-import org.hypertrace.core.graphql.utils.grpc.GraphQlGrpcContextBuilder;
 import org.hypertrace.core.graphql.utils.grpc.GrpcChannelRegistry;
+import org.hypertrace.core.graphql.utils.grpc.GrpcContextBuilder;
 import org.hypertrace.gateway.service.GatewayServiceGrpc;
 import org.hypertrace.gateway.service.GatewayServiceGrpc.GatewayServiceFutureStub;
 import org.hypertrace.gateway.service.v1.explore.ExploreResponse;
@@ -19,23 +21,26 @@
 @Singleton
 class GatewayServiceExplorerDao implements ExplorerDao {
   private final GatewayServiceFutureStub gatewayServiceStub;
-  private final GraphQlGrpcContextBuilder grpcContextBuilder;
+  private final GrpcContextBuilder grpcContextBuilder;
   private final GatewayServiceExploreRequestBuilder requestBuilder;
   private final GatewayServiceExploreResponseConverter responseConverter;
   private final GraphQlServiceConfig serviceConfig;
+  private final Scheduler boundedIoScheduler;
 
   @Inject
   GatewayServiceExplorerDao(
       GraphQlServiceConfig serviceConfig,
       CallCredentials credentials,
-      GraphQlGrpcContextBuilder grpcContextBuilder,
+      GrpcContextBuilder grpcContextBuilder,
       GrpcChannelRegistry grpcChannelRegistry,
       GatewayServiceExploreRequestBuilder requestBuilder,
-      GatewayServiceExploreResponseConverter responseConverter) {
+      GatewayServiceExploreResponseConverter responseConverter,
+      @BoundedIoScheduler Scheduler boundedIoScheduler) {
     this.grpcContextBuilder = grpcContextBuilder;
     this.requestBuilder = requestBuilder;
     this.responseConverter = responseConverter;
     this.serviceConfig = serviceConfig;
+    this.boundedIoScheduler = boundedIoScheduler;
 
     this.gatewayServiceStub =
         GatewayServiceGrpc.newFutureStub(
@@ -48,6 +53,7 @@ class GatewayServiceExplorerDao implements ExplorerDao {
   public Single<ExploreResultSet> explore(ExploreRequest request) {
     return this.requestBuilder
         .buildRequest(request)
+        .subscribeOn(this.boundedIoScheduler)
         .flatMap(serverRequest -> this.makeRequest(request.requestContext(), serverRequest))
         .flatMap(serverResponse -> this.responseConverter.convert(request, serverResponse));
   }
@@ -58,7 +64,7 @@ private Single<ExploreResponse> makeRequest(
     return Single.fromFuture(
         this.grpcContextBuilder
             .build(context)
-            .callInContext(
+            .call(
                 () ->
                     this.gatewayServiceStub
                         .withDeadlineAfter(