package cn.reghao.tnb.user.app.hystrix;

import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategy;
import lombok.extern.slf4j.Slf4j;

import java.lang.reflect.Field;
import java.util.concurrent.Callable;

/**
 * 将当前线程的 threadLocal 变量传给 callable
 * callable 在执行时会使用传入的调用者线程的 threadLocal 变量, 然后再执行任务
 * callable 执行前后需要保存 worker 线程自身的线程局部变量以供恢复
 *
 * @author reghao
 * @date 2025-07-22 16:04:38
 */
@Slf4j
public class MyHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy {
    @Override
    public <T> Callable<T> wrapCallable(Callable<T> callable) {
        // 获取当前线程的 threadLocalMap
        Object currentThreadLocalMap = getCurrentThreadLocalMap();

        Callable<T> finalCallable = new Callable<T>() {
            private Object callerThreadLocalMap = currentThreadLocalMap;
            private Callable<T> targetCallable = callable;

            @Override
            public T call() throws Exception {
                // 保存工作线程现有的线程变量
                Object oldThreadLocalMapOfWorkThread = getCurrentThreadLocalMap();
                // 将本线程的线程变量设置为 caller 的线程变量
                setCurrentThreadLocalMap(callerThreadLocalMap);
                try {
                    return targetCallable.call();
                } finally {
                    setCurrentThreadLocalMap(oldThreadLocalMapOfWorkThread);
                    log.info("restore work thread's threadlocal");
                }
            }
        };

        return finalCallable;
    }

    private Object getCurrentThreadLocalMap() {
        Thread thread = Thread.currentThread();
        try {
            Field field = Thread.class.getDeclaredField("threadLocals");
            field.setAccessible(true);
            Object o = field.get(thread);
            return o;
        } catch (NoSuchFieldException | IllegalAccessException e) {
            log.error("{}", e);
        }
        return null;
    }

    private void setCurrentThreadLocalMap(Object newThreadLocalMap) {
        Thread thread = Thread.currentThread();
        try {
            Field field = Thread.class.getDeclaredField("threadLocals");
            field.setAccessible(true);
            field.set(thread,newThreadLocalMap);

        } catch (NoSuchFieldException | IllegalAccessException e) {
            log.error("{}", e);
        }
    }
}
