定制Spring @Scheduled任务以实现线程上下文清理

import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.SchedulingConfigurer;
import org.springframework.scheduling.config.ScheduledTaskRegistrar;

@Configuration
@EnableScheduling // 启用Spring的调度功能
public class SchedulingConfiguration implements SchedulingConfigurer {

  @Override
  public void configureTasks(ScheduledTaskRegistrar taskRegistrar) {
    // 创建并初始化我们自定义的线程池任务调度器
    CustomThreadPoolTaskScheduler threadPoolTaskScheduler = new CustomThreadPoolTaskScheduler();
    threadPoolTaskScheduler.setPoolSize(5); // 设置线程池大小，可根据需求调整
    threadPoolTaskScheduler.setThreadNamePrefix("MyScheduledTask-"); // 设置线程名前缀
    threadPoolTaskScheduler.initialize(); // 必须调用initialize方法来初始化线程池

    // 将自定义的调度器设置给ScheduledTaskRegistrar
    taskRegistrar.setTaskScheduler(threadPoolTaskScheduler);
  }
}

import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;

public class CustomThreadPoolTaskScheduler extends ThreadPoolTaskScheduler {

  @Override
  protected ScheduledExecutorService createExecutor(
      int poolSize,
      ThreadFactory threadFactory,
      RejectedExecutionHandler rejectedExecutionHandler) {
    // 返回我们自定义的ScheduledThreadPoolExecutor
    return new CustomScheduledThreadPoolExecutor(poolSize, threadFactory, rejectedExecutionHandler);
  }
}

import java.util.concurrent.Callable;
import java.util.concurrent.Delayed;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.RunnableScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

public class CustomScheduledThreadPoolExecutor extends ScheduledThreadPoolExecutor {

  public CustomScheduledThreadPoolExecutor(
      int corePoolSize,
      ThreadFactory threadFactory,
      RejectedExecutionHandler handler) {
    super(corePoolSize, threadFactory, handler);
  }

  // 装饰Runnable类型的任务
  @Override
  protected <V> RunnableScheduledFuture<V> decorateTask(
      Runnable runnable, RunnableScheduledFuture<V> task) {
    return new CustomTask<>(task);
  }

  // 装饰Callable类型的任务
  @Override
  protected <V> RunnableScheduledFuture<V> decorateTask(
      Callable<V> callable, RunnableScheduledFuture<V> task) {
    return new CustomTask<>(task);
  }

  // 4. 自定义任务包装器 CustomTask
  private record CustomTask<V>(RunnableScheduledFuture<V> task)
      implements RunnableScheduledFuture<V> {

    @Override
    public void run() {
      try {
        // 在任务执行前可以添加自定义逻辑，例如设置MDC、初始化一些上下文
        // System.out.println("Before task execution: " + Thread.currentThread().getName());

        task.run(); // 执行原始任务

      } finally {
        // 在任务执行后清理线程上下文
        // 这是一个示例，需要替换为实际的清理逻辑
        // 例如：GeneralUtils.clearContext();
        // 如果使用了MDC，可以调用 MDC.clear();
        // 如果使用了ThreadLocal，需要手动移除或设置为null
        System.out.println("After task execution, clearing context for: " + Thread.currentThread().getName());
        // 假设有一个通用的上下文清理工具类
        // GeneralUtils.clearContext(); // 替换为你的实际清理方法
      }
    }

    // 以下方法都是简单地委托给原始任务，以保持其原有功能
    @Override
    public boolean cancel(boolean mayInterruptIfRunning) {
      return task.cancel(mayInterruptIfRunning);
    }

    @Override
    public boolean isCancelled() {
      return task.isCancelled();
    }

    @Override
    public boolean isDone() {
      return task.isDone();
    }

    @Override
    public V get() throws InterruptedException, ExecutionException {
      return task.get();
    }

    @Override
    public V get(long timeout, TimeUnit unit)
        throws InterruptedException, ExecutionException, TimeoutException {
      return task.get(timeout, unit);
    }

    @Override
    public long getDelay(TimeUnit unit) {
      return task.getDelay(unit);
    }

    @Override
    public int compareTo(Delayed o) {
      return task.compareTo(o);
    }

    @Override
    public boolean isPeriodic() {
      return task.isPeriodic();
    }
  }
}

import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

@Component
public class MyScheduledTasks {

    @Scheduled(fixedDelayString = "10000") // 每10秒执行一次
    public void doSomething() {
        System.out.println("Executing scheduled task: " + Thread.currentThread().getName());
        // 模拟任务中设置线程上下文
        // GeneralUtils.setContext("some_value");
        // System.out.println("Context set to: " + GeneralUtils.getContext());
    }
}