Initial OTel context propagation

Author: prdoyle

bosk-opentelemetry/README.md

@@ -0,0 +1,4 @@
+## bosk-opentelemetry
+
+This is the subproject for the published `bosk-opentelemetry` library,
+including support for propagating OpenTelemetry context via bosk diagnostic attributes.

bosk-opentelemetry/build.gradle

@@ -0,0 +1,29 @@
+
+plugins {
+	id 'bosk.development'
+	id 'bosk.maven-publish'
+	id 'com.github.spotbugs' version '6.1.2'
+}
+
+java {
+	toolchain {
+		languageVersion = JavaLanguageVersion.of(jdkVersion)
+	}
+}
+
+compileJava {
+	options.release = prodJavaVersion
+}
+
+compileTestJava {
+	options.release = null
+}
+
+repositories {
+	mavenCentral()
+}
+
+dependencies {
+	api project(":bosk-core")
+	implementation 'io.opentelemetry:opentelemetry-sdk:1.48.0'
+}

bosk-opentelemetry/src/main/java/works/bosk/opentelemetry/OpenTelemetryDriver.java

@@ -0,0 +1,26 @@
+package works.bosk.opentelemetry;
+
+import works.bosk.BoskDiagnosticContext;
+import works.bosk.BoskDriver;
+import works.bosk.DriverFactory;
+import works.bosk.DriverStack;
+import works.bosk.StateTreeNode;
+import works.bosk.drivers.DiagnosticScopeDriver;
+
+public sealed interface OpenTelemetryDriver extends BoskDriver permits ReceiverDriver {
+	/**
+	 * @return a {@link DriverFactory} that transmits OpenTelemetry context
+	 * across the given {@code subject} driver
+	 * via the {@link BoskDiagnosticContext diagnostic context}.
+	 *
+	 * @see OpenTelemetryRegistrar
+	 */
+	static <RR extends StateTreeNode> DriverFactory<RR> wrapping(DriverFactory<RR> subject) {
+		return DriverStack.of(
+			DiagnosticScopeDriver.factory(Utils::diagnosticScopeWithContextFromCurrentSpan),
+			subject,
+			ReceiverDriver.factory()
+		);
+	}
+
+}

bosk-opentelemetry/src/main/java/works/bosk/opentelemetry/OpenTelemetryRegistrar.java

@@ -0,0 +1,44 @@
+package works.bosk.opentelemetry;
+
+import works.bosk.BoskDiagnosticContext;
+import works.bosk.BoskHook;
+import works.bosk.HookRegistrar;
+import works.bosk.Reference;
+import works.bosk.RegistrarFactory;
+
+/**
+ * {@link HookRegistrar} that propagates
+ * OpenTelemetry context from the {@link BoskDiagnosticContext diagnostic context}
+ * into hooks.
+ * <p>
+ * Thread-locals are not automatically propagated into hooks,
+ * and so OpenTelemetry context must be explicitly propagated.
+ * This registrar retrieves diagnostic attributes placed there by {@link OpenTelemetryDriver}
+ * and propagate them into hooks.
+ */
+public final class OpenTelemetryRegistrar implements HookRegistrar {
+	final BoskDiagnosticContext diagnosticContext;
+	final HookRegistrar downstream;
+
+	OpenTelemetryRegistrar(BoskDiagnosticContext diagnosticContext, HookRegistrar downstream) {
+		this.diagnosticContext = diagnosticContext;
+		this.downstream = downstream;
+	}
+
+	/**
+	 * @return a {@link HookRegistrar} that propagates OpenTelemetry context
+	 * from the {@link BoskDiagnosticContext diagnostic context} into hooks.
+	 */
+	public static RegistrarFactory factory() {
+		return (b,d) -> new OpenTelemetryRegistrar(b.diagnosticContext(), d);
+	}
+
+	@Override
+	public <T> void registerHook(String name, Reference<T> scope, BoskHook<T> hook) {
+		downstream.registerHook(name, scope, ref -> {
+			try (var __ = Utils.contextFromDiagnosticAttributes(diagnosticContext).makeCurrent()) {
+				hook.onChanged(ref);
+			}
+		});
+	}
+}

bosk-opentelemetry/src/main/java/works/bosk/opentelemetry/ReceiverDriver.java

