From 4671770740b462ff010e8d53f6d386630f135476 Mon Sep 17 00:00:00 2001
From: Mona Mohamadinia <mona.mohamadinia@gmail.com>
Date: Thu, 2 Jul 2020 22:43:50 +0430
Subject: [PATCH] BAEL-4217: Added Thread Pool Aware MDC (#9554)

* Added Thread Pool Aware MDC

* Fixed Thread Pool
---
 .../java/com/baeldung/mdc/TransferDemo.java   | 18 ++++++-----
 .../mdc/pool/MdcAwareThreadPoolExecutor.java  | 31 +++++++++++++++++++
 .../com/baeldung/mdc/slf4j/Slf4jRunnable.java |  2 +-
 3 files changed, 42 insertions(+), 9 deletions(-)
 create mode 100644 logging-modules/log-mdc/src/main/java/com/baeldung/mdc/pool/MdcAwareThreadPoolExecutor.java

diff --git a/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/TransferDemo.java b/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/TransferDemo.java
index 259e9a8c5c..0cc9dee2b9 100644
--- a/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/TransferDemo.java
+++ b/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/TransferDemo.java
@@ -1,19 +1,21 @@
 package com.baeldung.mdc;
 
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-
-import org.apache.log4j.Logger;
-
-import com.baeldung.mdc.log4j.Log4JRunnable;
-import com.baeldung.mdc.log4j2.Log4J2Runnable;
+import com.baeldung.mdc.pool.MdcAwareThreadPoolExecutor;
 import com.baeldung.mdc.slf4j.Slf4jRunnable;
 
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor.AbortPolicy;
+
+import static java.util.concurrent.TimeUnit.MINUTES;
+
 public class TransferDemo {
 
     public static void main(String[] args) {
 
-        ExecutorService executor = Executors.newFixedThreadPool(3);
+        ExecutorService executor = new MdcAwareThreadPoolExecutor(3, 3, 0, MINUTES,
+          new LinkedBlockingQueue<>(), Thread::new, new AbortPolicy());
+
         TransactionFactory transactionFactory = new TransactionFactory();
 
         for (int i = 0; i < 10; i++) {
diff --git a/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/pool/MdcAwareThreadPoolExecutor.java b/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/pool/MdcAwareThreadPoolExecutor.java
new file mode 100644
index 0000000000..71463fe116
--- /dev/null
+++ b/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/pool/MdcAwareThreadPoolExecutor.java
@@ -0,0 +1,31 @@
+package com.baeldung.mdc.pool;
+
+import org.apache.logging.log4j.ThreadContext;
+import org.slf4j.MDC;
+
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.RejectedExecutionHandler;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+
+public class MdcAwareThreadPoolExecutor extends ThreadPoolExecutor {
+
+    public MdcAwareThreadPoolExecutor(int corePoolSize,
+      int maximumPoolSize,
+      long keepAliveTime,
+      TimeUnit unit,
+      BlockingQueue<Runnable> workQueue,
+      ThreadFactory threadFactory,
+      RejectedExecutionHandler handler) {
+        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
+    }
+
+    @Override
+    protected void afterExecute(Runnable r, Throwable t) {
+        System.out.println("Cleaning the MDC context");
+        MDC.clear();
+        org.apache.log4j.MDC.clear();
+        ThreadContext.clearAll();
+    }
+}
diff --git a/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/slf4j/Slf4jRunnable.java b/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/slf4j/Slf4jRunnable.java
index e30a28a3c7..f49640c527 100644
--- a/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/slf4j/Slf4jRunnable.java
+++ b/logging-modules/log-mdc/src/main/java/com/baeldung/mdc/slf4j/Slf4jRunnable.java
@@ -18,7 +18,7 @@ public class Slf4jRunnable implements Runnable {
 
         new Slf4TransferService().transfer(tx.getAmount());
 
-        MDC.clear();
+        // MDC.clear(); We don't need this with MdcAwareThreadPoolExecutor
 
     }
 }
\ No newline at end of file