JDK21下ScopedValue实现StructuredTaskScope上下文传播及MDC适配问询
问题:ScopedValue在StructuredTaskScope的上下文传播与MDC适配
背景与需求
运行环境:JDK21.0.1-preview
需要在请求链路中传播tracer/span、请求上下文状态,必须覆盖三种线程切换场景:
- 虚拟线程→平台线程
- 平台线程→虚拟线程
- 虚拟线程→虚拟线程
同时要求兼容ThreadLocal相关方案(如context-propagation、TransmittableThreadLocal)
已实现基于ScopedValue的部分代码,但未覆盖平台线程场景,且遇到Log4j2 MDC适配ScopedValue的问题,现有代码如下:
@SuppressWarnings({"unchecked", "rawtypes"}) static class PropagatingTaskDecorator implements TaskDecorator { public record Tracer(String traceId, String spanId, String parentId) { } private final ScopedValue[] keys; public PropagatingTaskDecorator(ScopedValue<?>... keys) { this.keys = keys; } static class ScopedValueMap { public static final ScopedValue<Tracer> TRACER_SCOPED_VALUE = ScopedValue.newInstance(); private final Map<ScopedValue<?>, Object> bindingsSnapshot; public ScopedValueMap() { bindingsSnapshot = Maps.newHashMap(); if (TRACER_SCOPED_VALUE.isBound()) { var latestTracer = TRACER_SCOPED_VALUE.get(); bindingsSnapshot.put(TRACER_SCOPED_VALUE, new Tracer(latestTracer.traceId(), UUID.randomUUID().toString(), latestTracer.spanId())); } else { bindingsSnapshot.put(TRACER_SCOPED_VALUE, new Tracer(UUID.randomUUID().toString(), UUID.randomUUID().toString(), null)); } } /** * 捕获指定ScopedValue的当前绑定值 * @param scopedValue 目标ScopedValue * @param <T> 值类型 */ <T> void put(ScopedValue<T> scopedValue) { if (scopedValue.isBound()) { bindingsSnapshot.put(scopedValue, scopedValue.get()); } } <T> T get(ScopedValue<T> scopedValue) { if (bindingsSnapshot.containsKey(scopedValue)) { return (T) bindingsSnapshot.get(scopedValue); } throw new RuntimeException(STR."ScopedValue not found: \{scopedValue}"); } } @Nonnull @Override public Runnable decorate(@Nonnull Runnable runnable) { // 捕获当前ScopedValue绑定快照 var bindingsSnapshot = new ScopedValueMap(); for (var key : keys) { bindingsSnapshot.put(key); } return () -> { ScopedValue.Carrier carrier = ScopedValue.where(TRACER_SCOPED_VALUE, bindingsSnapshot.get(TRACER_SCOPED_VALUE)); for (var key : keys) { carrier = carrier.where(key, bindingsSnapshot.get(key)); } carrier.run(runnable); }; } } @Slf4j class ServiceImpl { public void xxx() { try (var executor = new SimpleAsyncTaskExecutor("pg-")) { executor.setVirtualThreads(true); var decorator = new PropagatingTaskDecorator(); executor.setTaskDecorator(decorator); executor.setThreadFactory(Thread.ofVirtual().name("v-").factory()); executor.submit(() -> { log.info("{}", xxx); // 需要Log4j2通过%X{traceId}打印追踪信息,但MDC无法读取ScopedValue中的数据 // 目前尝试重写ThreadContextMap,区分虚拟/平台线程处理,但未完全实现 }); } } }
解决方案
一、跨线程类型的ScopedValue上下文传播兼容
要覆盖三种线程场景,需同时处理ScopedValue与ThreadLocal的双向同步,核心逻辑是在任务装饰器中同时捕获当前线程的ScopedValue绑定和ThreadLocal状态,执行任务时根据目标线程类型选择对应恢复策略:
改进后的通用任务装饰器
@SuppressWarnings({"unchecked", "rawtypes"}) static class UniversalPropagatingDecorator implements TaskDecorator { public record Tracer(String traceId, String spanId, String parentId) {} public static final ScopedValue<Tracer> TRACER_SCOPED_VALUE = ScopedValue.newInstance(); // 兼容ThreadLocal场景,可替换为TransmittableThreadLocal或context-propagation容器 private static final ThreadLocal<Tracer> TRACER_THREAD_LOCAL = new ThreadLocal<>(); private final ScopedValue<?>[] scopedKeys; public UniversalPropagatingDecorator(ScopedValue<?>... scopedKeys) { this.scopedKeys = scopedKeys; } // 上下文快照:同时包含ScopedValue绑定和ThreadLocal状态 private record ContextSnapshot( Map<ScopedValue<?>, Object> scopedBindings, Tracer threadLocalTracer ) {} // 捕获当前线程的完整上下文 private ContextSnapshot captureCurrentContext() { // 捕获ScopedValue绑定 Map<ScopedValue<?>, Object> scopedBindings = new HashMap<>(); if (TRACER_SCOPED_VALUE.isBound()) { scopedBindings.put(TRACER_SCOPED_VALUE, TRACER_SCOPED_VALUE.get()); } for (ScopedValue<?> key : scopedKeys) { if (key.isBound()) { scopedBindings.put(key, key.get()); } } // 捕获ThreadLocal状态 Tracer threadLocalTracer = TRACER_THREAD_LOCAL.get(); return new ContextSnapshot(scopedBindings, threadLocalTracer); } @Nonnull @Override public Runnable decorate(@Nonnull Runnable runnable) { ContextSnapshot snapshot = captureCurrentContext(); return () -> { boolean isVirtualThread = Thread.currentThread().isVirtual(); if (isVirtualThread) { // 虚拟线程:优先用ScopedValue恢复上下文,同步到ThreadLocal兼容旧代码 ScopedValue.Carrier carrier = ScopedValue.newCarrier(); snapshot.scopedBindings.forEach(carrier::where); carrier.run(() -> { try { if (snapshot.threadLocalTracer != null) { TRACER_THREAD_LOCAL.set(snapshot.threadLocalTracer); } runnable.run(); } finally { TRACER_THREAD_LOCAL.remove(); } }); } else { // 平台线程:优先恢复ThreadLocal,同步到ScopedValue兼容新代码 Tracer originalTracer = TRACER_THREAD_LOCAL.get(); try { if (snapshot.threadLocalTracer != null) { TRACER_THREAD_LOCAL.set(snapshot.threadLocalTracer); } ScopedValue.Carrier carrier = ScopedValue.newCarrier(); snapshot.scopedBindings.forEach(carrier::where); carrier.run(runnable); } finally { // 恢复原ThreadLocal状态 if (originalTracer != null) { TRACER_THREAD_LOCAL.set(originalTracer); } else { TRACER_THREAD_LOCAL.remove(); } } } }; } }
关键逻辑说明
- 双快照捕获:同时抓取ScopedValue绑定和ThreadLocal状态,确保新旧上下文方案无遗漏
- 线程类型判断:通过
Thread.currentThread().isVirtual()区分目标线程类型,执行对应恢复逻辑 - 双向同步:
- 虚拟线程:先通过ScopedValue.Carrier恢复上下文,再同步到ThreadLocal兼容依赖ThreadLocal的旧代码
- 平台线程:先恢复ThreadLocal状态,再绑定ScopedValue兼容依赖ScopedValue的新代码
- 资源清理:任务执行完成后恢复或清理ThreadLocal,避免内存泄漏
二、Log4j2 MDC适配ScopedValue
要让Log4j2的%X{traceId}能读取ScopedValue中的上下文,需自定义ThreadContextMap实现,优先从ScopedValue读取,降级到ThreadLocal:
自定义ScopedValueThreadContextMap
import org.apache.logging.log4j.ThreadContext; import org.apache.logging.log4j.spi.ThreadContextMap; import org.apache.logging.log4j.util.StringMap; import java.util.HashMap; import java.util.Map; import java.util.Set; public class ScopedValueThreadContextMap implements ThreadContextMap { private static final ScopedValue<Map<String, String>> MDC_SCOPED_VALUE = ScopedValue.newInstance(); // 平台线程的MDC备份 private static final ThreadLocal<Map<String, String>> MDC_THREAD_LOCAL = ThreadLocal.withInitial(HashMap::new); @Override public void put(String key, String value) { if (Thread.currentThread().isVirtual()) { Map<String, String> mdcMap = MDC_SCOPED_VALUE.isBound() ? new HashMap<>(MDC_SCOPED_VALUE.get()) : new HashMap<>(); mdcMap.put(key, value); // ScopedValue不可变,需重新绑定新的Map ScopedValue.where(MDC_SCOPED_VALUE, mdcMap).run(() -> {}); } else { MDC_THREAD_LOCAL.get().put(key, value); } } @Override public String get(String key) { if (Thread.currentThread().isVirtual()) { return MDC_SCOPED_VALUE.isBound() ? MDC_SCOPED_VALUE.get().get(key) : null; } else { return MDC_THREAD_LOCAL.get().get(key); } } @Override public void remove(String key) { if (Thread.currentThread().isVirtual()) { if (MDC_SCOPED_VALUE.isBound()) { Map<String, String> mdcMap = new HashMap<>(MDC_SCOPED_VALUE.get()); mdcMap.remove(key); ScopedValue.where(MDC_SCOPED_VALUE, mdcMap).run(() -> {}); } } else { MDC_THREAD_LOCAL.get().remove(key); } } @Override public void clear() { if (Thread.currentThread().isVirtual()) { ScopedValue.where(MDC_SCOPED_VALUE, new HashMap<>()).run(() -> {}); } else { MDC_THREAD_LOCAL.get().clear(); } } @Override public boolean containsKey(String key) { if (Thread.currentThread().isVirtual()) { return MDC_SCOPED_VALUE.isBound() && MDC_SCOPED_VALUE.get().containsKey(key); } else { return MDC_THREAD_LOCAL.get().containsKey(key); } } @Override public Map<String, String> getCopy() { if (Thread.currentThread().isVirtual()) { return MDC_SCOPED_VALUE.isBound() ? new HashMap<>(MDC_SCOPED_VALUE.get()) : new HashMap<>(); } else { return new HashMap<>(MDC_THREAD_LOCAL.get()); } } @Override public Map<String, String> getImmutableMapOrNull() { return getCopy(); } @Override public boolean isEmpty() { if (Thread.currentThread().isVirtual()) { return !MDC_SCOPED_VALUE.isBound() || MDC_SCOPED_VALUE.get().isEmpty(); } else { return MDC_THREAD_LOCAL.get().isEmpty(); } } @Override public StringMap getStringMap() { throw new UnsupportedOperationException("getStringMap未实现,可使用getCopy替代"); } @Override public void putAll(Map<String, String> map) { if (Thread.currentThread().isVirtual()) { Map<String, String> mdcMap = MDC_SCOPED_VALUE.isBound() ? new HashMap<>(MDC_SCOPED_VALUE.get()) : new HashMap<>(); mdcMap.putAll(map); ScopedValue.where(MDC_SCOPED_VALUE, mdcMap).run(() -> {}); } else { MDC_THREAD_LOCAL.get().putAll(map); } } @Override public Set<String> keySet() { if (Thread.currentThread().isVirtual()) { return MDC_SCOPED_VALUE.isBound() ? MDC_SCOPED_VALUE.get().keySet() : Set.of(); } else { return MDC_THREAD_LOCAL.get().keySet(); } } }
配置Log4j2使用自定义实现
在log4j2.properties中指定自定义ThreadContextMap:
log4j2.threadContextMap=com.yourpackage.ScopedValueThreadContextMap
与上下文传播的联动
在UniversalPropagatingDecorator中加入MDC快照捕获与恢复逻辑:
// 更新ContextSnapshot,加入MDC快照 private record ContextSnapshot( Map<ScopedValue<?>, Object> scopedBindings, Tracer threadLocalTracer, Map<String, String> mdcSnapshot ) {} // 捕获MDC快照 private ContextSnapshot captureCurrentContext() { // ... 原有代码 Map<String, String> mdcSnapshot = ThreadContext.getCopy(); return new ContextSnapshot(scopedBindings, threadLocalTracer, mdcSnapshot); } // 虚拟线程场景恢复MDC carrier.run(() -> { try { if (snapshot.threadLocalTracer != null) { TRACER_THREAD_LOCAL.set(snapshot.threadLocalTracer); } ThreadContext.putAll(snapshot.mdcSnapshot); runnable.run(); } finally { ThreadContext.clear(); TRACER_THREAD_LOCAL.remove(); } }); // 平台线程场景恢复MDC try { if (snapshot.threadLocalTracer != null) { TRACER_THREAD_LOCAL.set(snapshot.threadLocalTracer); } ThreadContext.putAll(snapshot.mdcSnapshot); ScopedValue.Carrier carrier = ScopedValue.newCarrier(); snapshot.scopedBindings.forEach(carrier::where); carrier.run(runnable); } finally { ThreadContext.clear(); if (originalTracer != null) { TRACER_THREAD_LOCAL.set(originalTracer); } else { TRACER_THREAD_LOCAL.remove(); } }
三、StructuredTaskScope中的集成
StructuredTaskScope会自动继承父线程的ScopedValue绑定,无需额外装饰器,但如果需要兼容ThreadLocal,可在提交任务时手动同步上下文:
try (var scope = new StructuredTaskScope.ShutdownOnFailure()) { ContextSnapshot snapshot = captureCurrentContext(); scope.fork(() -> { // 恢复ThreadLocal与MDC TRACER_THREAD_LOCAL.set(snapshot.threadLocalTracer); ThreadContext.putAll(snapshot.mdcSnapshot); try { // 执行业务逻辑 return doBusinessLogic(); } finally { // 清理资源 TRACER_THREAD_LOCAL.remove(); ThreadContext.clear(); } }); scope.join(); }
内容的提问来源于stack exchange,提问作者magical1989
相关产品推荐
相关产品推荐

