Tracing Asynchronous Database Transactions Across Thread Pools with OpenTelemetry and Java Agent — code snippets

snippet-1.txt

package com.example.observability;

import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class ManualContextPropagator {
    private static final ExecutorService executor = Executors.newFixedThreadPool(10);
    private final Tracer tracer;

    public ManualContextPropagator(Tracer tracer) {
        this.tracer = tracer;
    }

    public void executeAsynchronously(Runnable databaseTask) {
        Span parentSpan = tracer.spanBuilder("parent-transaction-step").startSpan();
        try (Scope scope = parentSpan.makeCurrent()) {
            // Context.current() captures the active span (parentSpan) from the caller thread
            Context currentContext = Context.current();
            
            // Wrap the runnable to propagate the context automatically to the executor thread
            Runnable wrappedTask = currentContext.wrap(() -> {
                // Inside the executor thread, the parentSpan is now active in ThreadLocal
                Span childSpan = tracer.spanBuilder("async-db-query")
                        .setParent(Context.current())
                        .startSpan();
                try (Scope childScope = childSpan.makeCurrent()) {
                    databaseTask.run();
                } catch (Exception e) {
                    childSpan.recordException(e);
                    throw e;
                } finally {
                    childSpan.end();
                }
            });

            executor.submit(wrappedTask);
        } finally {
            parentSpan.end();
        }
    }
}

snippet-2.txt

# System properties and agent configuration for OpenTelemetry Java Agent
# Add these flags to your JVM startup options

# Enable debugging of instrumentation to verify class transformations
otel.javaagent.debug=false
otel.javaagent.logging=application

# Explicitly enable context propagation for custom executor classes
otel.instrumentation.executors.enabled=true
otel.instrumentation.executors.include=com.example.observability.CustomThreadPoolExecutor,org.apache.tomcat.util.threads.ThreadPoolExecutor

# Enable JDBC and Connection Pool instrumentation
otel.instrumentation.jdbc.enabled=true
otel.instrumentation.hikari.enabled=true

# Enable Reactor context propagation for reactive boundaries
otel.instrumentation.reactor.enabled=true

# Define target collector endpoint
otel.exporter.otlp.endpoint=http://otel-collector.monitoring.svc.cluster.local:4317
otel.exporter.otlp.protocol=grpc

snippet-3.txt

package com.example.observability;

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.instrumentation.jdbc.datasource.OpenTelemetryDataSource;
import javax.sql.DataSource;

public class DataSourceConfig {

    public static DataSource createInstrumentedDataSource(OpenTelemetry openTelemetry) {
        HikariConfig config = new HikariConfig();
        config.setJdbcUrl("jdbc:postgresql://localhost:5432/production_db");
        config.setUsername("db_user");
        config.setPassword("super_secure_password");
        config.setMaximumPoolSize(25);
        config.setMinimumIdle(5);
        config.setPoolName("HikariAsyncPool");
        config.setConnectionTimeout(30000); // 30 seconds
        config.setIdleTimeout(600000); // 10 minutes

        HikariDataSource hikariDataSource = new HikariDataSource(config);

        // Wrap HikariDataSource with OpenTelemetryDataSource to intercept connection checkouts
        // and link active OpenTelemetry trace context to the underlying SQL spans.
        return new OpenTelemetryDataSource(hikariDataSource);
    }
}

snippet-4.txt

package com.example.observability;

import io.opentelemetry.api.trace.Span;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import reactor.core.publisher.Mono;
import reactor.util.context.ContextView;

public class ReactorTraceHelper {

    // Wrap a Mono block so that its internal operations execute within the active OTel context
    public static <T> Mono<T> withTraceContext(Mono<T> source) {
        return Mono.deferContextual(contextView -> {
            // Extract OTel Context from Reactor's ContextView, defaulting to the current thread context
            Context otelContext = contextView.getOrDefault(Context.class, Context.current());
            
            return source.contextWrite(ctx -> ctx.put(Context.class, otelContext))
                    .doOnEach(signal -> {
                        // Ensure context is active during onNext or onError processing
                        if (signal.isOnNext() || signal.isOnError()) {
                            try (Scope scope = otelContext.makeCurrent()) {
                                // Downstream operators can now read the active span from ThreadLocal
                            }
                        }
                    });
        });
    }

