清除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 CustomSchedulingConfiguration implements SchedulingConfigurer {

  @Override
  public void configureTasks(ScheduledTaskRegistrar taskRegistrar) {
    // 创建并初始化自定义的ThreadPoolTaskScheduler
    var customThreadPoolTaskScheduler = new CustomThreadPoolTaskScheduler();
    customThreadPoolTaskScheduler.initialize(); // 必须调用initialize()方法

    // 将自定义的TaskScheduler设置给任务注册器
    taskRegistrar.setTaskScheduler(customThreadPoolTaskScheduler);
  }
}

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.RejectedExecutionHandler;
import java.util.concurrent.RunnableScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;

public class CustomScheduledThreadPoolExecutor extends ScheduledThreadPoolExecutor {

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

  @Override
  protected <V> RunnableScheduledFuture<V> decorateTask(
      Callable<V> callable, RunnableScheduledFuture<V> task) {
    // 对于Callable类型的任务，也包装成CustomTask
    return new CustomTask<>(task);
  }

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

import com.demo.utils.GeneralUtils; // 假设这是你的工具类，包含clearContext()方法

import java.util.concurrent.Delayed;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.RunnableScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

// 使用Java 16+的record简化代码，如果使用旧版本Java，可定义为普通类
public record CustomTask<V>(RunnableScheduledFuture<V> task)
    implements RunnableScheduledFuture<V> {

  @Override
  public void run() {
    try {
      // TODO: 任务执行前的自定义逻辑，例如设置上下文、日志记录等
      // System.out.println("Scheduled task started: " + Thread.currentThread().getName());

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

    } finally {
      // TODO: 任务执行后的自定义逻辑，例如清理线程上下文
      // 假设GeneralUtils.clearContext()用于清理ThreadLocal等
      GeneralUtils.clearContext();
      // System.out.println("Scheduled task finished and context cleared: " + Thread.currentThread().getName());
    }
  }

  // 以下方法委托给原始任务执行，确保ScheduledExecutorService的正常行为
  @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;

import com.demo.utils.GeneralUtils; // 假设你有一个设置上下文的工具类

@Component
public class MyScheduledTasks {

    // 模拟一个ThreadLocal上下文
    private static final ThreadLocal<String> CONTEXT = new ThreadLocal<>();

    @Scheduled(fixedDelayString = "5000") // 每5秒执行一次
    public void doSomething() {
        // 在任务中设置上下文
        CONTEXT.set("TaskContext_" + System.currentTimeMillis());
        System.out.println("Task 'doSomething' executed. Context: " + CONTEXT.get());

        // 模拟一些工作
        try {
            Thread.sleep(1000);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        // 理论上，这里不需要手动清理，因为CustomTask的finally块会处理
        // 但为了演示，我们可以模拟在任务内部清理，不过通常不推荐
        // CONTEXT.remove();
    }

    @Scheduled(cron = "0 * * * * ?") // 每分钟执行一次
    public void anotherTask() {
        CONTEXT.set("AnotherTaskContext_" + System.currentTimeMillis());
        System.out.println("Task 'anotherTask' executed. Context: " + CONTEXT.get());
        // ...
    }
}