Spring @Scheduled 任务线程上下文清理的定制化方案

package com.example.config;

import com.example.scheduler.CustomThreadPoolTaskScheduler;
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) {
    // 创建并初始化我们自定义的ThreadPoolTaskScheduler
    CustomThreadPoolTaskScheduler threadPoolTaskScheduler = new CustomThreadPoolTaskScheduler();
    threadPoolTaskScheduler.setPoolSize(5); // 设置线程池大小，可根据需求调整
    threadPoolTaskScheduler.setThreadNamePrefix("custom-scheduled-task-"); // 设置线程名称前缀
    threadPoolTaskScheduler.initialize(); // 必须调用initialize方法来启动调度器

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

package com.example.scheduler;

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

public class CustomThreadPoolTaskScheduler extends ThreadPoolTaskScheduler {

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

package com.example.scheduler;

import com.example.utils.GeneralUtils; // 假设这是你的上下文清理工具类
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;
import org.springframework.lang.Nullable;

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类型的任务
    return new CustomTask<>(task);
  }

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

  /**
   * 自定义任务包装器，用于在任务执行前后插入清理逻辑。
   */
  private record CustomTask<V>(RunnableScheduledFuture<V> task)
      implements RunnableScheduledFuture<V> {

    @Override
    public void run() {
      try {
        // --- 在任务执行前执行的逻辑 ---
        // 例如：初始化MDC、设置一些线程局部变量等
        // GeneralUtils.initializeContext();

        task.run(); // 执行原始任务
      } finally {
        // --- 在任务执行后（无论成功失败）执行的清理逻辑 ---
        // 这是清理线程上下文的关键位置
        GeneralUtils.clearContext(); // 调用你的清理方法
        // 例如：MDC.clear(); ThreadLocal.remove();
      }
    }

    // 以下是RunnableScheduledFuture接口的其他方法，直接委托给被包装的任务
    @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, @Nullable TimeUnit unit)
        throws InterruptedException, ExecutionException, TimeoutException {
      return task.get(timeout, unit);
    }

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

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

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

package com.example.utils;

import org.slf4j.MDC;

public class GeneralUtils {

    private static final ThreadLocal<String> CURRENT_USER = new ThreadLocal<>();
    // 假设有其他ThreadLocal变量...

    public static void setCurrentUser(String user) {
        CURRENT_USER.set(user);
        MDC.put("user", user); // 也可以清理MDC
    }

    public static String getCurrentUser() {
        return CURRENT_USER.get();
    }

    public static void clearContext() {
        System.out.println("Clearing thread context for thread: " + Thread.currentThread().getName());
        CURRENT_USER.remove(); // 清理ThreadLocal变量
        MDC.clear(); // 清理MDC
        // 清理其他任何线程局部变量或上下文信息
    }
}

package com.example.tasks;

import com.example.utils.GeneralUtils;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

@Component
public class MyScheduledTasks {

    private int taskCounter = 0;

    @Scheduled(fixedDelayString = "10000") // 每10秒执行一次
    public void doSomething() {
        taskCounter++;
        String currentUser = "user-" + taskCounter;
        GeneralUtils.setCurrentUser(currentUser); // 设置上下文信息

        System.out.println(
            "Task " + taskCounter + " running in thread: " + Thread.currentThread().getName() +
            ", Context User: " + GeneralUtils.getCurrentUser() +
            ", MDC User: " + MDC.get("user")
        );

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

        // 任务结束，finally块中的GeneralUtils.clearContext()会自动执行
    }
}