Java 多线程上下文传递在复杂场景下的实践

TransmittableThreadLocal（TTL）这个时候就派上用场了。这是阿里开源的一个组件，我们来看看它怎么解决线程池的问题，先来一段代码，在上面的基础上修改一下，使用TransmittableThreadLocal。

static TransmittableThreadLocal<String> transmittableThreadLocal = new TransmittableThreadLocal<>();// 使用TransmittableThreadLocal


public static void main(String[] args) throws InterruptedException {

    ExecutorService executorService = Executors.newFixedThreadPool(1);
    executorService = TtlExecutors.getTtlExecutorService(executorService); // 用TtlExecutors装饰线程池

    transmittableThreadLocal.set("i am a transmittable parent");
    executorService.execute(new Runnable() {
        @Override
        public void run() {

            System.out.println(transmittableThreadLocal.get());
            transmittableThreadLocal.set("i am a old transmittable parent");// 子线程设置新的值

        }
    });
    System.out.println(transmittableThreadLocal.get());

    TimeUnit.SECONDS.sleep(1);
    transmittableThreadLocal.set("i am a new transmittable parent");// 主线程设置新的值

    executorService.execute(new Runnable() {
        @Override
        public void run() {

            System.out.println(transmittableThreadLocal.get());
        }
    });
}


i am a transmittable parent
i am a transmittable parent
i am a new transmittable parent

private TtlRunnable(Runnable runnable, boolean releaseTtlValueReferenceAfterRun) {
    this.capturedRef = new AtomicReference<Object>(capture());
    this.runnable = runnable;
    this.releaseTtlValueReferenceAfterRun = releaseTtlValueReferenceAfterRun;
}

com.alibaba.ttl.TtlRunnable#run
/**
 * wrap method {@link Runnable#run()}.
 */
@Override
public void run() {
    Object captured = capturedRef.get();// 获取线程的ThreadLocalMap
    if (captured == null || releaseTtlValueReferenceAfterRun && !capturedRef.compareAndSet(captured, null)) {
        throw new IllegalStateException("TTL value reference is released after run!");
    }

    Object backup = replay(captured);// 暂存当前子线程的ThreadLocalMap到backup
    try {
        runnable.run();
    } finally {
        restore(backup);// 恢复线程执行时被改版的Threadlocal对应的值
    }
}


com.alibaba.ttl.TransmittableThreadLocal.Transmitter#replay


/**
 * Replay the captured {@link TransmittableThreadLocal} values from {@link #capture()},
 * and return the backup {@link TransmittableThreadLocal} values in current thread before replay.
 *
 * @param captured captured {@link TransmittableThreadLocal} values from other thread from {@link #capture()}
 * @return the backup {@link TransmittableThreadLocal} values before replay
 * @see #capture()
 * @since 2.3.0
 */
public static Object replay(Object captured) {
    @SuppressWarnings("unchecked")
    Map<TransmittableThreadLocal<?>, Object> capturedMap = (Map<TransmittableThreadLocal<?>, Object>) captured;
    Map<TransmittableThreadLocal<?>, Object> backup = new HashMap<TransmittableThreadLocal<?>, Object>();

    for (Iterator<? extends Map.Entry<TransmittableThreadLocal<?>, ?>> iterator = holder.get().entrySet().iterator();
         iterator.hasNext(); ) {
        Map.Entry<TransmittableThreadLocal<?>, ?> next = iterator.next();
        TransmittableThreadLocal<?> threadLocal = next.getKey();

        // backup
        backup.put(threadLocal, threadLocal.get());

        // clear the TTL value only in captured
        // avoid extra TTL value in captured, when run task.
        if (!capturedMap.containsKey(threadLocal)) {
            iterator.remove();
            threadLocal.superRemove();
        }
    }

    // set value to captured TTL
    for (Map.Entry<TransmittableThreadLocal<?>, Object> entry : capturedMap.entrySet()) {
        @SuppressWarnings("unchecked")
        TransmittableThreadLocal<Object> threadLocal = (TransmittableThreadLocal<Object>) entry.getKey();
        threadLocal.set(entry.getValue());
    }

    // call beforeExecute callback
    doExecuteCallback(true);

    return backup;
}



