Spring Boot开启虚拟线程ScopedValue上下文传递的使用方式

spring:
  main:
    # 保证 JVM 在全是虚拟线程情况下不会提前退出
    keep-alive: true
  # 全局虚拟线程开关（推荐方式）
  threads:
    virtual:
      # 启用虚拟线程，覆盖 TaskExecutor、@Async、@Scheduled、Web Server
      enabled: true

import lombok.Builder;
/**
 * 虚拟线程上下文传递参数
 *
 * @param traceId  链路ID（分布式微服务传递追踪）
 * @param userId   用户ID
 * @param tenantId 租户ID
 */
@Builder
public record RequestContext(
    String traceId,
    String userId,
    String tenantId) {
}

import lombok.NoArgsConstructor;
/**
 * ScopedValue工具类
 */
@NoArgsConstructor
public final class ContextKeys {
    // 链路ID
    public static final String TRACE_ID = "traceId";
    /**
     * WEB请求上下文传递
     */
    public static final ScopedValue<RequestContext> REQUEST_CONTEXT = ScopedValue.newInstance();
}

import lombok.NoArgsConstructor;
import org.slf4j.MDC;
import java.util.Optional;
import java.util.concurrent.Callable;
/**
 * 获取上下文业务参数
 */
@NoArgsConstructor
public final class RequestContextHolder {
    /**
     * 获取完整上下文
     */
    public static Optional<RequestContext> getOptional() {
        return ContextKeys.REQUEST_CONTEXT.isBound() ? Optional.of(ContextKeys.REQUEST_CONTEXT.get()) : Optional.empty();
    }
    /**
     * 获取 traceId
     */
    public static String getTraceId() {
        return getOptional().map(RequestContext::traceId).orElse(null);
    }
    /**
     * 获取 userId
     */
    public static String getUserId() {
        return getOptional().map(RequestContext::userId).orElse(null);
    }
    /**
     * 获取 tenantId
     */
    public static String getTenantId() {
        return getOptional().map(RequestContext::tenantId).orElse(null);
    }
    /**
     * 绑定上下文并运行 Runnable
     */
    public static void with(RequestContext ctx, Runnable task) {
        ScopedValue.where(ContextKeys.REQUEST_CONTEXT, ctx).run(() -> {
            try {
                // MDC桥接
                MDC.put(ContextKeys.TRACE_ID, ctx.traceId());
                task.run();
            } finally {
                MDC.clear();
            }
        });
    }
    /**
     * 绑定上下文并运行 Callable
     */
    public static <T> T with(RequestContext ctx, Callable<T> task) throws Exception {
        return ScopedValue.where(ContextKeys.REQUEST_CONTEXT, ctx).call(() -> {
            try {
                // MDC桥接
                MDC.put(ContextKeys.TRACE_ID, ctx.traceId());
                return task.call();
            } finally {
                MDC.clear();
            }
        });
    }
}

import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.jetbrains.annotations.NotNull;
import org.slf4j.MDC;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;
import java.io.IOException;
import java.util.Optional;
import java.util.UUID;
/**
 * HTTP请求上下文初始化
 */
@Component
@Order(Ordered.HIGHEST_PRECEDENCE)
public class ContextInitFilter extends OncePerRequestFilter {
    @Override
    protected void doFilterInternal(@NotNull HttpServletRequest request, @NotNull HttpServletResponse response, @NotNull FilterChain filterChain) throws ServletException, IOException {
        RequestContext ctx = buildContext(request);
        // ScopedValue 绑定上下文
        try {
            ScopedValue.where(ContextKeys.REQUEST_CONTEXT, ctx).run(() -> {
                try {
                    // MDC桥接
                    MDC.put(ContextKeys.TRACE_ID, ctx.traceId());
                    filterChain.doFilter(request, response);
                } catch (IOException | ServletException e) {
                    throw new RuntimeException(e);
                } finally {
                    MDC.clear();
                }
            });
        } catch (RuntimeException e) {
            // 拆包，保持Servlet语义
            if (e.getCause() instanceof IOException io) throw io;
            if (e.getCause() instanceof ServletException se) throw se;
            throw e;
        }
    }
    /**
     * 构建请求上下文
     */
    private RequestContext buildContext(HttpServletRequest request) {
        // 链路ID
        String traceId = Optional.ofNullable(request.getHeader("X-Trace-Id")).orElse(UUID.randomUUID().toString());
        // 用户ID
        String userId = request.getHeader("X-User-Id");
        // 租户ID
        String tenantId = request.getHeader("X-Tenant-Id");
        return RequestContext.builder().traceId(traceId).userId(userId).tenantId(tenantId).build();
    }
}

import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@RestController
public class DemoController {
    @GetMapping("/test")
    public String test() {
        // 获取参数
        String tenantId = RequestContextHolder.getTenantId();
        String userId = RequestContextHolder.getUserId();
        // 新虚拟线程执行
        RequestContext ctx = RequestContextHolder.getOptional().orElseThrow();
        ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor();
        executor.submit(() -> {
            RequestContextHolder.with(ctx, () -> {
                // 异步任务中依然可以获取 traceId / userId
                System.out.println("traceId=" + RequestContextHolder.getTraceId());
            });
        });
        return "tenantId=" + tenantId + ", userId=" + userId;
    }
    /**
     * StructuredTaskScope实现：同步写法 + 并发执行 + 自动失败传播
     * 非常类似WebFlux/Reactor的：Mono.zip(callA(), callB()).map(tuple -> combine(tuple.getT1(), tuple.getT2()));
     */
    @Transactional
    public void service() {
        try (var scope = new StructuredTaskScope.ShutdownOnFailure()) {
            var a = scope.fork(this::taskA);
            var b = scope.fork(this::taskB);
            // 等待所有
            scope.join();
            // 有失败就抛
            scope.throwIfFailed();
            return combine(a.get(), b.get());
        }
    }
}