From 78fd452b0405440fbe806feaedb6bfc1aeac40d5 Mon Sep 17 00:00:00 2001
From: Greg Meldrum <gregmeldrum@gmail.com>
Date: Thu, 7 Dec 2023 13:35:51 -0500
Subject: [PATCH] Move MDC to setOutputListener (#139)

---
 .../terraform/client/ProcessLauncher.java     | 13 +++++++++++--
 .../terraform/manager/TerraformManager.java   | 19 ++-----------------
 2 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/service/terraform-plugin/src/main/java/com/solace/maas/ep/event/management/agent/plugin/terraform/client/ProcessLauncher.java b/service/terraform-plugin/src/main/java/com/solace/maas/ep/event/management/agent/plugin/terraform/client/ProcessLauncher.java
index d1b1d5e2..ca9e34b8 100644
--- a/service/terraform-plugin/src/main/java/com/solace/maas/ep/event/management/agent/plugin/terraform/client/ProcessLauncher.java
+++ b/service/terraform-plugin/src/main/java/com/solace/maas/ep/event/management/agent/plugin/terraform/client/ProcessLauncher.java
@@ -33,6 +33,7 @@
 
 import com.solace.maas.ep.event.management.agent.plugin.jacoco.ExcludeFromJacocoGeneratedReport;
 import lombok.extern.slf4j.Slf4j;
+import org.slf4j.MDC;
 
 import java.io.BufferedReader;
 import java.io.File;
@@ -64,13 +65,21 @@ final class ProcessLauncher {
     }
 
     void setOutputListener(Consumer<String> listener) {
+        Map<String, String> mdcContext = MDC.getCopyOfContextMap();
         assert process == null;
-        outputListener = listener;
+        outputListener = (log) -> {
+            MDC.setContextMap(mdcContext);
+            listener.accept(log);
+        };
     }
 
     void setErrorListener(Consumer<String> listener) {
+        Map<String, String> mdcContext = MDC.getCopyOfContextMap();
         assert process == null;
-        errorListener = listener;
+        errorListener = (error) -> {
+            MDC.setContextMap(mdcContext);
+            listener.accept(error);
+        };
     }
 
     void setInheritIO(boolean inheritIO) {
diff --git a/service/terraform-plugin/src/main/java/com/solace/maas/ep/event/management/agent/plugin/terraform/manager/TerraformManager.java b/service/terraform-plugin/src/main/java/com/solace/maas/ep/event/management/agent/plugin/terraform/manager/TerraformManager.java
index 9bed967c..53c6c3fc 100644
--- a/service/terraform-plugin/src/main/java/com/solace/maas/ep/event/management/agent/plugin/terraform/manager/TerraformManager.java
+++ b/service/terraform-plugin/src/main/java/com/solace/maas/ep/event/management/agent/plugin/terraform/manager/TerraformManager.java
@@ -45,12 +45,6 @@ public void execute(CommandRequest request, Command command, Map<String, String>
 
         MDC.put(RouteConstants.COMMAND_CORRELATION_ID, request.getCorrelationId());
         MDC.put(RouteConstants.MESSAGING_SERVICE_ID, request.getServiceId());
-        String traceId = MDC.get(RouteConstants.TRACE_ID);
-        String spanId = MDC.get("spanId");
-        String commandCorrelationId = MDC.get(RouteConstants.COMMAND_CORRELATION_ID);
-        String actorId = MDC.get(RouteConstants.ACTOR_ID);
-        String scheduleId = MDC.get(RouteConstants.SCHEDULE_ID);
-        String messagingServiceId = MDC.get(RouteConstants.MESSAGING_SERVICE_ID);
 
         log.debug("Executing command {} for serviceId {} correlationId {} context {}", command.getCommand(), request.getServiceId(),
                 request.getCorrelationId(), request.getContext());
@@ -58,8 +52,7 @@ public void execute(CommandRequest request, Command command, Map<String, String>
         try (TerraformClient terraformClient = terraformClientFactory.createClient()) {
 
             Path configPath = createConfigPath(request);
-            List<String> logOutput = setupTerraformClient(terraformClient, configPath, traceId, spanId, commandCorrelationId,
-                    actorId, scheduleId, messagingServiceId);
+            List<String> logOutput = setupTerraformClient(terraformClient, configPath);
             String commandVerb = executeTerraformCommand(command, envVars, configPath, terraformClient);
             processTerraformResponse(request, command, commandVerb, logOutput);
         } catch (InterruptedException e) {
@@ -71,21 +64,13 @@ public void execute(CommandRequest request, Command command, Map<String, String>
         }
     }
 
-    private static List<String> setupTerraformClient(TerraformClient terraformClient, Path configPath,
-                                                     String traceId, String spanId, String commandCorrelationId, String actorId,
-                                                     String scheduleId, String messagingServiceId) {
+    private static List<String> setupTerraformClient(TerraformClient terraformClient, Path configPath) {
         terraformClient.setWorkingDirectory(configPath.toFile());
         List<String> output = new ArrayList<>();
 
         // Write each terraform to the output list so that it can be processed later
         // Also write the output to the main log to be streamed back to EP
         terraformClient.setOutputListener(tfLog -> {
-            MDC.put(RouteConstants.TRACE_ID, traceId);
-            MDC.put("spanId", spanId);
-            MDC.put(RouteConstants.COMMAND_CORRELATION_ID, commandCorrelationId);
-            MDC.put(RouteConstants.ACTOR_ID, actorId);
-            MDC.put(RouteConstants.SCHEDULE_ID, scheduleId);
-            MDC.put(RouteConstants.MESSAGING_SERVICE_ID, messagingServiceId);
             output.add(tfLog);
             log.debug("Terraform output: {}", tfLog);
         });