com.alibaba.ttl.TransmittableThreadLocal.Transmitter#restore

/**
 * Restore the backup {@link TransmittableThreadLocal} values from {@link Transmitter#replay(Object)}.
 *
 * @param backup the backup {@link TransmittableThreadLocal} values from {@link Transmitter#replay(Object)}
 * @since 2.3.0
 */
public static void restore(Object backup) {
    @SuppressWarnings("unchecked")
    Map<TransmittableThreadLocal<?>, Object> backupMap = (Map<TransmittableThreadLocal<?>, Object>) backup;
    // call afterExecute callback
    doExecuteCallback(false);

    for (Iterator<? extends Map.Entry<TransmittableThreadLocal<?>, ?>> iterator = holder.get().entrySet().iterator();
         iterator.hasNext(); ) {
        Map.Entry<TransmittableThreadLocal<?>, ?> next = iterator.next();
        TransmittableThreadLocal<?> threadLocal = next.getKey();

        // clear the TTL value only in backup
        // avoid the extra value of backup after restore
        if (!backupMap.containsKey(threadLocal)) {
            iterator.remove();
            threadLocal.superRemove();
        }
    }

    // restore TTL value
    for (Map.Entry<TransmittableThreadLocal<?>, Object> entry : backupMap.entrySet()) {
        @SuppressWarnings("unchecked")
        TransmittableThreadLocal<Object> threadLocal = (TransmittableThreadLocal<Object>) entry.getKey();
        threadLocal.set(entry.getValue());
    }
}

可以看下整个过程的完整时序图：

• HTTP 请求

public class ShopShardingHelperUtil {

    private static TransmittableThreadLocal<String> countrySet = new TransmittableThreadLocal<>();

    /**
     * 获取threadLocal中设置的国家标志
     * @return
     */
    public static String getCountry() {
        return countrySet.get();
    }

    /**
     * 设置threadLocal中设置的国家
     */
    public static void setCountry (String country) {
        countrySet.set(country.toLowerCase());
    }


    /**
     * 清除标志
     */
    public static void clear () {
        countrySet.remove();
    }
}



/** 拦截器对cookie和url综合判断国家信息，放入到TransmittableThreadLocal中 **/
// 设置线程中的国家标志
String country = localeContext.getLocale().getCountry().toLowerCase();

ShopShardingHelperUtil.setCountry(country);


/** 自定义线程池，用TtlExecutors包装原有自定义线程池 **/
public static Executor getExecutor() {

    if (executor == null) {
        synchronized (TransmittableExecutor.class) {
            if (executor == null) {
                executor = TtlExecutors.getTtlExecutor(initExecutor());// 用TtlExecutors装饰Executor，结合TransmittableThreadLocal解决异步线程threadlocal传递问题
            }
        }
    }

    return executor;
}


/** 实际使用线程池的地方,直接调用执行即可**/
TransmittableExecutor.getExecutor().execute(new BatchExeRunnable(param1,param2));



