基于Kafka的Spring Batch远程分区多租户：JobRepository初始化前设租户

可以在Worker读取Kafka消息阶段获取租户值 #

Spring Batch远程分区通过Kafka传递的请求消息本质是StepExecutionRequest对象，其中包含携带租户编码的ExecutionContext。你可以在Kafka入站流中添加消息处理器，提前解析出租户编码并存入上下文，确保JobRepository初始化时能获取到该值。

具体实现步骤 #

1. 定义租户上下文Holder
   用ThreadLocal存储租户编码，确保多线程环境下的隔离性：

   public class TenantContextHolder {
       private static final ThreadLocal<String> tenantCodeHolder = new ThreadLocal<>();
   
       public static void setTenantCode(String tenantCode) {
           tenantCodeHolder.set(tenantCode);
       }
   
       public static String getTenantCode() {
           return tenantCodeHolder.get();
       }
   
       public static void clear() {
           tenantCodeHolder.remove();
       }
   }
   

2. 修改Kafka入站流，添加租户提取逻辑
   在消息进入处理通道前，解析StepExecutionRequest并设置租户编码到上下文：

   @Bean // request coming from manager
   public IntegrationFlow inboundFlow(ConsumerFactory consumerFactory) {
       return IntegrationFlows
               .from(Kafka.inboundChannelAdapter(consumerFactory, new ConsumerProperties("requestForWorkers")))
               .transform(this::extractAndSetTenant)
               .channel(requestForWorkers())
               .get();
   }
   
   private Object extractAndSetTenant(Object payload) {
       if (payload instanceof StepExecutionRequest) {
           StepExecutionRequest request = (StepExecutionRequest) payload;
           ExecutionContext executionContext = request.getStepExecution().getExecutionContext();
           // 替换为你在Manager端存入的租户编码key
           String tenantCode = executionContext.getString("tenantCode");
           TenantContextHolder.setTenantCode(tenantCode);
       }
       return payload; // 传递原消息，不影响后续Spring Batch处理
   }
   

3. 在JobRepository初始化时获取租户编码
   自定义JobRepository配置，从TenantContextHolder中读取租户信息，执行初始化逻辑：

   @Bean
   public JobRepository jobRepository(DataSource dataSource, PlatformTransactionManager transactionManager) throws Exception {
       JobRepositoryFactoryBean factory = new JobRepositoryFactoryBean();
       factory.setDataSource(dataSource);
       factory.setTransactionManager(transactionManager);
   
       String tenantCode = TenantContextHolder.getTenantCode();
       if (tenantCode != null) {
           // 根据租户编码执行初始化操作，比如设置表前缀、切换数据源等
           factory.setTablePrefix(tenantCode + "_BATCH_");
       }
   
       factory.afterPropertiesSet();
       return factory.getObject();
   }
   

4. 清理ThreadLocal，避免内存泄漏
   添加Step执行监听器，在Step完成后清理租户上下文：

   @Bean
   public StepExecutionListener tenantCleanupListener() {
       return new StepExecutionListener() {
           @Override
           public void beforeStep(StepExecution stepExecution) {}
   
           @Override
           public ExitStatus afterStep(StepExecution stepExecution) {
               TenantContextHolder.clear();
               return ExitStatus.COMPLETED;
           }
       };
   }
   

注意事项 #

• 确保Manager端确实将租户编码存入了StepExecution的ExecutionContext中，key要和Worker端解析时的key一致。
• Kafka消费者是多线程模型，ThreadLocal的清理必须执行，防止内存泄漏。

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