refactor: add ability to restore context

Author: aaron-steinfeld

hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/GrpcContextBuilder.java

@@ -1,8 +1,11 @@
 package org.hypertrace.core.graphql.utils.grpc;
 
+import java.util.Optional;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
 import org.hypertrace.core.grpcutils.context.RequestContext;
 
 public interface GrpcContextBuilder {
   RequestContext build(GraphQlRequestContext requestContext);
+
+  Optional<GraphQlRequestContext> tryRestore(RequestContext requestContext);
 }

hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/PlatformGrpcContextBuilder.java

@@ -4,6 +4,9 @@
 import static org.hypertrace.core.grpcutils.context.RequestContextConstants.REQUEST_ID_HEADER_KEY;
 import static org.hypertrace.core.grpcutils.context.RequestContextConstants.TENANT_ID_HEADER_KEY;
 
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Map;
@@ -16,8 +19,12 @@
 
 class PlatformGrpcContextBuilder implements GrpcContextBuilder {
 
+  private final Cache<String, GraphQlRequestContext> contextCache =
+      CacheBuilder.newBuilder().expireAfterAccess(Duration.ofMinutes(1)).maximumSize(1000).build();
+
   @Override
   public RequestContext build(GraphQlRequestContext requestContext) {
+    this.contextCache.put(requestContext.getRequestId(), requestContext);
     Map<String, String> grpcHeaders =
         this.mergeMaps(
             requestContext.getTracingContextHeaders(),
@@ -30,6 +37,11 @@ public RequestContext build(GraphQlRequestContext requestContext) {
     return this.build(grpcHeaders);
   }
 
+  @Override
+  public Optional<GraphQlRequestContext> tryRestore(RequestContext requestContext) {
+    return requestContext.getRequestId().map(this.contextCache::getIfPresent);
+  }
+
   private RequestContext build(@Nonnull Map<String, String> headers) {
     RequestContext platformContext = new RequestContext();
     headers.forEach(platformContext::add);

hypertrace-core-graphql-grpc-utils/src/test/java/org/hypertrace/core/graphql/utils/grpc/PlatformGrpcContextBuilderTest.java

@@ -9,6 +9,7 @@
 import java.util.Map;
 import java.util.Optional;
 import org.hypertrace.core.graphql.context.GraphQlRequestContext;
+import org.hypertrace.core.grpcutils.context.RequestContext;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
@@ -61,4 +62,12 @@ void passesAuthHeaderToPlatformContextIfPresent() {
         "auth header",
         this.builder.build(this.mockRequestContext).getRequestHeaders().get("authorization"));
   }
+
+  @Test
+  void testRestoreContext() {
+    when(this.mockRequestContext.getRequestId()).thenReturn("request id");
+    RequestContext resultContext = this.builder.build(this.mockRequestContext);
+    assertEquals(Optional.empty(), this.builder.tryRestore(RequestContext.forTenantId("other")));
+    assertEquals(Optional.of(this.mockRequestContext), this.builder.tryRestore(resultContext));
+  }
 }