package com.microservices.common.core.threadPool;

import com.microservices.common.core.context.SecurityContextHolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.concurrent.*;

/**
 * @author OTTO
 */
public class ThreadPoolExecutorWrap extends ThreadPoolExecutor {
    private static final Logger logger = LoggerFactory.getLogger(ThreadPoolExecutorWrap.class);

    public ThreadPoolExecutorWrap(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    @Override
    public void execute(Runnable task) {
        if (task == null) {
            throw new NullPointerException();
        }
        super.execute(newTaskFor(new RunnableWrap(task), null));
    }

    static class RunnableWrap implements Runnable {
        private final Runnable task;
        private final Map<String, Object> contextValue;

        public RunnableWrap(Runnable task) {
            contextValue = SecurityContextHolder.getLocalMap();
            this.task = task;
        }

        @Override
        public void run() {
            try {
                SecurityContextHolder.setLocalMap(contextValue);
                // 用户任务逻辑
                task.run();
            } catch (Exception e) {
                logger.error("线程池任务执行失败：", e);
            }
        }
    }
}
