Skip to content

Commit eb6b0d1

Browse files
aaron-steinfeldaaron-steinfeld
authored
feat: add context-aware cache key (#16)
* feat: add context-aware cache key * test: add tests * refactor: add no data run methods to key
1 parent 614ba57 commit eb6b0d1

File tree

5 files changed

+222
-8
lines changed

5 files changed

+222
-8
lines changed

‎grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/ContextualKey.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package org.hypertrace.core.grpcutils.context;
2+
3+
import java.util.function.Consumer;
4+
import java.util.function.Function;
5+
import java.util.function.Supplier;
6+
7+
public interface ContextualKey<T> {
8+
RequestContext getContext();
9+
10+
T getData();
11+
12+
/**
13+
* Calls the function in the key's context and providing the key's data as an argument, returning
14+
* any result
15+
*/
16+
<R> R callInContext(Function<T, R> function);
17+
18+
<R> R callInContext(Supplier<R> supplier);
19+
20+
/**
21+
* Calls the function in the key's context and providing the key's data as an argument, returning
22+
* no result
23+
*/
24+
void runInContext(Consumer<T> consumer);
25+
26+
void runInContext(Runnable runnable);
27+
}

‎grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/DefaultContextualKey.java‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
package org.hypertrace.core.grpcutils.context;
2+
3+
import java.util.Map;
4+
import java.util.Map.Entry;
5+
import java.util.Objects;
6+
import java.util.function.Consumer;
7+
import java.util.function.Function;
8+
import java.util.function.Supplier;
9+
import java.util.stream.Collectors;
10+
11+
class DefaultContextualKey<T> implements ContextualKey<T> {
12+
private final RequestContext context;
13+
private final T data;
14+
private final Map<String, String> meaningfulContextHeaders;
15+
16+
DefaultContextualKey(RequestContext context, T data) {
17+
this.context = context;
18+
this.data = data;
19+
this.meaningfulContextHeaders = this.extractMeaningfulHeaders(context.getRequestHeaders());
20+
}
21+
22+
@Override
23+
public RequestContext getContext() {
24+
return this.context;
25+
}
26+
27+
@Override
28+
public T getData() {
29+
return this.data;
30+
}
31+
32+
@Override
33+
public <R> R callInContext(Function<T, R> function) {
34+
return this.context.call(() -> function.apply(this.getData()));
35+
}
36+
37+
@Override
38+
public <R> R callInContext(Supplier<R> supplier) {
39+
return this.context.call(supplier::get);
40+
}
41+
42+
@Override
43+
public void runInContext(Consumer<T> consumer) {
44+
this.context.run(() -> consumer.accept(this.getData()));
45+
}
46+
47+
@Override
48+
public void runInContext(Runnable runnable) {
49+
this.context.run(runnable);
50+
}
51+
52+
@Override
53+
public boolean equals(Object o) {
54+
if (this == o) return true;
55+
if (o == null || getClass() != o.getClass()) return false;
56+
DefaultContextualKey<?> that = (DefaultContextualKey<?>) o;
57+
return Objects.equals(getData(), that.getData())
58+
&& meaningfulContextHeaders.equals(that.meaningfulContextHeaders);
59+
}
60+
61+
@Override
62+
public int hashCode() {
63+
return Objects.hash(getData(), meaningfulContextHeaders);
64+
}
65+
66+
@Override
67+
public String toString() {
68+
return "DefaultContextualKey{"
69+
+ "data="
70+
+ data
71+
+ ", meaningfulContextHeaders="
72+
+ meaningfulContextHeaders
73+
+ '}';
74+
}
75+
76+
private Map<String, String> extractMeaningfulHeaders(Map<String, String> allHeaders) {
77+
return allHeaders.entrySet().stream()
78+
.filter(
79+
entry ->
80+
RequestContextConstants.CACHE_MEANINGFUL_HEADERS.contains(
81+
entry.getKey().toLowerCase()))
82+
.collect(Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue));
83+
}
84+
}

‎grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContext.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,4 +79,12 @@ public <V> V call(@Nonnull Callable<V> callable) {
7979
public void run(@Nonnull Runnable runnable) {
8080
Context.current().withValue(RequestContext.CURRENT, this).run(runnable);
8181
}
82+
83+
public <T> ContextualKey<T> buildContextualKey(T data) {
84+
return new DefaultContextualKey<>(this, data);
85+
}
86+
87+
public ContextualKey<Void> buildContextualKey() {
88+
return new DefaultContextualKey<>(this, null);
89+
}
8290
}

‎grpc-context-utils/src/main/java/org/hypertrace/core/grpcutils/context/RequestContextConstants.java‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
package org.hypertrace.core.grpcutils.context;
22

3+
import static io.grpc.Metadata.ASCII_STRING_MARSHALLER;
4+
35
import io.grpc.Metadata;
46
import java.util.Set;
57

