diff --git a/hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/GrpcContextBuilder.java b/hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/GrpcContextBuilder.java index 48c7b332..26373636 100644 --- a/hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/GrpcContextBuilder.java +++ b/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 tryRestore(RequestContext requestContext); } diff --git a/hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/PlatformGrpcContextBuilder.java b/hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/PlatformGrpcContextBuilder.java index 7c090220..880ec704 100644 --- a/hypertrace-core-graphql-grpc-utils/src/main/java/org/hypertrace/core/graphql/utils/grpc/PlatformGrpcContextBuilder.java +++ b/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 contextCache = + CacheBuilder.newBuilder().expireAfterAccess(Duration.ofMinutes(1)).maximumSize(1000).build(); + @Override public RequestContext build(GraphQlRequestContext requestContext) { + this.contextCache.put(requestContext.getRequestId(), requestContext); Map grpcHeaders = this.mergeMaps( requestContext.getTracingContextHeaders(), @@ -30,6 +37,11 @@ public RequestContext build(GraphQlRequestContext requestContext) { return this.build(grpcHeaders); } + @Override + public Optional tryRestore(RequestContext requestContext) { + return requestContext.getRequestId().map(this.contextCache::getIfPresent); + } + private RequestContext build(@Nonnull Map headers) { RequestContext platformContext = new RequestContext(); headers.forEach(platformContext::add); diff --git a/hypertrace-core-graphql-grpc-utils/src/test/java/org/hypertrace/core/graphql/utils/grpc/PlatformGrpcContextBuilderTest.java b/hypertrace-core-graphql-grpc-utils/src/test/java/org/hypertrace/core/graphql/utils/grpc/PlatformGrpcContextBuilderTest.java index 7b0efaf1..20ef7a77 100644 --- a/hypertrace-core-graphql-grpc-utils/src/test/java/org/hypertrace/core/graphql/utils/grpc/PlatformGrpcContextBuilderTest.java +++ b/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)); + } }