@@ -0,0 +1,80 @@
+package works.bosk.opentelemetry;
+
+import java.io.IOException;
+import java.lang.reflect.Type;
+import works.bosk.BoskDiagnosticContext;
+import works.bosk.BoskDriver;
+import works.bosk.DriverFactory;
+import works.bosk.Identifier;
+import works.bosk.Reference;
+import works.bosk.StateTreeNode;
+import works.bosk.exceptions.InvalidTypeException;
+
+/**
+ * Propagates OpenTelemetry context from the {@link BoskDiagnosticContext doagnostic context}
+ * into the downstream driver.
+ */
+final class ReceiverDriver implements OpenTelemetryDriver {
+
+	private final BoskDiagnosticContext diagnosticContext;
+	private final BoskDriver downstream;
+
+	public static <R extends StateTreeNode> DriverFactory<R> factory() {
+		return (b, d) -> new ReceiverDriver(b.diagnosticContext(), d);
+	}
+
+	ReceiverDriver(BoskDiagnosticContext diagnosticContext1, BoskDriver downstream) {
+		this.diagnosticContext = diagnosticContext1;
+		this.downstream = downstream;
+	}
+
+	@Override
+	public StateTreeNode initialRoot(Type rootType) throws InvalidTypeException, IOException, InterruptedException {
+		try (var __ = Utils.contextFromDiagnosticAttributes(diagnosticContext).makeCurrent()) {
+			return downstream.initialRoot(rootType);
+		}
+	}
+
+	@Override
+	public <T> void submitReplacement(Reference<T> target, T newValue) {
+		try (var __ = Utils.contextFromDiagnosticAttributes(diagnosticContext).makeCurrent()) {
+			downstream.submitReplacement(target, newValue);
+		}
+	}
+
+	@Override
+	public <T> void submitConditionalReplacement(Reference<T> target, T newValue, Reference<Identifier> precondition, Identifier requiredValue) {
+		try (var __ = Utils.contextFromDiagnosticAttributes(diagnosticContext).makeCurrent()) {
+			downstream.submitConditionalReplacement(target, newValue, precondition, requiredValue);
+		}
+	}
+
+	@Override
+	public <T> void submitConditionalCreation(Reference<T> target, T newValue) {
+		try (var __ = Utils.contextFromDiagnosticAttributes(diagnosticContext).makeCurrent()) {
+			downstream.submitConditionalCreation(target, newValue);
+		}
+	}
+
+	@Override
+	public <T> void submitDeletion(Reference<T> target) {
+		try (var __ = Utils.contextFromDiagnosticAttributes(diagnosticContext).makeCurrent()) {
+			downstream.submitDeletion(target);
+		}
+	}
+
+	@Override
+	public <T> void submitConditionalDeletion(Reference<T> target, Reference<Identifier> precondition, Identifier requiredValue) {
+		try (var __ = Utils.contextFromDiagnosticAttributes(diagnosticContext).makeCurrent()) {
+			downstream.submitConditionalDeletion(target, precondition, requiredValue);
+		}
+	}
+
+	@Override
+	public void flush() throws IOException, InterruptedException {
+		try (var __ = Utils.contextFromDiagnosticAttributes(diagnosticContext).makeCurrent()) {
+			downstream.flush();
+		}
+	}
+
+}

bosk-opentelemetry/src/main/java/works/bosk/opentelemetry/Utils.java

@@ -0,0 +1,86 @@
+package works.bosk.opentelemetry;
+
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.TextMapGetter;
+import io.opentelemetry.context.propagation.TextMapSetter;
+import java.util.HashMap;
+import java.util.Map;
+import works.bosk.BoskDiagnosticContext;
+import works.bosk.MapValue;
+
+public class Utils {
+	/**
+	 * Prefix for OpenTelemetry-related diagnostic attributes.
+	 */
+	static final String OTEL_PREFIX = "otel.";
+
+	/**
+	 * We use the w3c context propagation format.
+	 * Future-proof this putting the attributes under a sub-prefix so we can evolve this later.
+	 */
+	static final String W3C_SUB_PREFIX = "w3c.";
+
+	/**
+	 * On the setter side, we'll be given plain old OTel keys and need to add the prefixes.
+	 * We add {@link #W3C_SUB_PREFIX} to them here, and then the {@link #OTEL_PREFIX} prefix
+	 * is added by {@link BoskDiagnosticContext#withReplacedPrefix}.
+	 */
+	static final TextMapSetter<Map<String, String>> W3C_MAP_SETTER = (carrier, key, value) -> {
+		if (carrier != null) {
+			carrier.put(W3C_SUB_PREFIX + key, value);
+		}
+	};
+
+	/**
+	 * On the getter side, it's simplest to deal with plain old OTel keys.
+	 * Any prefixes are stripped off while we're building the carrier map.
+	 */
+	private static final TextMapGetter<Map<String, String>> MAP_GETTER = new TextMapGetter<>() {
+		@Override
+		public Iterable<String> keys(Map<String, String> carrier) {
+			return carrier.keySet();
+		}
+
+		@Override
+		public String get(Map<String, String> carrier, String key) {
+			if (carrier == null) {
+				return null;
+			} else {
+				return carrier.get(key);
+			}
+		}
+	};
+
+	static Context contextFromDiagnosticAttributes(BoskDiagnosticContext diagnosticContext) {
+		Map<String, String> otelAttributes = new HashMap<>();
+		diagnosticContext.getAttributes().forEach((k, v) -> {
+			if (k.startsWith(OTEL_PREFIX)) {
+				if (k.startsWith(OTEL_PREFIX + W3C_SUB_PREFIX)) {
+					otelAttributes.put(k.substring(OTEL_PREFIX.length() + W3C_SUB_PREFIX.length()), v);
+				} else {
+					throw new IllegalStateException("Only w3c context propagation is supported; unexpected OpenTelemetry key from diagnostic context: " + k);
+				}
+			}
+		});
+		return W3CTraceContextPropagator.getInstance()
+			.extract(Context.current(), otelAttributes, MAP_GETTER);
+	}
+
+	/**
+	 * Stashes OpenTelemetry context into the {@link BoskDiagnosticContext doagnostic context},
+	 * where it can be retrieved by
+	 * {@link ReceiverDriver} to propagate the context downstream,
+	 * and by {@link OpenTelemetryRegistrar} to propagate the context into hooks.
+	 */
+	static BoskDiagnosticContext.DiagnosticScope diagnosticScopeWithContextFromCurrentSpan(BoskDiagnosticContext dc) {
+		return dc.withReplacedPrefix(OTEL_PREFIX, w3cContextFromCurrentSpan());
+	}
+
+	static MapValue<String> w3cContextFromCurrentSpan() {
+		Map<String, String> result = new HashMap<>();
+		W3CTraceContextPropagator.getInstance()
+			.inject(Context.current(), result, W3C_MAP_SETTER);
+		return MapValue.copyOf(result);
+	}
+}

