Add tracing support with OpenTelemetry

pom.xml

@@ -54,6 +54,11 @@
             <artifactId>oshi-core</artifactId>
             <version>6.9.2</version>
         </dependency>
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-sdk</artifactId>
+            <version>1.58.0</version>
+        </dependency>
         <dependency>
             <groupId>com.github.spotbugs</groupId>
             <artifactId>spotbugs-annotations</artifactId>

src/main/java/io/apitally/common/ApitallyClient.java

@@ -69,6 +69,7 @@ public enum HubRequestStatus {
 
     public final RequestCounter requestCounter;
     public final RequestLogger requestLogger;
+    public final SpanCollector spanCollector;
     public final ValidationErrorCounter validationErrorCounter;
     public final ServerErrorCounter serverErrorCounter;
     public final ConsumerRegistry consumerRegistry;
@@ -85,6 +86,10 @@ public ApitallyClient(String clientId, String env, RequestLoggingConfig requestL
 
         this.requestCounter = new RequestCounter();
         this.requestLogger = new RequestLogger(requestLoggingConfig);
+        this.spanCollector =
+                new SpanCollector(
+                        requestLoggingConfig.isEnabled()
+                                && requestLoggingConfig.isTracingEnabled());
         this.validationErrorCounter = new ValidationErrorCounter();
         this.serverErrorCounter = new ServerErrorCounter();
         this.consumerRegistry = new ConsumerRegistry();

src/main/java/io/apitally/common/ApitallyAppender.java → src/main/java/io/apitally/common/LogAppender.java

@@ -9,8 +9,8 @@
 import java.util.List;
 import org.slf4j.LoggerFactory;
 
-public class ApitallyAppender extends AppenderBase<ILoggingEvent> {
-    private static final String NAME = "ApitallyAppender";
+public class LogAppender extends AppenderBase<ILoggingEvent> {
+    private static final String NAME = "ApitallyLogAppender";
     private static final int MAX_BUFFER_SIZE = 1000;
     private static final int MAX_MESSAGE_LENGTH = 2048;
 
@@ -26,7 +26,7 @@ public static synchronized void register() {
             return;
         }
 
-        ApitallyAppender appender = new ApitallyAppender();
+        LogAppender appender = new LogAppender();
         appender.setContext(loggerContext);
         appender.setName(NAME);
         appender.start();

src/main/java/io/apitally/common/RequestLogger.java

@@ -9,6 +9,7 @@
 import io.apitally.common.dto.Request;
 import io.apitally.common.dto.RequestLogItem;
 import io.apitally.common.dto.Response;
+import io.apitally.common.dto.SpanData;
 import java.io.IOException;
 import java.net.MalformedURLException;
 import java.net.URL;
@@ -139,7 +140,12 @@ public void setSuspendUntil(long timestamp) {
     }
 
     public void logRequest(
-            Request request, Response response, Exception exception, List<LogRecord> logs) {
+            Request request,
+            Response response,
+            Exception exception,
+            List<LogRecord> logs,
+            List<SpanData> spans,
+            String traceId) {
         if (!enabled || suspendUntil != null && suspendUntil > System.currentTimeMillis()) {
             return;
         }
@@ -179,7 +185,13 @@ public void logRequest(
                 logs = null;
             }
 
-            RequestLogItem item = new RequestLogItem(request, response, exceptionDto, logs);
+            if (!config.isTracingEnabled()) {
+                spans = null;
+                traceId = null;
+            }
+
+            RequestLogItem item =
+                    new RequestLogItem(request, response, exceptionDto, logs, spans, traceId);
             pendingWrites.add(item);
 
             if (pendingWrites.size() > MAX_PENDING_WRITES) {
@@ -290,6 +302,12 @@ public void writeToFile() throws IOException {
                 if (item.getLogs() != null && !item.getLogs().isEmpty()) {
                     itemNode.set("logs", objectMapper.valueToTree(item.getLogs()));
                 }
+                if (item.getSpans() != null && !item.getSpans().isEmpty()) {
+                    itemNode.set("spans", objectMapper.valueToTree(item.getSpans()));
+                }
+                if (item.getTraceId() != null && !item.getTraceId().isEmpty()) {
+                    itemNode.put("trace_id", item.getTraceId());
+                }
 
                 String serializedItem = objectMapper.writeValueAsString(itemNode);
                 currentFile.writeLine(serializedItem.getBytes(StandardCharsets.UTF_8));

src/main/java/io/apitally/common/RequestLoggingConfig.java

@@ -12,6 +12,7 @@ public class RequestLoggingConfig {
     private boolean responseBodyIncluded = false;
     private boolean exceptionIncluded = true;
     private boolean logCaptureEnabled = false;
+    private boolean tracingEnabled = false;
     private List<String> queryParamMaskPatterns = new ArrayList<>();
     private List<String> headerMaskPatterns = new ArrayList<>();
     private List<String> bodyFieldMaskPatterns = new ArrayList<>();
@@ -82,6 +83,14 @@ public void setLogCaptureEnabled(boolean logCaptureEnabled) {
         this.logCaptureEnabled = logCaptureEnabled;
     }
 
+    public boolean isTracingEnabled() {
+        return tracingEnabled;
+    }
+
+    public void setTracingEnabled(boolean tracingEnabled) {
+        this.tracingEnabled = tracingEnabled;
+    }
+
     public List<String> getQueryParamMaskPatterns() {
         return queryParamMaskPatterns;
     }

src/main/java/io/apitally/common/SpanCollector.java

@@ -0,0 +1,180 @@
+package io.apitally.common;
+
+import io.apitally.common.dto.SpanData;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+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 io.opentelemetry.sdk.trace.ReadWriteSpan;
+import io.opentelemetry.sdk.trace.ReadableSpan;
+import io.opentelemetry.sdk.trace.SpanProcessor;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedQueue;
+
+public class SpanCollector implements SpanProcessor {
+    private final boolean enabled;
+    private volatile Tracer tracer;
+    private final Map<String, Set<String>> includedSpanIds = new ConcurrentHashMap<>();
+    private final Map<String, ConcurrentLinkedQueue<SpanData>> collectedSpans =
+            new ConcurrentHashMap<>();
+
+    public SpanCollector(boolean enabled) {
+        this.enabled = enabled;
+    }
+
+    public void setTracer(Tracer tracer) {
+        this.tracer = tracer;
+    }
+
+    public SpanHandle startCollection() {
+        if (!enabled || tracer == null) {
+            return null;
+        }
+
+        Span span = tracer.spanBuilder("root").setSpanKind(SpanKind.INTERNAL).startSpan();
+        Scope scope = Context.current().with(span).makeCurrent();
+        SpanContext spanContext = span.getSpanContext();
+        String traceId = spanContext.getTraceId();
+
+        Set<String> spanIds = ConcurrentHashMap.newKeySet();
+        spanIds.add(spanContext.getSpanId());
+        includedSpanIds.put(traceId, spanIds);
+        collectedSpans.put(traceId, new ConcurrentLinkedQueue<>());
+
+        return new SpanHandle(traceId, span, scope, this);
+    }
+
+    List<SpanData> getAndClearSpans(String traceId) {
+        if (traceId == null) {
+            return null;
+        }
+
+        includedSpanIds.remove(traceId);
+        ConcurrentLinkedQueue<SpanData> spans = collectedSpans.remove(traceId);
+        return spans != null ? new ArrayList<>(spans) : null;
+    }
+
+    @Override
+    public void onStart(Context parentContext, ReadWriteSpan span) {
+        if (!enabled) {
+            return;
+        }
+
+        SpanContext spanContext = span.getSpanContext();
+        String traceId = spanContext.getTraceId();
+        String spanId = spanContext.getSpanId();
+
+        Set<String> included = includedSpanIds.get(traceId);
+        if (included == null) {
+            return;
+        }
+
+        SpanContext parentSpanContext = span.getParentSpanContext();
+        if (parentSpanContext.isValid() && included.contains(parentSpanContext.getSpanId())) {
+            included.add(spanId);
+        }
+    }
+
+    @Override
+    public boolean isStartRequired() {
+        return enabled;
+    }
+
+    @Override
+    public void onEnd(ReadableSpan span) {
+        if (!enabled) {
+            return;
+        }
+
+        SpanContext spanContext = span.getSpanContext();
+        String traceId = spanContext.getTraceId();
+        String spanId = spanContext.getSpanId();
+
+        Set<String> included = includedSpanIds.get(traceId);
+        if (included == null || !included.contains(spanId)) {
+            return;
+        }
+
+        SpanData data = serializeSpan(span);
+        ConcurrentLinkedQueue<SpanData> spans = collectedSpans.get(traceId);
+        if (spans != null) {
+            spans.add(data);
+        }
+    }
+
+    @Override
+    public boolean isEndRequired() {
+        return enabled;
+    }
+
+    private SpanData serializeSpan(ReadableSpan span) {
+        io.opentelemetry.sdk.trace.data.SpanData spanData = span.toSpanData();
+        SpanContext spanContext = spanData.getSpanContext();
+        SpanContext parentSpanContext = spanData.getParentSpanContext();
+
+        String parentSpanId = parentSpanContext.isValid() ? parentSpanContext.getSpanId() : null;
+        String status =
+                spanData.getStatus().getStatusCode() != StatusCode.UNSET
+                        ? spanData.getStatus().getStatusCode().name()
+                        : null;
+
+        Map<String, Object> attributes = null;
+        if (!spanData.getAttributes().isEmpty()) {
+            Map<String, Object> attrMap = new HashMap<>();
+            spanData.getAttributes().forEach((key, value) -> attrMap.put(key.getKey(), value));
+            attributes = attrMap;
+        }
+
+        return new SpanData(
+                spanContext.getSpanId(),
+                parentSpanId,
+                spanData.getName(),
+                spanData.getKind().name(),
+                spanData.getStartEpochNanos(),
+                spanData.getEndEpochNanos(),
+                status,
+                attributes);
+    }
+
+    public static class SpanHandle {
+        private final String traceId;
+        private final Span span;
+        private final Scope scope;
+        private final SpanCollector collector;
+
+        SpanHandle(String traceId, Span span, Scope scope, SpanCollector collector) {
+            this.traceId = traceId;
+            this.span = span;
+            this.scope = scope;
+            this.collector = collector;
+        }
+
+        public String getTraceId() {
+            return traceId;
+        }
+
+        public void setName(String name) {
+            span.updateName(name);
+        }
+
+        public List<SpanData> end() {
+            scope.close();
+            span.end();
+            return collector.getAndClearSpans(traceId);
+        }
+    }
+
+    void resetForTest() {
+        this.tracer = null;
+        this.includedSpanIds.clear();
+        this.collectedSpans.clear();
+    }
+}

src/main/java/io/apitally/common/dto/RequestLogItem.java

@@ -10,14 +10,23 @@ public class RequestLogItem extends BaseDto {
     private final Response response;
     private final ExceptionDto exception;
     private final List<LogRecord> logs;
+    private final List<SpanData> spans;
+    private final String traceId;
 
     public RequestLogItem(
-            Request request, Response response, ExceptionDto exception, List<LogRecord> logs) {
+            Request request,
+            Response response,
+            ExceptionDto exception,
+            List<LogRecord> logs,
+            List<SpanData> spans,
+            String traceId) {
         this.uuid = UUID.randomUUID().toString();
         this.request = request;
         this.response = response;
         this.exception = exception;
         this.logs = logs;
+        this.spans = spans;
+        this.traceId = traceId;
     }
 
     @JsonProperty("uuid")
@@ -44,4 +53,14 @@ public ExceptionDto getException() {
     public List<LogRecord> getLogs() {
         return logs;
     }
+
+    @JsonProperty("spans")
+    public List<SpanData> getSpans() {
+        return spans;
+    }
+
+    @JsonProperty("trace_id")
+    public String getTraceId() {
+        return traceId;
+    }
 }