How to transfer traceId of asynchronous thread

foreword

In daily troubleshooting, we often query the log link of the request in ELK according to the traceId. In the synchronous request, it is very cool to stop at the end according to the traceId, but what should we do if it is an asynchronous request? The asynchronous requests in the project are combined with the thread pool to start the asynchronous thread. The following combines the MDC and the thread pool in slf4j to realize the traceId transfer of the asynchronous thread.

Rewrite the method in ThreadPoolTaskExecutor

The following tool classes set the request mapping context through MDC.setContextMap(context) before the execution of Callable and Runnable asynchronous tasks

import org.slf4j.MDC;
import org.springframework.util.CollectionUtils;

import java.util.Map;
import java.util.concurrent.Callable;

/**
 * @desc: 定义MDC工具类，支持Runnable和Callable两种，目的就是为了把父线程的traceId设置给子线程
 */
public class MdcUtil {

    public static <T> Callable<T> wrap(final Callable<T> callable, final Map<String, String> context) {
        return () -> {
            if (CollectionUtils.isEmpty(context)) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            try {
                return callable.call();
            } finally {
                // 清除子线程的，避免内存溢出，就和ThreadLocal.remove()一个原因
                MDC.clear();
            }
        };
    }

    public static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
        return () -> {
            if (CollectionUtils.isEmpty(context)) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            try {
                runnable.run();
            } finally {
                MDC.clear();
            }
        };
    }
}

 The following defines a ThreadPoolMdcExecutor class to inherit the ThreadPoolTaskExecutor class and rewrite the execute and submit methods

import java.util.concurrent.Callable;
import java.util.concurrent.Future;

/**
 * @desc: 把当前的traceId透传到子线程特意加的实现。
 *   重点就是 MDC.getCopyOfContextMap()，此方法获取当前线程（父线程）的traceId
 */
public class ThreadPoolMdcExecutor extends ThreadPoolTaskExecutor {
    @Override
    public void execute(Runnable task) {
        super.execute(MdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public Future<?> submit(Runnable task) {
        return super.submit(MdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return super.submit(MdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }
}

Define the thread pool below, you can use ThreadPoolMdcExecutor

    @Bean(name = "callBackExecutorConfig")
    public Executor callBackExecutorConfig() {
        ThreadPoolTaskExecutor executor = new ThreadPoolMdcExecutor();
        // 配置核心线程数
        executor.setCorePoolSize(10);
        // 配置最大线程数
        executor.setMaxPoolSize(20);
        // 配置队列大小
        executor.setQueueCapacity(200);
        // 配置线程池中的线程的名称前缀
        executor.setThreadNamePrefix("async-Thread-");
        // rejection-policy：当pool已经达到max size的时候，如何处理新任务
        // abort：在调用executor执行的方法中抛出异常 RejectedExecutionException
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
        // 执行初始化
        executor.initialize();
        return executor;
    }

After defining the thread pool, we can use the callBackExecutorConfig thread pool to perform asynchronous tasks to avoid the loss of traceId in the asynchronous thread.

Thread pool enhancement

The above is to rewrite the execute and submit methods by inheriting ThreadPoolTaskExecutor, and set MDC.setContextMap(context) to set the context. We can also enhance the thread pool by implementing the TaskDecorator interface

public class TaskDecoratorForMdc implements TaskDecorator {
    @Override
    public Runnable decorate(Runnable runnable) {
        try {
            Optional<Map<String, String>> contextMapOptional =ofNullable(MDC.getCopyOfContextMap());
            return () -> {
                try {
                    contextMapOptional.ifPresent(MDC::setContextMap);
                    runnable.run();
                } finally {
                    MDC.clear();
                }
            };
        } catch (Exception e) {
            return runnable;
        }
    }
}

Next, define the thread pool and enhance the thread pool

    @Bean(name = "callBackExecutorConfig")
    public Executor callBackExecutorConfig() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor ();
        // 配置核心线程数
        executor.setCorePoolSize(10);
        // 配置最大线程数
        executor.setMaxPoolSize(20);
        // 配置队列大小
        executor.setQueueCapacity(200);
        // 配置线程池中的线程的名称前缀
        executor.setThreadNamePrefix("async-Thread-");
        // rejection-policy：当pool已经达到max size的时候，如何处理新任务
        // abort：在调用executor执行的方法中抛出异常 RejectedExecutionException
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
        //线程池增强
        threadPoolTaskExecutor.setTaskDecorator(new TaskDecoratorForMdc());
        // 执行初始化
        executor.initialize();
        return executor;
    }

Summarize

The above two methods actually use Mdc to synchronize the traceId between asynchronous threads. You can look at the source code of Mdc. In the end, the child thread obtains the parent thread information through InheritableThreadLocal

public class BasicMDCAdapter implements MDCAdapter {
    private InheritableThreadLocal<Map<String, String>> inheritableThreadLocal = 
        new InheritableThreadLocal<Map<String, String>>() {
        protected Map<String, String> childValue(Map<String, String> parentValue) {
            return parentValue == null ? null : new HashMap(parentValue);
        }
    };
    
    //省略若干
    ......
}