From 9571ff0f7bee8a07b5726c8a80dba047abc9df81 Mon Sep 17 00:00:00 2001
From: Nidhi <nidhi@appsmith.com>
Date: Thu, 4 Apr 2024 12:02:06 +0530
Subject: [PATCH] chore: Added trace and span info to context (#32399)

---
 .../appsmith/server/configurations/MDCConfig.java | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/app/server/appsmith-server/src/main/java/com/appsmith/server/configurations/MDCConfig.java b/app/server/appsmith-server/src/main/java/com/appsmith/server/configurations/MDCConfig.java
index 4d36213555..2cb991b090 100644
--- a/app/server/appsmith-server/src/main/java/com/appsmith/server/configurations/MDCConfig.java
+++ b/app/server/appsmith-server/src/main/java/com/appsmith/server/configurations/MDCConfig.java
@@ -2,6 +2,8 @@ package com.appsmith.server.configurations;
 
 import com.appsmith.server.filters.MDCFilter;
 import com.appsmith.server.helpers.LogHelper;
+import io.micrometer.observation.Observation;
+import io.micrometer.tracing.handler.TracingObservationHandler;
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
 import org.reactivestreams.Subscription;
@@ -13,11 +15,15 @@ import reactor.core.publisher.Operators;
 import reactor.util.context.Context;
 
 import java.util.Map;
+import java.util.Optional;
 
 @Configuration
 public class MDCConfig {
 
     private static final String MDC_CONTEXT_REACTOR_KEY = "MDCConfig";
+    private static final String OBSERVATION = "micrometer.observation";
+    private static final String TRACE_ID = "traceId";
+    private static final String SPAN_ID = "spanId";
 
     @PostConstruct
     void contextOperatorHook() {
@@ -78,6 +84,15 @@ public class MDCConfig {
                 Map<String, String> map = context.get(LogHelper.CONTEXT_MAP);
 
                 map.put(MDCFilter.THREAD, Thread.currentThread().getName());
+                Optional<Observation> observationOptional = context.getOrEmpty(OBSERVATION);
+                observationOptional.ifPresent(observation -> {
+                    TracingObservationHandler.TracingContext tracingContext =
+                            observation.getContext().get(TracingObservationHandler.TracingContext.class);
+                    if (tracingContext != null && tracingContext.getSpan() != null) {
+                        map.put(TRACE_ID, tracingContext.getSpan().context().traceId());
+                        map.put(SPAN_ID, tracingContext.getSpan().context().spanId());
+                    }
+                });
                 MDC.setContextMap(map);
             } else {
                 MDC.clear();