diff --git a/backend/common/src/main/java/ai/verta/modeldb/common/futures/FutureJdbi.java b/backend/common/src/main/java/ai/verta/modeldb/common/futures/FutureJdbi.java index 79c6031549..a05bc25a66 100644 --- a/backend/common/src/main/java/ai/verta/modeldb/common/futures/FutureJdbi.java +++ b/backend/common/src/main/java/ai/verta/modeldb/common/futures/FutureJdbi.java @@ -4,6 +4,7 @@ import org.jdbi.v3.core.HandleConsumer; import org.jdbi.v3.core.Jdbi; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; @@ -17,18 +18,24 @@ public FutureJdbi(Jdbi jdbi, Executor executor) { } public InternalFuture withHandle(HandleCallback callback) { - CompletableFuture promise = new CompletableFuture(); - - executor.execute( + return InternalFuture.trace( () -> { - try { - promise.complete(jdbi.withHandle(callback)); - } catch (Throwable e) { - promise.completeExceptionally(e); - } - }); - - return InternalFuture.from(promise); + CompletableFuture promise = new CompletableFuture(); + + executor.execute( + () -> { + try { + promise.complete(jdbi.withHandle(callback)); + } catch (Throwable e) { + promise.completeExceptionally(e); + } + }); + + return InternalFuture.from(promise); + }, + "jdbi.withHandle", + Map.of("caller", String.format("%s:%d", Thread.currentThread().getStackTrace()[2].getFileName(), Thread.currentThread().getStackTrace()[2].getLineNumber())), + executor); } public InternalFuture withHandleCompose( @@ -37,18 +44,24 @@ public InternalFuture withHandleCompose( } public InternalFuture useHandle(final HandleConsumer consumer) { - CompletableFuture promise = new CompletableFuture(); - - executor.execute( + return InternalFuture.trace( () -> { - try { - jdbi.useHandle(consumer); - promise.complete(null); - } catch (Throwable e) { - promise.completeExceptionally(e); - } - }); - - return InternalFuture.from(promise); + CompletableFuture promise = new CompletableFuture(); + + executor.execute( + () -> { + try { + jdbi.useHandle(consumer); + promise.complete(null); + } catch (Throwable e) { + promise.completeExceptionally(e); + } + }); + + return InternalFuture.from(promise); + }, + "jdbi.useHandle", + Map.of("caller", String.format("%s:%d", Thread.currentThread().getStackTrace()[2].getFileName(), Thread.currentThread().getStackTrace()[2].getLineNumber())), + executor); } } diff --git a/backend/common/src/main/java/ai/verta/modeldb/common/futures/InternalFuture.java b/backend/common/src/main/java/ai/verta/modeldb/common/futures/InternalFuture.java index be25a9d534..5464d2d404 100644 --- a/backend/common/src/main/java/ai/verta/modeldb/common/futures/InternalFuture.java +++ b/backend/common/src/main/java/ai/verta/modeldb/common/futures/InternalFuture.java @@ -6,6 +6,7 @@ import io.opentracing.contrib.grpc.ActiveSpanContextSource; import io.opentracing.contrib.grpc.ActiveSpanSource; import io.opentracing.tag.Tags; +import io.opentracing.util.GlobalTracer; import net.bytebuddy.implementation.bytecode.Throw; import org.apache.logging.log4j.util.TriConsumer; @@ -48,26 +49,33 @@ static private Span createSpanFromParent(Tracer tracer, SpanContext parentSpanCo spanBuilder = tracer.buildSpan(operationName).asChildOf(parentSpanContext); } - for (var entry : tags.entrySet()) { - spanBuilder = spanBuilder.withTag(entry.getKey(), entry.getValue()); + if (tags != null) { + for (var entry : tags.entrySet()) { + spanBuilder = spanBuilder.withTag(entry.getKey(), entry.getValue()); + } } return spanBuilder.start(); } - public static InternalFuture trace(Supplier> supplier, String operationName, Map tags, Tracer tracer, Executor executor) { + public static InternalFuture trace(Supplier> supplier, String operationName, Map tags, Executor executor) { + if (!GlobalTracer.isRegistered()) + return supplier.get(); + + final var tracer = GlobalTracer.get(); + final var spanContext = getActiveSpanContext(tracer); final var span = createSpanFromParent(tracer, spanContext, operationName, tags); final var promise = new CompletableFuture(); - supplier.get().stage.whenComplete((v, t) -> { + supplier.get().stage.whenCompleteAsync((v, t) -> { span.finish(); if (t != null) { promise.completeExceptionally(t); } else { promise.complete(v); } - }); + }, executor); return InternalFuture.from(promise); }