/** mybatis的Interceptor代码, 使用TransmittableThreadLocal的国家信息，改造原有sql，加上国家参数，在增删改查sql中区分国家数据 **/
public Object intercept(Invocation invocation) throws Throwable {

    StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
    BoundSql boundSql = statementHandler.getBoundSql();

    String originalSql = boundSql.getSql();

    Statement statement = (Statement) CCJSqlParserUtil.parse(originalSql);

    String threadCountry = ShopShardingHelperUtil.getCountry();

    // 线程中的国家不为空才进行处理
    if (StringUtils.isNotBlank(threadCountry)) {

        if (statement instanceof Select) {

            Select selectStatement = (Select) statement;
            VivoSelectVisitor vivoSelectVisitor = new VivoSelectVisitor(threadCountry);
            vivoSelectVisitor.init(selectStatement);
        } else if (statement instanceof Insert) {

            Insert insertStatement = (Insert) statement;
            VivoInsertVisitor vivoInsertVisitor = new VivoInsertVisitor(threadCountry);
            vivoInsertVisitor.init(insertStatement);

        } else if (statement instanceof Update) {

            Update updateStatement = (Update) statement;
            VivoUpdateVisitor vivoUpdateVisitor = new VivoUpdateVisitor(threadCountry);
            vivoUpdateVisitor.init(updateStatement);

        } else if (statement instanceof Delete) {

            Delete deleteStatement = (Delete) statement;
            VivoDeleteVisitor vivoDeleteVisitor = new VivoDeleteVisitor(threadCountry);
            vivoDeleteVisitor.init(deleteStatement);
        }


        Field boundSqlField = BoundSql.class.getDeclaredField("sql");
        boundSqlField.setAccessible(true);
        boundSqlField.set(boundSql, statement.toString());
    } else {

        logger.error("----------- intercept not-add-country sql.... ---------" + statement.toString());
    }

    logger.info("----------- intercept query new sql.... ---------" + statement.toString());
    // 调用方法，实际上就是拦截的方法
    Object result = invocation.proceed();

    return result;
}

对于 Dubbo 接口和无法判断国家信息的 HTTP 接口，在入参部分增加国家信息参数，通过拦截器或者手动set国家信息到TransmittableThreadLocal。

@Bean
public ThreadPoolTaskExecutor threadPoolTaskExecutor(){
    return TtlThreadPoolExecutors.getAsyncExecutor();
}


public class TtlThreadPoolExecutors {

    private static final String COMMON_BUSINESS = "COMMON_EXECUTOR";

    public static final int QUEUE_CAPACITY = 20000;

    public static ExecutorService getExecutorService() {
        return TtlExecutorServiceMananger.getExecutorService(COMMON_BUSINESS);
    }

    public static ExecutorService getExecutorService(String threadGroupName) {
        return TtlExecutorServiceMananger.getExecutorService(threadGroupName);
    }

    public static ThreadPoolTaskExecutor getAsyncExecutor() {
        // 用TtlExecutors装饰Executor，结合TransmittableThreadLocal解决异步线程threadlocal传递问题
        return getTtlThreadPoolTaskExecutor(initTaskExecutor());
    }

    private static ThreadPoolTaskExecutor initTaskExecutor () {
        return initTaskExecutor(TtlThreadPoolFactory.DEFAULT_CORE_SIZE, TtlThreadPoolFactory.DEFAULT_POOL_SIZE, QUEUE_CAPACITY);
    }

    private static ThreadPoolTaskExecutor initTaskExecutor (int coreSize, int poolSize, int executorQueueCapacity) {
        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
        taskExecutor.setCorePoolSize(coreSize);
        taskExecutor.setMaxPoolSize(poolSize);
        taskExecutor.setQueueCapacity(executorQueueCapacity);
        taskExecutor.setKeepAliveSeconds(120);
        taskExecutor.setAllowCoreThreadTimeOut(true);
        taskExecutor.setThreadNamePrefix("TaskExecutor-ttl");
        taskExecutor.initialize();
        return taskExecutor;
    }

    private static ThreadPoolTaskExecutor getTtlThreadPoolTaskExecutor(ThreadPoolTaskExecutor executor) {
        if (null == executor || executor instanceof ThreadPoolTaskExecutorWrapper) {
            return executor;
        }
        return new ThreadPoolTaskExecutorWrapper(executor);
    }
}




/**
 * @ClassName : LocaleContextHolder
 * @Description : 本地化信息上下文holder
 */
public class LocalizationContextHolder {
    private static TransmittableThreadLocal<LocalizationContext> localizationContextHolder = new TransmittableThreadLocal<>();
    private static LocalizationInfo defaultLocalizationInfo = new LocalizationInfo();

    private LocalizationContextHolder(){}

    public static LocalizationContext getLocalizationContext() {
        return localizationContextHolder.get();
    }

    public static void resetLocalizationContext () {
        localizationContextHolder.remove();
    }

