MDC+VirtualThread

CustomVirtualThreadTaskExecutor

@Slf4j
public class CustomVirtualThreadTaskExecutor /*extends TaskExecutorAdapter*/ implements ExecutorService {
    private final ExecutorService exsInner = Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("virtual-async#", 1).factory());
//    public CustomVirtualThreadTaskExecutor() {
//        super(Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("virtual-async#", 1).factory()));
//    }

    @Override
    public void execute(Runnable task) {
        exsInner.execute(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public void shutdown() {
        exsInner.shutdown();
    }

    @Override
    public List<Runnable> shutdownNow() {
        log.warn("CustomVirtualThreadTaskExecutor.shutdownNow() is called");
        return exsInner.shutdownNow();
    }

    @Override
    public boolean isShutdown() {
        return exsInner.isShutdown();
    }

    @Override
    public boolean isTerminated() {
        return exsInner.isTerminated();
    }

    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        return exsInner.awaitTermination(timeout, unit);
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return exsInner.submit(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
        return exsInner.submit(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()), result);
    }

    @Override
    public Future<?> submit(Runnable task) {
        return exsInner.submit(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) throws InterruptedException {
        return exsInner.invokeAll(ThreadMdcUtil.wrap(tasks, MDC.getCopyOfContextMap()));
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException {
        return exsInner.invokeAll(ThreadMdcUtil.wrap(tasks, MDC.getCopyOfContextMap()), timeout, unit);
    }

    @Override
    public <T> T invokeAny(Collection<? extends Callable<T>> tasks) throws InterruptedException, ExecutionException {
        return exsInner.invokeAny(ThreadMdcUtil.wrap(tasks, MDC.getCopyOfContextMap()));
    }

    @Override
    public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        return exsInner.invokeAny(ThreadMdcUtil.wrap(tasks, MDC.getCopyOfContextMap()), timeout, unit);
    }

    @Override
    public void close() {
        log.warn("CustomVirtualThreadTaskExecutor.close() is called");
        exsInner.close();
    }


}

ThreadMdcUtil

@Slf4j
public class ThreadMdcUtil {
    public static <T> Callable<T> wrap(final Callable<T> callable, final Map<String, String> context) {
        return () -> {
            setContext(context);
            try {
                return callable.call();
            } catch (Throwable e) {
                saveErrorLog(e);

                throw e;
            } finally {
                MDC.clear();
//                TtlThreadLocalContext.remove();
            }
        };
    }

    public static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
        return () -> {
            setContext(context);
            //setTraceIdIfAbsent();
            try {
                runnable.run();
            } catch (Throwable e) {
                saveErrorLog(e);

                throw e;
            } finally {
                MDC.clear();
//                TtlThreadLocalContext.remove();
            }
        };
    }

    public static <T> Collection<? extends Callable<T>> wrap(Collection<? extends Callable<T>> tasks, final Map<String, String> context) {
         return tasks.stream().map(tCallable -> wrap(tCallable, context)).collect(Collectors.toList());
    }

    private static void setContext(Map<String, String> context) {
        if (CollectionUtils.isEmpty(context)) {
            MDC.clear();
        } else {
            MDC.setContextMap(context);
        }
    }

    private static void saveErrorLog(Throwable e) {
//        ICustomLogService operateLogService = SpringUtil.getBean(ICustomLogService.class);
//        String err = ExceptionInfoHandler.compileError(e);
//        operateLogService.asyncSaveOperErrorData("异步线程内异常", err);
    }

}