    // Capture the current OpenTelemetry context and convert it into a Reactor context
    public static reactor.util.context.Context currentOtelContext() {
        return reactor.util.context.Context.of(Context.class, Context.current());
    }
}

snippet-5.txt

package com.example.observability;

import io.opentelemetry.context.Context;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;

public class ContextPropagatingExecutor implements Executor {
    private final Executor delegate;

    public ContextPropagatingExecutor(Executor delegate) {
        this.delegate = delegate;
    }

    @Override
    public void execute(Runnable command) {
        // Capture active trace context on the caller thread
        Context currentContext = Context.current();
        
        // Submit wrapped command that restores context on the worker thread
        delegate.execute(() -> {
            try (Scope scope = currentContext.makeCurrent()) {
                command.run();
            } // Scope is closed here, preventing ThreadLocal context leakage
        });
    }

    // Wrap an entire ExecutorService using OpenTelemetry's built-in task wrapper
    public static ExecutorService wrapExecutorService(ExecutorService executorService) {
        return Context.taskWrapping(executorService);
    }
}

snippet-6.txt

package com.example.observability;

import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.api.trace.StatusCode;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import java.sql.Connection;
import java.sql.PreparedStatement;
import javax.sql.DataSource;

public class TracedTransactionExecutor {
    private final Tracer tracer;
    private final DataSource dataSource;

    public TracedTransactionExecutor(Tracer tracer, DataSource dataSource) {
        this.tracer = tracer;
        this.dataSource = dataSource;
    }

    public void executeTracedTransaction(String accountId, double amount) {
        // Start parent span for the database transaction block
        Span txSpan = tracer.spanBuilder("db.transaction")
                .setSpanKind(SpanKind.CLIENT)
                .setAttribute("db.system", "postgresql")
                .setAttribute("db.operation", "UPDATE")
                .startSpan();

        try (Scope txScope = txSpan.makeCurrent()) {
            Span connSpan = tracer.spanBuilder("db.connection_acquisition").startSpan();
            Connection conn;
            try {
                conn = dataSource.getConnection();
            } catch (Exception e) {
                connSpan.recordException(e);
                connSpan.setStatus(StatusCode.ERROR, "Failed to acquire connection");
                throw new RuntimeException(e);
            } finally {
                connSpan.end();
            }

            conn.setAutoCommit(false);
            try {
                Span querySpan = tracer.spanBuilder("db.query.update_balance")
                        .setAttribute("db.statement", "UPDATE accounts SET balance = balance + ? WHERE id = ?")
                        .startSpan();
                
                try (Scope queryScope = querySpan.makeCurrent();
                     PreparedStatement stmt = conn.prepareStatement("UPDATE accounts SET balance = balance + ? WHERE id = ?")) {
                    stmt.setDouble(1, amount);
                    stmt.setString(2, accountId);
                    stmt.executeUpdate();
                } catch (Exception e) {
                    querySpan.recordException(e);
                    querySpan.setStatus(StatusCode.ERROR, e.getMessage());
                    throw e;
                } finally {
                    querySpan.end();
                }

                conn.commit();
                txSpan.setAttribute("db.transaction.status", "COMMITTED");
            } catch (Exception e) {
                conn.rollback();
                txSpan.setAttribute("db.transaction.status", "ROLLED_BACK");
                txSpan.recordException(e);
                txSpan.setStatus(StatusCode.ERROR, "Transaction rolled back due to error");
                throw new RuntimeException("Transaction failed, rolled back", e);
            } finally {
                conn.close();
            }
        } finally {
            txSpan.end();
        }
    }
}