6-
import static io.grpc.Metadata.ASCII_STRING_MARSHALLER;
7-
88
/**
9-
* GRPC request context constants used to propagate the tenantId, authorization token, tracing headers etc
10-
* in the platform services.
9+
* GRPC request context constants used to propagate the tenantId, authorization token, tracing
10+
* headers etc in the platform services.
1111
*/
1212
public class RequestContextConstants {
1313
public static final String TENANT_ID_HEADER_KEY = "x-tenant-id";
@@ -17,10 +17,20 @@ public class RequestContextConstants {
1717

1818
public static final String AUTHORIZATION_HEADER = "authorization";
1919

20+
/** The values in this set are looked up with case insensitivity. */
21+
public static final Set<String> HEADER_PREFIXES_TO_BE_PROPAGATED =
22+
Set.of(
23+
TENANT_ID_HEADER_KEY,
24+
"X-B3-",
25+
"grpc-trace-bin",
26+
"traceparent",
27+
"tracestate",
28+
AUTHORIZATION_HEADER);
29+
2030
/**
21-
* The values in this set are looked up with case insensitivity.
31+
* These headers may affect returned results and should be accounted for in any cached remote
32+
* results
2233
*/
23-
public static final Set<String> HEADER_PREFIXES_TO_BE_PROPAGATED =
24-
Set.of(TENANT_ID_HEADER_KEY, "X-B3-", "grpc-trace-bin",
25-
"traceparent", "tracestate", AUTHORIZATION_HEADER);
34+
static final Set<String> CACHE_MEANINGFUL_HEADERS =
35+
Set.of(TENANT_ID_HEADER_KEY, AUTHORIZATION_HEADER);
2636
}

‎grpc-context-utils/src/test/java/org/hypertrace/core/grpcutils/context/DefaultContextualKeyTest.java‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
package org.hypertrace.core.grpcutils.context;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNotEquals;
5+
import static org.junit.jupiter.api.Assertions.assertSame;
6+
import static org.mockito.ArgumentMatchers.any;
7+
import static org.mockito.ArgumentMatchers.eq;
8+
import static org.mockito.Mockito.doAnswer;
9+
import static org.mockito.Mockito.mock;
10+
import static org.mockito.Mockito.times;
11+
import static org.mockito.Mockito.verify;
12+
13+
import java.util.function.Consumer;
14+
import java.util.function.Function;
15+
import java.util.function.Supplier;
16+
import org.junit.jupiter.api.Test;
17+
18+
class DefaultContextualKeyTest {
19+
20+
@Test
21+
void callsProvidedMethodsInContext() {
22+
RequestContext testContext = RequestContext.forTenantId("test-tenant");
23+
ContextualKey<String> key = new DefaultContextualKey<>(testContext, "input");
24+
25+
Function<String, String> testFunction =
26+
value ->
27+
"returned: "
28+
+ value
29+
+ " for "
30+
+ RequestContext.CURRENT.get().getTenantId().orElseThrow();
31+
32+
assertEquals("returned: input for test-tenant", key.callInContext(testFunction));
33+
34+
Supplier<String> testSupplier =
35+
() -> "returned for " + RequestContext.CURRENT.get().getTenantId().orElseThrow();
36+
37+
assertEquals("returned for test-tenant", key.callInContext(testSupplier));
38+
}
39+
40+
@Test
41+
void runsProvidedMethodInContext() {
42+
RequestContext testContext = RequestContext.forTenantId("test-tenant");
43+
ContextualKey<String> key = new DefaultContextualKey<>(testContext, "input");
44+
45+
Consumer<String> testConsumer = mock(Consumer.class);
46+
47+
doAnswer(
48+
invocation -> {
49+
assertSame(testContext, RequestContext.CURRENT.get());
50+
return null;
51+
})
52+
.when(testConsumer)
53+
.accept(any());
54+
key.runInContext(testConsumer);
55+
verify(testConsumer, times(1)).accept(eq("input"));
56+
57+
Runnable testRunnable = mock(Runnable.class);
58+
key.runInContext(testRunnable);
59+
verify(testRunnable, times(1)).run();
60+
}
61+
62+
@Test
63+
void matchesEquivalentKeysOnly() {
64+
RequestContext tenant1Context = RequestContext.forTenantId("first");
65+
RequestContext alternateTenant1Context = RequestContext.forTenantId("first");
66+
alternateTenant1Context.add("other", "value");
67+
RequestContext tenant2Context = RequestContext.forTenantId("second");
68+
69+
assertEquals(
70+
new DefaultContextualKey<>(tenant1Context, "input"),
71+
new DefaultContextualKey<>(tenant1Context, "input"));
72+
73+
assertEquals(
74+
new DefaultContextualKey<>(tenant1Context, "input"),
75+
new DefaultContextualKey<>(alternateTenant1Context, "input"));
76+
77+
assertNotEquals(
78+
new DefaultContextualKey<>(tenant1Context, "input"),
79+
new DefaultContextualKey<>(tenant2Context, "input"));
80+
81+
assertNotEquals(
82+
new DefaultContextualKey<>(tenant1Context, "input"),
83+
new DefaultContextualKey<>(tenant1Context, "other input"));
84+
}
85+
}

0 commit comments

Comments
 (0)