    public static void setLocalizationContext (LocalizationContext localizationContext) {
        if(localizationContext == null) {
            resetLocalizationContext();
        } else {
            localizationContextHolder.set(localizationContext);
        }
    }

    public static void setLocalizationInfo (LocalizationInfo localizationInfo) {
        LocalizationContext localizationContext = getLocalizationContext();
        String brand = (localizationContext instanceof BrandLocalizationContext ?
                ((BrandLocalizationContext) localizationContext).getBrand() : null);
        if(StringUtils.isNotEmpty(brand)) {
            localizationContext = new SimpleBrandLocalizationContext(localizationInfo, brand);
        } else if(localizationInfo != null) {
            localizationContext = new SimpleLocalizationContext(localizationInfo);
        } else {
            localizationContext = null;
        }
        setLocalizationContext(localizationContext);
    }

    public static void setDefaultLocalizationInfo(@Nullable LocalizationInfo localizationInfo) {
        LocalizationContextHolder.defaultLocalizationInfo = localizationInfo;
    }

    public static LocalizationInfo getLocalizationInfo () {
        LocalizationContext localizationContext = getLocalizationContext();
        if(localizationContext != null) {
            LocalizationInfo localizationInfo = localizationContext.getLocalizationInfo();
            if(localizationInfo != null) {
                return localizationInfo;
            }
        }
        return defaultLocalizationInfo;
    }

    public static String getCountry(){
        return getLocalizationInfo().getCountry();
    }

    public static String getTimezone(){
        return getLocalizationInfo().getTimezone();
    }

    public static String getBrand(){
        return getBrand(getLocalizationContext());
    }

    public static String getBrand(LocalizationContext localizationContext) {
        if(localizationContext == null) {
            return null;
        }
        if(localizationContext instanceof BrandLocalizationContext) {
            return ((BrandLocalizationContext) localizationContext).getBrand();
        }
        throw new LocaleException("unsupported localizationContext type");
    }
}
    @Override
    public LocaleContext resolveLocaleContext(final HttpServletRequest request) {
        parseLocaleCookieIfNecessary(request);
        LocaleContext localeContext = new TimeZoneAwareLocaleContext() {
            @Override
            public Locale getLocale() {
                return (Locale) request.getAttribute(LOCALE_REQUEST_ATTRIBUTE_NAME);
            }
            @Override
            public TimeZone getTimeZone() {
                return (TimeZone) request.getAttribute(TIME_ZONE_REQUEST_ATTRIBUTE_NAME);
            }
        };
        // 设置线程中的国家标志
        setLocalizationInfo(request, localeContext.getLocale());
        return localeContext;
    }

    private void setLocalizationInfo(HttpServletRequest request, Locale locale) {
        String country = locale!=null?locale.getCountry():null;
        String language = locale!=null?(locale.getLanguage() + "_" + locale.getVariant()):null;
        LocaleRequestMessage localeRequestMessage = localeRequestParser.parse(request);
        final String countryStr = country;
        final String languageStr = language;
        final String brandStr = localeRequestMessage.getBrand();
        LocalizationContextHolder.setLocalizationContext(new BrandLocalizationContext() {
            @Override
            public String getBrand() {
                return brandStr;
            }

            @Override
            public LocalizationInfo getLocalizationInfo() {
                return LocalizationInfoAssembler.assemble(countryStr, languageStr);
            }
        });
    }

对于定时任务job，因为所有国家都需要执行，所以会把所有国家进行遍历执行，这也可以通过简单的注解和AOP来解决。

本文从业务拓展的角度阐述了在复杂业务场景下如何通过ThreadLocal，过渡到InheritableThreadLocal，再通过TransmittableThreadLocal解决实际业务问题。因为海外的业务在不断的探索中前进，技术也在不断的探索中演进，面对这种复杂多变的情况，我们的应对策略是先做国际化，再做本地化，more global才能more local，多国家的隔离只是国际化最基本的起点，未来还有很多业务和技术等着我们去挑战。