package com.zhgd.xmgl.config;

import com.alibaba.fastjson.JSONObject;
import com.zhgd.xmgl.util.RequestIdUtil;
import com.zhgd.xmgl.util.ThreadLocalUtil;
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;

@Slf4j
public class MdcThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {

    /**
     * 接口请求开启的异步线程会调用下述方法
     */
    @Override
    public void execute(Runnable task) {
        Map<String, String> context = MDC.getCopyOfContextMap();    //复制主线程MDC
        JSONObject tlJo = ThreadLocalUtil.get();
        log.info("MDC异步多线程...");
        super.execute(() -> {
            if (null != context) {
                MDC.setContextMap(context);     //主线程MDC赋予子线程
            } else {
                RequestIdUtil.setRequestId();   //主线程没有MDC就自己生成一个
            }
            ThreadLocalUtil.set(tlJo);
            try {
                task.run();
            } finally {
                try {
                    RequestIdUtil.clear();
                    ThreadLocalUtil.remove();
                } catch (Exception e) {
                    log.warn("MDC clear exception：{}", e.getMessage());
                }
            }
        });
    }

    /**
     * 定时任务会调用下述方法
     */
    @Override
    public <T> Future<T> submit(Callable<T> task) {
        Map<String, String> context = MDC.getCopyOfContextMap();    //复制主线程MDC
        JSONObject tlJo = ThreadLocalUtil.get();
        log.info("MDC异步多线程...");
        return super.submit(() -> {
            if (null != context) {
                MDC.setContextMap(context);     //主线程MDC赋予子线程
            } else {
                RequestIdUtil.setRequestId();   //主线程没有MDC就自己生成一个
            }
            ThreadLocalUtil.set(tlJo);
            try {
                return task.call();
            } finally {
                try {
                    RequestIdUtil.clear();
                    ThreadLocalUtil.remove();
                } catch (Exception e) {
                    log.warn("MDC clear exception：{}", e.getMessage());
                }
            }
        });
    }
}