Spring DeferredResult：如何在响应HTTP线程中设置MDC上下文？

问题背景 #

在异步HTTP场景中涉及三类线程：

• 请求线程：在请求过滤器中设置MDC
• 异步线程：已实现从请求线程复制线程上下文
• 响应线程：需要实现从异步线程复制MDC上下文，但现有DeferredResultProcessingInterceptor类的所有方法要么在响应线程启动前执行，要么在响应完全执行后执行，无法在响应线程启动时注入上下文。

相关代码如下：

@RestController
public class ServiceChangeStatusPollingApi {

    private final StatusChangePolling statusChangePolling;

    private final Executor taskExecutor;

    @Autowired
    public ServiceChangeStatusPollingApi(StatusChangePolling statusChangePolling,
                                         @Qualifier(TaskConfiguration.ASYNC_EXECUTOR_NAME)
                                         Executor taskExecutor) {
        this.statusChangePolling = statusChangePolling;
        this.taskExecutor = taskExecutor;
    }

    @GetMapping(value = "/status/polling",
            produces = MediaType.APPLICATION_JSON_VALUE)
    @ResponseBody
    public DeferredResult<ServiceDeploymentState> getLatestServiceDeploymentStatus(
            @Parameter(name = "id", description = "ID of the service")
            @RequestParam(name = "id") UUID serviceId,
            @Parameter(name = "previousServiceDeploymentState", description = "Previous known service status to client")
            @RequestParam(name = "previousServiceDeploymentState", required = false)
            ServiceDeploymentState previousServiceDeploymentState
    ) {
        DeferredResult<ServiceDeploymentState> stateDeferredResult = new DeferredResult<>();
        taskExecutor.execute(() ->
                this.statusChangePolling.waitForUsersWithStatus(
                        stateDeferredResult,
                        serviceId,
                        previousServiceDeploymentState));
        return stateDeferredResult;
    }

}

解决方案 #

通过包装DeferredResult的回调逻辑，在异步线程设置结果时捕获MDC上下文，再在响应线程处理结果时恢复该上下文，具体实现如下：

1. 自定义MDC感知的DeferredResult #

创建MdcAwareDeferredResult类，重写结果设置方法，捕获异步线程的MDC上下文并在响应回调中恢复：

public class MdcAwareDeferredResult<T> extends DeferredResult<T> {
    @Override
    public void setResult(T result) {
        // 保存异步线程的MDC上下文
        Map<String, String> mdcContext = MDC.getCopyOfContextMap();
        super.setResult(result);
        // 注册响应完成回调，恢复MDC到响应线程
        this.onCompletion(() -> {
            if (mdcContext != null) {
                MDC.setContextMap(mdcContext);
            } else {
                MDC.clear();
            }
        });
    }

    @Override
    public void setErrorResult(Object result) {
        Map<String, String> mdcContext = MDC.getCopyOfContextMap();
        super.setErrorResult(result);
        this.onCompletion(() -> {
            if (mdcContext != null) {
                MDC.setContextMap(mdcContext);
            } else {
                MDC.clear();
            }
        });
    }
}

2. 替换Controller中的DeferredResult实现 #

在接口方法中使用自定义的MdcAwareDeferredResult替代原生类：

@GetMapping(value = "/status/polling",
        produces = MediaType.APPLICATION_JSON_VALUE)
@ResponseBody
public DeferredResult<ServiceDeploymentState> getLatestServiceDeploymentStatus(
        @RequestParam(name = "id") UUID serviceId,
        @RequestParam(name = "previousServiceDeploymentState", required = false)
        ServiceDeploymentState previousServiceDeploymentState
) {
    MdcAwareDeferredResult<ServiceDeploymentState> stateDeferredResult = new MdcAwareDeferredResult<>();
    taskExecutor.execute(() ->
            this.statusChangePolling.waitForUsersWithStatus(
                    stateDeferredResult,
                    serviceId,
                    previousServiceDeploymentState));
    return stateDeferredResult;
}

3. 确保异步线程池的MDC复制（若未实现） #

如果异步线程池还未处理MDC复制，可自定义包装类实现请求线程到异步线程的上下文传递：

public class MdcAwareExecutor implements Executor {
    private final Executor delegate;

    public MdcAwareExecutor(Executor delegate) {
        this.delegate = delegate;
    }

    @Override
    public void execute(Runnable command) {
        Map<String, String> mdcContext = MDC.getCopyOfContextMap();
        delegate.execute(() -> {
            try {
                if (mdcContext != null) {
                    MDC.setContextMap(mdcContext);
                }
                command.run();
            } finally {
                MDC.clear();
            }
        });
    }
}

在配置类中包装原异步线程池：

@Configuration
public class TaskConfiguration {
    public static final String ASYNC_EXECUTOR_NAME = "asyncExecutor";

    @Bean(name = ASYNC_EXECUTOR_NAME)
    public Executor asyncExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        executor.setCorePoolSize(5);
        executor.setMaxPoolSize(10);
        executor.setQueueCapacity(20);
        executor.setThreadNamePrefix("Async-");
        executor.initialize();
        return new MdcAwareExecutor(executor);
    }
}

原理说明 #

• 异步线程执行时，MdcAwareExecutor会复制请求线程的MDC上下文并在异步任务期间生效。
• 当异步线程调用setResult/setErrorResult时，MdcAwareDeferredResult捕获当前异步线程的MDC上下文，通过onCompletion注册的回调会在响应线程处理结果时执行，从而将MDC上下文恢复到响应线程中。

内容的提问来源于stack exchange，提问作者Swaroop Raghupathy