package com.sckw.log.TraceLog;

import com.sckw.core.model.constant.StringConstant;
import com.sckw.core.utils.UUIDUtils;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

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

/**
 * @author lfdc
 * @description
 * @date 2023-11-21 16:11:05
 */
@Slf4j
public class MdcTaskExecutor extends ThreadPoolTaskExecutor {

    private static final long serialVersionUID = 832576625683993818L;

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        log.info("mdc thread pool task executor submit");
        Map<String, String> context = MDC.getCopyOfContextMap();
        return super.submit(() -> {
            T result;
            if (context != null) {
                //将父线程的MDC内容传给子线程
                MDC.setContextMap(context);
            } else {
                //直接给子线程设置MDC
                MDC.put(StringConstant.TRACE_ID, UUIDUtils.get32UUID() + StringConstant.RUNG + StringConstant.SUB_THREAD);
            }
            try {
                //执行任务
                result = task.call();
            } finally {
                try {
                    MDC.clear();
                } catch (Exception e) {
                    log.warn("MDC clear exception", e);
                }
            }
            return result;
        });
    }

    @Override
    public void execute(Runnable task) {
        log.info("mdc thread pool task executor execute");
        Map<String, String> context = MDC.getCopyOfContextMap();
        super.execute(() -> {
            if (context != null) {
                //将父线程的MDC内容传给子线程
                MDC.setContextMap(context);
            } else {
                //直接给子线程设置MDC
                MDC.put(StringConstant.TRACE_ID, UUIDUtils.get32UUID() + StringConstant.RUNG + StringConstant.SUB_THREAD);
            }
            try {
                //执行任务
                task.run();
            } finally {
                try {
                    MDC.clear();
                } catch (Exception e) {
                    log.warn("MDC clear exception", e);
                }
            }
        });
    }
}
