BFD-4683: Add MDC-aware thread pool executor (#3116)

Author: aschey-forpeople

apps/bfd-server-ng/README.md

@@ -82,10 +82,18 @@ mvn spring-boot:run
 
 ### Profiling SQL queries
 
-The `sqlprofile` profile can be used to log SQL queries and output metrics.
+The `sql-profile` profile can be used to log SQL queries and output metrics.
 
 ```sh
-mvn spring-boot run -Dspring-boot.run.profiles=local,sqlprofile
+mvn spring-boot run -Dspring-boot.run.profiles=local,sql-profile
+```
+
+### Structured Logging
+
+The `structured-log` profile can be used to test structured logs.
+
+```sh
+mvn spring-boot run -Dspring-boot.run.profiles=local,structured-log
 ```
 
 ## Swagger

apps/bfd-server-ng/src/main/java/gov/cms/bfd/server/ng/Application.java

@@ -16,6 +16,7 @@
 import org.springframework.boot.web.servlet.ServletRegistrationBean;
 import org.springframework.context.annotation.Bean;
 import org.springframework.scheduling.annotation.EnableAsync;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 import org.springframework.transaction.annotation.EnableTransactionManagement;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 
@@ -104,4 +105,14 @@ public AuditLogger auditLogger(Configuration configuration, ObjectMapper objectM
   public DynamoDbClient dynamoDbClient(Configuration configuration) {
     return configuration.getDynamoDbClient();
   }
+
+  /**
+   * Creates the custom task executor.
+   *
+   * @return task executor
+   */
+  @Bean
+  public ThreadPoolTaskExecutor taskExecutor() {
+    return new MdcAwareThreadPoolExecutor();
+  }
 }

apps/bfd-server-ng/src/main/java/gov/cms/bfd/server/ng/Configuration.java

@@ -38,9 +38,11 @@
 @ConfigurationProperties(prefix = "bfd")
 public class Configuration implements Serializable {
   // Unfortunately, constructor injection doesn't work with @ConfigurationProperties
+  private static final String LOCAL_ENV = "local";
 
   /** Identifies which Spring profiles indicate that the server is being run on a local machine. */
-  private static final List<String> ALLOWED_LOCAL_PROFILES = List.of("local", "sqlprofile");
+  private static final List<String> ALLOWED_LOCAL_PROFILES =
+      List.of(LOCAL_ENV, "sql-profile", "structured-log");
 
   // Getters should only be generated for configuration properties, not dependencies
   @Getter(value = AccessLevel.NONE)
@@ -57,12 +59,11 @@ public class Configuration implements Serializable {
   @Autowired(required = false)
   private JdbcConnectionDetails jdbcConnectionDetails;
 
-  private static final String BFD_ENV_LOCAL = "local";
   private String dynamoLocalUrl = "http://localhost:8000";
 
   // Default to local configuration, this should be overridden on deployment with the correct
   // environment.
-  private String env = BFD_ENV_LOCAL;
+  private String env = LOCAL_ENV;
   private String dbIdentifier = "";
   private Local local = new Local();
   private Sensitive sensitive = new Sensitive();
@@ -129,7 +130,7 @@ public AuditLogger getAuditLogger(ObjectMapper objectMapper) {
   }
 
   boolean isLocal() {
-    return env.equalsIgnoreCase(BFD_ENV_LOCAL);
+    return env.equalsIgnoreCase(LOCAL_ENV);
   }
 
   /** Represents possible types of audit logging. */

apps/bfd-server-ng/src/main/java/gov/cms/bfd/server/ng/MdcAwareThreadPoolExecutor.java

@@ -0,0 +1,31 @@
+package gov.cms.bfd.server.ng;
+
+import org.slf4j.MDC;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
+
+/**
+ * Custom task executor that preserves MDC info, ensuring all relevant info is logged for each
+ * request.
+ */
+public class MdcAwareThreadPoolExecutor extends ThreadPoolTaskExecutor {
+
+  static Runnable wrapWithMdcContext(Runnable task) {
+    // save the current MDC context
+    var contextMap = MDC.getCopyOfContextMap();
+    return () -> {
+      MDC.clear();
+      MDC.setContextMap(contextMap);
+      try {
+        task.run();
+      } finally {
+        // once the task is complete, clear MDC
+        MDC.clear();
+      }
+    };
+  }
+
+  @Override
+  public void execute(Runnable command) {
+    super.execute(wrapWithMdcContext(command));
+  }
+}

…ain/resources/application-aws.properties …esources/application-deployed.propertiesapps/bfd-server-ng/src/main/resources/application-aws.properties renamed to apps/bfd-server-ng/src/main/resources/application-deployed.properties apps/bfd-server-ng/src/main/resources/application-aws.properties renamed to apps/bfd-server-ng/src/main/resources/application-deployed.properties

@@ -1,6 +1,4 @@
 spring.config.import=aws-parameterstore:/bfd/${BFD_ENV}/server-ng?prefix=bfd.
 bfd.env=${BFD_ENV}
-# Use Elastic Common Schema for a sensible JSON-based log format
-logging.structured.format.console=ecs
 # Native support for Forwarded headers
 server.forward-headers-strategy=NATIVE

…ources/application-sqlprofile.properties …urces/application-sql-profile.propertiesapps/bfd-server-ng/src/main/resources/application-sqlprofile.properties renamed to apps/bfd-server-ng/src/main/resources/application-sql-profile.properties apps/bfd-server-ng/src/main/resources/application-sqlprofile.properties renamed to apps/bfd-server-ng/src/main/resources/application-sql-profile.properties

@@ -0,0 +1,2 @@
+# Use Elastic Common Schema for a sensible JSON-based log format
+logging.structured.format.console=ecs

apps/bfd-server-ng/src/main/resources/application-structured-log.properties

@@ -13,3 +13,4 @@ spring.mvc.favicon.enabled=false
 # This is still behind mtls though so it's fine to include here
 management.endpoints.web.exposure.include=metrics
 spring.threads.virtual.enabled=true
+spring.profiles.group.aws=structured-log,deployed

apps/bfd-server-ng/src/main/resources/application.properties

@@ -9,7 +9,7 @@
         <include resource="org/springframework/boot/logging/logback/console-appender.xml" />
         <logger name="org.springframework.web" level="DEBUG" />
     </springProfile>
-    <springProfile name="aws">
+    <springProfile name="structured-log">
         <!-- use structured logging in the deployed environment only -->
         <include resource="org/springframework/boot/logging/logback/structured-console-appender.xml" />
     </springProfile>