You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

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();
                    }
                }
            }
        };
    }
}

关键逻辑说明

  1. 双快照捕获:同时抓取ScopedValue绑定和ThreadLocal状态,确保新旧上下文方案无遗漏
  2. 线程类型判断:通过Thread.currentThread().isVirtual()区分目标线程类型,执行对应恢复逻辑
  3. 双向同步:
    • 虚拟线程:先通过ScopedValue.Carrier恢复上下文,再同步到ThreadLocal兼容依赖ThreadLocal的旧代码
    • 平台线程:先恢复ThreadLocal状态,再绑定ScopedValue兼容依赖ScopedValue的新代码
  4. 资源清理:任务执行完成后恢复或清理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

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.07.03 07:13:10