Skip to content

Commit a41844e

Browse files
PerfectSlayerPerfectSlayer
committed
feat(context): Introduce context helpers API
1 parent 0dc9c08 commit a41844e

File tree

4 files changed

+354
-3
lines changed

4 files changed

+354
-3
lines changed

‎components/context/src/main/java/datadog/context/ContextHelpers.java

+143
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
package datadog.context;
2+
3+
import static java.lang.Math.max;
4+
import static java.util.Arrays.copyOfRange;
5+
import static java.util.Objects.requireNonNull;
6+
7+
import java.util.ArrayList;
8+
import java.util.List;
9+
import java.util.function.BinaryOperator;
10+
11+
/**
12+
* Static helpers to manipulate context collections.
13+
*
14+
* <p>Typical usages include:
15+
*
16+
* <pre>{@code
17+
* // Finding a context value from multiple sources:
18+
* Span span = findFirst(spanKey, message, request, CURRENT)
19+
* // Find all context values from different sources:
20+
* List<Error> errors = findAll(errorKey, message, request, CURRENT)
21+
* // Capture multiple contexts in a single one:
22+
* Context aggregate = combine(message, request, CURRENT)
23+
* // Combine multiple contexts into a single one using custom merge rules:
24+
* Context combined = combine(
25+
* (current, next) -> {
26+
* var metric = current.get(metricKey);
27+
* var nextMetric = next.get(metricKey);
28+
* return current.with(metricKey, metric.add(nextMetric));
29+
* }, message, request, CURRENT);
30+
* }</pre>
31+
*
32+
* where {@link #CURRENT} denotes a carrier with the current context.
33+
*/
34+
public final class ContextHelpers {
35+
/** A helper object carrying the {@link Context#current()} context. */
36+
public static final Object CURRENT = new Object();
37+
38+
private ContextHelpers() {}
39+
40+
/**
41+
* Find the first context value from given context carriers.
42+
*
43+
* @param key The key used to store the value.
44+
* @param carriers The carrier to get context and value from.
45+
* @param <T> The type of the value to look for.
46+
* @return The first context value found, {@code null} if not found.
47+
*/
48+
public static <T> T findFirst(ContextKey<T> key, Object... carriers) {
49+
requireNonNull(key, "key cannot be null");
50+
for (Object carrier : carriers) {
51+
requireNonNull(carrier, "carrier cannot be null");
52+
Context context = carrier == CURRENT ? Context.current() : Context.from(carrier);
53+
T value = context.get(key);
54+
if (value != null) {
55+
return value;
56+
}
57+
}
58+
return null;
59+
}
60+
61+
/**
62+
* Find all the context values from the given context carriers.
63+
*
64+
* @param key The key used to store the value.
65+
* @param carriers The carriers to get context and value from.
66+
* @param <T> The type of the values to look for.
67+
* @return A list of all values found, in context order.
68+
*/
69+
public static <T> List<T> findAll(ContextKey<T> key, Object... carriers) {
70+
requireNonNull(key, "key cannot be null");
71+
List<T> values = new ArrayList<>(carriers.length);
72+
for (Object carrier : carriers) {
73+
requireNonNull(carrier, "carrier cannot be null");
74+
Context context = carrier == CURRENT ? Context.current() : Context.from(carrier);
75+
T value = context.get(key);
76+
if (value != null) {
77+
values.add(value);
78+
}
79+
}
80+
return values;
81+
}
82+
83+
/**
84+
* Combine contexts and their values, keeping the first founds.
85+
*
86+
* @param contexts The contexts to combine.
87+
* @return A context containing all the values from all the given context, keeping the first value
88+
* found for a given key.
89+
*/
90+
public static Context combine(Context... contexts) {
91+
return combine(ContextHelpers::combineKeepingFirst, contexts);
92+
}
93+
94+
/**
95+
* Combine multiple contexts into a single one.
96+
*
97+
* @param combiner The context combiner, taking already combined context as first parameter, any
98+
* following one as second parameter, and returning the combined context.
99+
* @param contexts The contexts to combine.
100+
* @return The combined context.
101+
*/
102+
public static Context combine(BinaryOperator<Context> combiner, Context... contexts) {
103+
requireNonNull(combiner, "combiner cannot be null");
104+
Context result = new IndexedContext(new Object[0]);
105+
for (Context context : contexts) {
106+
requireNonNull(context, "context cannot be null");
107+
result = combiner.apply(result, context);
108+
}
109+
return result;
110+
}
111+
112+
private static Context combineKeepingFirst(Context current, Context next) {
113+
if (!(current instanceof IndexedContext)) {
114+
throw new IllegalStateException("Left context is supposed to be an IndexedContext");
115+
}
116+
IndexedContext currentIndexed = (IndexedContext) current;
117+
if (next instanceof EmptyContext) {
118+
return current;
119+
} else if (next instanceof SingletonContext) {
120+
SingletonContext nextSingleton = (SingletonContext) next;
121+
Object[] store =
122+
copyOfRange(
123+
currentIndexed.store, 0, max(currentIndexed.store.length, nextSingleton.index + 1));
124+
if (store[nextSingleton.index] == null) {
125+
store[nextSingleton.index] = nextSingleton.value;
126+
}
127+
return new IndexedContext(store);
128+
} else if (next instanceof IndexedContext) {
129+
IndexedContext nextIndexed = (IndexedContext) next;
130+
Object[] store =
131+
copyOfRange(
132+
currentIndexed.store, 0, max(currentIndexed.store.length, nextIndexed.store.length));
133+
for (int i = 0; i < nextIndexed.store.length; i++) {
134+
Object nextValue = nextIndexed.store[i];
135+
if (nextValue != null && store[i] == null) {
136+
store[i] = nextValue;
137+
}
138+
}
139+
return new IndexedContext(store);
140+
}
141+
throw new IllegalStateException("Unsupported context type: " + next.getClass().getName());
142+
}
143+
}

