From 1db565369bbd20f81d09a761c567207067929b02 Mon Sep 17 00:00:00 2001
From: "ievgen.degtiarenko" <ievgen.degtiarenko@elastic.co>
Date: Fri, 7 Jun 2024 11:51:31 +0200
Subject: [PATCH] Utility to create spans for synchronous blocks of code

---
 .../apm/internal/tracing/APMTracer.java       |  5 ++
 .../org/elasticsearch/telemetry/Trace.java    | 54 +++++++++++++++++++
 .../telemetry/tracing/Tracer.java             | 11 ++++
 3 files changed, 70 insertions(+)
 create mode 100644 server/src/main/java/org/elasticsearch/telemetry/Trace.java

diff --git a/modules/apm/src/main/java/org/elasticsearch/telemetry/apm/internal/tracing/APMTracer.java b/modules/apm/src/main/java/org/elasticsearch/telemetry/apm/internal/tracing/APMTracer.java
index 0b020a24eeffb..b252450f60e1a 100644
--- a/modules/apm/src/main/java/org/elasticsearch/telemetry/apm/internal/tracing/APMTracer.java
+++ b/modules/apm/src/main/java/org/elasticsearch/telemetry/apm/internal/tracing/APMTracer.java
@@ -103,6 +103,11 @@ public void setEnabled(boolean enabled) {
         }
     }
 
+    @Override
+    public boolean isEnabled() {
+        return enabled;
+    }
+
     public void setIncludeNames(List<String> includeNames) {
         this.includeNames = includeNames;
         this.filterAutomaton = buildAutomaton(includeNames, excludeNames);
diff --git a/server/src/main/java/org/elasticsearch/telemetry/Trace.java b/server/src/main/java/org/elasticsearch/telemetry/Trace.java
new file mode 100644
index 0000000000000..ce10ac9aa73e6
--- /dev/null
+++ b/server/src/main/java/org/elasticsearch/telemetry/Trace.java
@@ -0,0 +1,54 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.telemetry;
+
+import org.elasticsearch.common.UUIDs;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.telemetry.tracing.Traceable;
+import org.elasticsearch.telemetry.tracing.Tracer;
+import org.elasticsearch.threadpool.ThreadPool;
+
+import java.util.Map;
+
+public class Trace {
+
+    public static Releasable trace(ThreadPool threadPool, Tracer tracer, String name) {
+        return trace(threadPool, tracer, name, Map.of());
+    }
+
+    /**
+     * Creates a new trace for a synchronous block of code.
+     * @return a span that needs to be released once block of code is completed
+     */
+    public static Releasable trace(ThreadPool threadPool, Tracer tracer, String name, Map<String, Object> attributes) {
+        if (tracer.isEnabled() == false) {
+            return () -> {};
+        }
+        var span = Span.create();
+        var ctx = threadPool.getThreadContext().newTraceContext();
+        tracer.startTrace(threadPool.getThreadContext(), span, name, attributes);
+
+        return () -> {
+            tracer.stopTrace(span);
+            ctx.restore();
+        };
+    }
+
+    private record Span(String spanId) implements Traceable {
+
+        @Override
+        public String getSpanId() {
+            return spanId;
+        }
+
+        public static Span create() {
+            return new Span(UUIDs.randomBase64UUID());
+        }
+    }
+}
diff --git a/server/src/main/java/org/elasticsearch/telemetry/tracing/Tracer.java b/server/src/main/java/org/elasticsearch/telemetry/tracing/Tracer.java
index 6f2c98dda4e2b..7d651d73776d8 100644
--- a/server/src/main/java/org/elasticsearch/telemetry/tracing/Tracer.java
+++ b/server/src/main/java/org/elasticsearch/telemetry/tracing/Tracer.java
@@ -32,6 +32,11 @@
  */
 public interface Tracer {
 
+    /**
+     * @return {@code if service is enabled}
+     */
+    boolean isEnabled();
+
     /**
      * Called when a span starts.
      * @param traceContext the current context. Required for tracing parent/child span activity.
@@ -144,6 +149,12 @@ public interface Tracer {
      * in order to avoid null checks everywhere.
      */
     Tracer NOOP = new Tracer() {
+
+        @Override
+        public boolean isEnabled() {
+            return false;
+        }
+
         @Override
         public void startTrace(TraceContext traceContext, Traceable traceable, String name, Map<String, Object> attributes) {}