Trace jdbi futures (#2515)

* Trace jdbi * Add caller * Fix index
VertaAI · Jul 26, 2021 · b713ced · b713ced
1 parent 20765ab
commit b713ced
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 28 deletions.
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
@@ -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 <R, T extends Exception> InternalFuture<R> withHandle(HandleCallback<R, T> callback) {
-    CompletableFuture<R> promise = new CompletableFuture<R>();
-
-    executor.execute(
+    return InternalFuture.trace(
         () -> {
-          try {
-            promise.complete(jdbi.withHandle(callback));
-          } catch (Throwable e) {
-            promise.completeExceptionally(e);
-          }
-        });
-
-    return InternalFuture.from(promise);
+          CompletableFuture<R> promise = new CompletableFuture<R>();
+
+          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 <R, T extends Exception> InternalFuture<R> withHandleCompose(
@@ -37,18 +44,24 @@ public <R, T extends Exception> InternalFuture<R> withHandleCompose(
   }
 
   public <T extends Exception> InternalFuture<Void> useHandle(final HandleConsumer<T> consumer) {
-    CompletableFuture<Void> promise = new CompletableFuture<Void>();
-
-    executor.execute(
+    return InternalFuture.trace(
         () -> {
-          try {
-            jdbi.useHandle(consumer);
-            promise.complete(null);
-          } catch (Throwable e) {
-            promise.completeExceptionally(e);
-          }
-        });
-
-    return InternalFuture.from(promise);
+          CompletableFuture<Void> promise = new CompletableFuture<Void>();
+
+          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
@@ -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 <T> InternalFuture<T> trace(Supplier<InternalFuture<T>> supplier, String operationName, Map<String,String> tags, Tracer tracer, Executor executor) {
+    public static <T> InternalFuture<T> trace(Supplier<InternalFuture<T>> supplier, String operationName, Map<String,String> 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<T>();
-        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);
     }