bosk-opentelemetry/src/main/java/works/bosk/opentelemetry/package-info.java

@@ -0,0 +1,16 @@
+/**
+ * Provides OpenTelemetry context propagation for Bosk.
+ * <p>
+ * To use this, configure your bosk instance to do the following two things:
+ * <ol>
+ *     <li>
+ *         Use {@link works.bosk.opentelemetry.OpenTelemetryDriver#wrapping} on any driver
+ *         that does not implicitly propagate thread context by calling its downstream driver
+ *         synchronously on the same thread.
+ *     </li>
+ *     <li>
+ *         Use {@link works.bosk.opentelemetry.OpenTelemetryRegistrar#factory()}.
+ *     </li>
+ * </ol>
+ */
+package works.bosk.opentelemetry;

bosk-opentelemetry/src/test/java/works/bosk/opentelemetry/OpenTelemetryDriverTest.java

@@ -0,0 +1,87 @@
+package works.bosk.opentelemetry;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.TraceId;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.propagation.ContextPropagators;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
+import io.opentelemetry.sdk.trace.SdkTracerProvider;
+import java.io.IOException;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import works.bosk.Bosk;
+import works.bosk.DriverStack;
+import works.bosk.StateTreeNode;
+import works.bosk.drivers.BufferingDriver;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class OpenTelemetryDriverTest {
+
+	private static OpenTelemetrySdk openTelemetry;
+
+	public record Root(Integer revision) implements StateTreeNode {}
+
+	@BeforeAll
+	static void setup() {
+		openTelemetry = OpenTelemetrySdk.builder()
+			.setTracerProvider(SdkTracerProvider.builder().build())
+			.setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
+			.buildAndRegisterGlobal();
+	}
+
+	@AfterAll
+	static void teardown() {
+		openTelemetry.close();
+	}
+
+	@Test
+	void wrapping_propagatesTraceId() throws InterruptedException, IOException {
+		var bosk = new Bosk<Root>(
+			"test-bosk",
+			Root.class,
+			_ -> new Root(0),
+			DriverStack.of(
+				OpenTelemetryDriver.wrapping(
+					// Use a driver that does not call its downstream driver synchronously on the same thread
+					// so that the OpenTelemetry thread context is not propagated implicitly.
+					// Otherwise, this isn't much of a test.
+					BufferingDriver.factory()
+				)
+			),
+			OpenTelemetryRegistrar.factory());
+		record Observation(int revision, String traceID) { }
+		BlockingQueue<Observation> observations = new LinkedBlockingQueue<>();
+		bosk.registerHook("attribute observer", bosk.rootReference(), ref -> {
+			observations.add(new Observation(ref.value().revision(),
+				Span.current().getSpanContext().getTraceId()));
+		});
+
+		// There's an initial observation from the time the hook was registered.
+		// At that time, there was no active trace
+		assertEquals(new Observation(0, TraceId.getInvalid()), observations.take());
+
+		// Submit a change with an active OpenTelemetry span and observe the attributes
+		Observation expected;
+		var span = openTelemetry.getTracer("test-scope")
+			.spanBuilder("test-span")
+			.setParent(Context.current())
+			.startSpan();
+		try (var _ = span.makeCurrent()) {
+			bosk.driver().submitReplacement(bosk.rootReference(), new Root(1));
+			expected = new Observation(1, Span.current().getSpanContext().getTraceId());
+		} finally {
+			span.end();
+		}
+
+		bosk.driver().flush();
+		// Check that the trace context was propagated all the way into the hook
+		assertEquals(expected, observations.take());
+		assert observations.isEmpty();
+	}
+
+}