‎components/context/src/main/java/datadog/context/IndexedContext.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
/** {@link Context} containing many values. */
1212
@ParametersAreNonnullByDefault
1313
final class IndexedContext implements Context {
14-
private final Object[] store;
14+
final Object[] store;
1515

1616
IndexedContext(Object[] store) {
1717
this.store = store;

‎components/context/src/main/java/datadog/context/SingletonContext.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@
1010
/** {@link Context} containing a single value. */
1111
@ParametersAreNonnullByDefault
1212
final class SingletonContext implements Context {
13-
private final int index;
14-
private final Object value;
13+
final int index;
14+
final Object value;
1515

1616
SingletonContext(int index, Object value) {
1717
this.index = index;

‎components/context/src/test/java/datadog/context/ContextHelpersTest.java

+208
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
package datadog.context;
2+
3+
import static datadog.context.Context.current;
4+
import static datadog.context.Context.root;
5+
import static datadog.context.ContextHelpers.CURRENT;
6+
import static datadog.context.ContextHelpers.combine;
7+
import static datadog.context.ContextHelpers.findAll;
8+
import static datadog.context.ContextHelpers.findFirst;
9+
import static datadog.context.ContextTest.BOOLEAN_KEY;
10+
import static datadog.context.ContextTest.FLOAT_KEY;
11+
import static datadog.context.ContextTest.STRING_KEY;
12+
import static java.util.Arrays.asList;
13+
import static java.util.Collections.emptyList;
14+
import static java.util.Collections.singleton;
15+
import static java.util.logging.Level.ALL;
16+
import static java.util.logging.Level.INFO;
17+
import static java.util.logging.Level.SEVERE;
18+
import static java.util.logging.Level.WARNING;
19+
import static org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertIterableEquals;
21+
import static org.junit.jupiter.api.Assertions.assertNotNull;
22+
import static org.junit.jupiter.api.Assertions.assertNull;
23+
import static org.junit.jupiter.api.Assertions.assertThrows;
24+
import static org.junit.jupiter.params.provider.Arguments.arguments;
25+
26+
import java.util.function.BinaryOperator;
27+
import java.util.logging.Level;
28+
import java.util.stream.Stream;
29+
import org.junit.jupiter.api.BeforeAll;
30+
import org.junit.jupiter.api.Test;
31+
import org.junit.jupiter.params.ParameterizedTest;
32+
import org.junit.jupiter.params.provider.Arguments;
33+
import org.junit.jupiter.params.provider.MethodSource;
34+
35+
class ContextHelpersTest {
36+
private static final Object CARRIER_1 = new Object();
37+
private static final Object CARRIER_2 = new Object();
38+
private static final Object UNSET_CARRIER = new Object();
39+
private static final Object NON_CARRIER = new Object();
40+
private static final String VALUE_1 = "value1";
41+
private static final String VALUE_2 = "value2";
42+
43+
@BeforeAll
44+
static void init() {
45+
Context context1 = root().with(STRING_KEY, VALUE_1);
46+
context1.attachTo(CARRIER_1);
47+
48+
Context context2 = root().with(STRING_KEY, VALUE_2);
49+
context2.attachTo(CARRIER_2);
50+
51+
root().attachTo(UNSET_CARRIER);
52+
}
53+
54+
@ParameterizedTest
55+
@MethodSource("findFirstArguments")
56+
void testFindFirst(Object[] carriers, String expected) {
57+
assertEquals(expected, findFirst(STRING_KEY, carriers), "Cannot find first value");
58+
}
59+
60+
static Stream<Arguments> findFirstArguments() {
61+
return Stream.of(
62+
arguments(emptyArray(), null),
63+
arguments(arrayOf(NON_CARRIER), null),
64+
arguments(arrayOf(UNSET_CARRIER), null),
65+
arguments(arrayOf(CARRIER_1), VALUE_1),
66+
arguments(arrayOf(CARRIER_1, CARRIER_2), VALUE_1),
67+
arguments(arrayOf(NON_CARRIER, CARRIER_1), VALUE_1),
68+
arguments(arrayOf(UNSET_CARRIER, CARRIER_1), VALUE_1),
69+
arguments(arrayOf(CARRIER_1, NON_CARRIER), VALUE_1),
70+
arguments(arrayOf(CARRIER_1, UNSET_CARRIER), VALUE_1));
71+
}
72+
73+
@ParameterizedTest
74+
@MethodSource("findAllArguments")
75+
void testFindAll(Object[] carriers, Iterable<String> expected) {
76+
assertIterableEquals(expected, findAll(STRING_KEY, carriers), "Cannot find all values");
77+
}
78+
79+
static Stream<Arguments> findAllArguments() {
80+
return Stream.of(
81+
arguments(emptyArray(), emptyList()),
82+
arguments(arrayOf(CARRIER_1), singleton(VALUE_1)),
83+
arguments(arrayOf(CARRIER_1, CARRIER_2), asList(VALUE_1, VALUE_2)),
84+
arguments(arrayOf(NON_CARRIER, CARRIER_1), singleton(VALUE_1)),
85+
arguments(arrayOf(UNSET_CARRIER, CARRIER_1), singleton(VALUE_1)),
86+
arguments(arrayOf(CARRIER_1, NON_CARRIER), singleton(VALUE_1)),
87+
arguments(arrayOf(CARRIER_1, UNSET_CARRIER), singleton(VALUE_1)));
88+
}
89+
90+
@Test
91+
void testNullCarriers() {
92+
assertThrows(
93+
NullPointerException.class, () -> findFirst(null, CARRIER_1), "Should fail on null key");
94+
assertThrows(
95+
NullPointerException.class,
96+
() -> findFirst(STRING_KEY, (Object) null),
97+
"Should fail on null context");
98+
assertThrows(
99+
NullPointerException.class,
100+
() -> findFirst(STRING_KEY, null, CARRIER_1),
101+
"Should fail on null context");
102+
assertThrows(
103+
NullPointerException.class, () -> findAll(null, CARRIER_1), "Should fail on null key");
104+
assertThrows(
105+
NullPointerException.class,
106+
() -> findAll(STRING_KEY, (Object) null),
107+
"Should fail on null context");
108+
assertThrows(
109+
NullPointerException.class,
110+
() -> findAll(STRING_KEY, null, CARRIER_1),
111+
"Should fail on null context");
112+
}
113+
114+
@Test
115+
void testCurrent() {
116+
assertEquals(root(), current(), "Current context is already set");
117+
Context context = root().with(STRING_KEY, VALUE_1);
118+
try (ContextScope ignored = context.attach()) {
119+
assertEquals(
120+
VALUE_1, findFirst(STRING_KEY, CURRENT), "Failed to get value from current context");
121+
assertIterableEquals(
122+
singleton(VALUE_1),
123+
findAll(STRING_KEY, CURRENT),
124+
"Failed to get value from current context");
125+
}
126+
assertEquals(root(), current(), "Current context stayed attached");
127+
}
128+
129+
@Test
130+
void testCombine() {
131+
Context context1 = root().with(STRING_KEY, VALUE_1).with(BOOLEAN_KEY, true);
132+
Context context2 = root().with(STRING_KEY, VALUE_2).with(FLOAT_KEY, 3.14F);
133+
Context context3 = root();
134+
Context context4 = root().with(FLOAT_KEY, 567F);
135+
136+
Context combined = combine(context1, context2, context3, context4);
137+
assertEquals(VALUE_1, combined.get(STRING_KEY), "First duplicate value should be kept");
138+
assertEquals(true, combined.get(BOOLEAN_KEY), "Values from first context should be kept");
139+
assertEquals(3.14f, combined.get(FLOAT_KEY), "Values from second context should be kept");
140+
}
141+
142+
@Test
143+
void testCombiner() {
144+
ContextKey<ErrorStats> errorKey = ContextKey.named("error");
145+
Context context1 = root().with(errorKey, ErrorStats.from(INFO, 12)).with(STRING_KEY, VALUE_1);
146+
Context context2 = root().with(errorKey, ErrorStats.from(SEVERE, 1)).with(FLOAT_KEY, 3.14F);
147+
Context context3 = root().with(errorKey, ErrorStats.from(WARNING, 6)).with(BOOLEAN_KEY, true);
148+
149+
BinaryOperator<Context> errorStatsMerger =
150+
(left, right) -> {
151+
ErrorStats mergedStats = ErrorStats.merge(left.get(errorKey), right.get(errorKey));
152+
return left.with(errorKey, mergedStats);
153+
};
154+
Context combined = combine(errorStatsMerger, context1, context2, context3);
155+
ErrorStats combinedStats = combined.get(errorKey);
156+
assertNotNull(combinedStats, "Failed to combined error stats");
157+
assertEquals(19, combinedStats.errorCount, "Failed to combine error stats");
158+
assertEquals(SEVERE, combinedStats.maxLevel, "Failed to combine error stats");
159+
assertNull(combined.get(STRING_KEY), "Combiner should drop any other context values");
160+
assertNull(combined.get(FLOAT_KEY), "Combiner should drop any other context values");
161+
assertNull(combined.get(BOOLEAN_KEY), "Combiner should drop any other context values");
162+
}
163+
164+
@Test
165+
void testNullCombine() {
166+
assertThrows(
167+
NullPointerException.class,
168+
() -> combine((BinaryOperator<Context>) null, root()),
169+
"Should fail on null combiner");
170+
assertThrows(
171+
NullPointerException.class,
172+
() -> combine((left, right) -> left, (Context) null),
173+
"Should fail on null context");
174+
}
175+
176+
private static class ErrorStats {
177+
int errorCount;
178+
Level maxLevel;
179+
180+
public ErrorStats() {
181+
this.errorCount = 0;
182+
this.maxLevel = ALL;
183+
}
184+
185+
public static ErrorStats from(Level logLevel, int count) {
186+
ErrorStats stats = new ErrorStats();
187+
stats.errorCount = count;
188+
stats.maxLevel = logLevel;
189+
return stats;
190+
}
191+
192+
public static ErrorStats merge(ErrorStats a, ErrorStats b) {
193+
if (a == null) {
194+
return b;
195+
}
196+
Level maxLevel = a.maxLevel.intValue() > b.maxLevel.intValue() ? a.maxLevel : b.maxLevel;
197+
return from(maxLevel, a.errorCount + b.errorCount);
198+
}
199+
}
200+
201+
private static Object[] emptyArray() {
202+
return new Object[0];
203+
}
204+
205+
private static Object[] arrayOf(Object... objects) {
206+
return objects;
207+
}
208+
}

0 commit comments

Comments
 (0)