ai-security-xdr/haobang-security-dm/syslog-serve/src/main/java/com/haobang/apipull/ApiPullScheduler.java

package com.haobang.apipull;

import com.common.mapper.DeviceCollectApiMapper;
import com.haobang.config.AppProperties;
import com.haobang.dispatch.DeviceContextResolver;
import com.haobang.apipull.persistence.DeviceCollectApiCheckpoint;
import com.haobang.apipull.persistence.DeviceCollectApiConfig;
import com.haobang.apipull.persistence.DeviceCollectApiRun;
import com.haobang.wal.SyslogWalService;
import java.net.URI;
import java.time.LocalDateTime;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.*;
import java.util.concurrent.ConcurrentHashMap;
import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

/**
 * API 拉取任务调度器和并发闸门。
 *
 * <p>并发限制分三层：同一 JVM 内同一任务只能运行一次、全局并发不超过线程池大小、同一设备 主机不超过 per-host 上限。跨实例互斥由达梦租约和单调递增的 fencing token
 * 保证。
 */
@Component
public class ApiPullScheduler {
  private static final Logger log = LoggerFactory.getLogger(ApiPullScheduler.class);
  private final ApiPullProperties properties;
  private final DeviceCollectApiMapper mapper;
  private final ApiPullRunner runner;
  private final SyslogWalService wal;
  private final DeviceContextResolver deviceResolver;
  private final AppProperties appProperties;
  private ExecutorService executor;
  private Semaphore global;
  private final ConcurrentHashMap<String, Semaphore> hosts = new ConcurrentHashMap<>();
  private final Set<Integer> localRunning = ConcurrentHashMap.newKeySet();

  public ApiPullScheduler(
      ApiPullProperties properties,
      DeviceCollectApiMapper mapper,
      ApiPullRunner runner,
      SyslogWalService wal,
      DeviceContextResolver deviceResolver,
      AppProperties appProperties) {
    this.properties = properties;
    this.mapper = mapper;
    this.runner = runner;
    this.wal = wal;
    this.deviceResolver = deviceResolver;
    this.appProperties = appProperties;
  }

  @PostConstruct
  public void init() {
    properties.validate();
    global = new Semaphore(properties.getGlobalConcurrency());
    executor =
        Executors.newFixedThreadPool(
            properties.getGlobalConcurrency(),
            r -> {
              Thread t = new Thread(r, "device-api-pull");
              t.setDaemon(true);
              return t;
            });
    log.info(
        "设备 API 拉取调度器配置: enabled={}, globalConcurrency={}, perHostConcurrency={}, leaseSeconds={},"
            + " maxPagesPerSlice={}, slowPageThresholdMs={}",
        properties.isEnabled(),
        properties.getGlobalConcurrency(),
        properties.getPerHostConcurrency(),
        properties.getLeaseSeconds(),
        properties.getMaxPagesPerSlice(),
        properties.getSlowPageThresholdMs());
  }

  /** 扫描到期任务；总开关关闭时不查询数据库，也不创建任何设备请求。 */
  @Scheduled(fixedDelayString = "${device.api.scheduler.scan-interval-ms:5000}")
  public void scan() {
    if (!properties.isEnabled()) return;
    for (DeviceCollectApiConfig c
        : mapper.selectDueConfigs(appProperties.getAppServiceDeviceCollectId()))
      try {
        trySubmit(c, "SCHEDULED");
      } catch (Throwable failure) {
        log.warn("设备 API 调度提交失败: taskId={}, error={}", c.getCollectTaskId(), failure.toString());
      }
  }

  /** 提交一次人工运行；手工操作不会绕过任务启用状态或资源保护。 */
  public DeviceCollectApiRun trigger(Integer taskId) {
    if (!properties.isEnabled())
      throw new ApiPullException(
          "API_PULL_DISABLED",
          "设备 API 拉取总开关未开启",
          org.springframework.http.HttpStatus.SERVICE_UNAVAILABLE);
    if (wal.getFatalFailure() != null || wal.isDiskCritical() || wal.shouldPauseTcp())
      throw new ApiPullException(
          "WAL_UNAVAILABLE",
          "WAL 不可用或处于接入保护水位",
          org.springframework.http.HttpStatus.SERVICE_UNAVAILABLE);
    DeviceCollectApiConfig c = mapper.selectConfig(taskId);
    if (c == null)
      throw new ApiPullException(
          "TASK_NOT_FOUND", "API 拉取任务不存在", org.springframework.http.HttpStatus.NOT_FOUND);
    if (!c.isEnabled())
      throw new ApiPullException(
          "TASK_NOT_PAUSED", "任务处于暂停状态，不能手工触发", org.springframework.http.HttpStatus.CONFLICT);
    DeviceCollectApiRun run = trySubmit(c, "MANUAL");
    if (run == null)
      throw new ApiPullException(
          "LEASE_CONFLICT", "任务正在运行或无可用并发槽", org.springframework.http.HttpStatus.CONFLICT);
    return run;
  }

  private DeviceCollectApiRun trySubmit(DeviceCollectApiConfig c, String trigger) {
    if (!c.isEnabled()) return null;
    if (!requireApiCollectionMode(c, "MANUAL".equals(trigger))) return null;
    if (wal.getFatalFailure() != null || wal.isDiskCritical() || wal.shouldPauseTcp()) {
      mapper.scheduleNext(c.getCollectTaskId(), c.getConfigVersion(), c.getPollIntervalSeconds());
      return null;
    }
    // 获取顺序固定为任务互斥 -> 全局许可 -> 主机许可 -> DB 租约，失败时按持有情况释放。
    if (!localRunning.add(c.getCollectTaskId())) return null;
    if (!global.tryAcquire()) {
      localRunning.remove(c.getCollectTaskId());
      return null;
    }
    String host;
    try {
      host = new URI(c.getBaseUrl()).getHost();
    } catch (Exception e) {
      releaseLocal(c.getCollectTaskId(), null);
      return null;
    }
    Semaphore hostPermit =
        hosts.computeIfAbsent(host, h -> new Semaphore(properties.getPerHostConcurrency()));
    if (!hostPermit.tryAcquire()) {
      releaseLocal(c.getCollectTaskId(), null);
      return null;
    }
    DeviceCollectApiCheckpoint cp = null;
    DeviceCollectApiRun submittedRun = null;
    try {
      // DB 租约是跨实例最终互斥；本地 semaphore 只用于资源控制，不能代替 fencing。
      if (mapper.acquireLease(
              c.getCollectTaskId(),
              c.getConfigVersion(),
              properties.getInstanceId(),
              properties.getLeaseSeconds(),
              appProperties.getAppServiceDeviceCollectId())
          != 1) {
        releaseLocal(c.getCollectTaskId(), hostPermit);
        return null;
      }
      cp = mapper.selectCheckpoint(c.getCollectTaskId());
      DeviceCollectApiRun run = new DeviceCollectApiRun();
      run.setRunId(UUID.randomUUID().toString());
      run.setCollectTaskId(c.getCollectTaskId());
      run.setDeviceId(c.getDeviceId());
      run.setAdapterType(c.getAdapterType());
      run.setAdapterVersion(c.getAdapterVersion());
      run.setStreamKey(c.getStreamKey());
      run.setInstanceId(properties.getInstanceId());
      run.setConfigVersion(c.getConfigVersion());
      run.setFencingToken(cp.getFencingToken());
      run.setTriggerType(trigger);
      run.setStatus("RUNNING");
      run.setScheduledAt(LocalDateTime.now(ApiPullTime.DATABASE_ZONE));
      mapper.insertRun(run);
      submittedRun = run;
      final DeviceCollectApiCheckpoint leasedCheckpoint = cp;
      // 成功提交后，许可的所有权交给工作线程并由其 finally 统一释放。
      executor.execute(
          () -> {
            try {
              // 模式可能在排队期间变化；执行 HTTP 前再次校验，避免已提交线程绕过新策略。
              if (requireApiCollectionMode(c, false)) {
                runner.execute(c, leasedCheckpoint, run);
              } else {
                run.setStatus("FAILED");
                run.setErrorCode("COLLECTION_MODE_DENIED");
                run.setErrorSummary("设备采集模式不允许 API_PULL");
                run.setFinishedAt(LocalDateTime.now(ApiPullTime.DATABASE_ZONE));
                mapper.updateRun(run);
                mapper.releaseLease(
                    c.getCollectTaskId(), properties.getInstanceId(), leasedCheckpoint.getFencingToken());
              }
            } finally {
              releaseLocal(c.getCollectTaskId(), hostPermit);
            }
          });
      return run;
    } catch (Throwable e) {
      if (submittedRun != null) {
        submittedRun.setStatus("FAILED");
        submittedRun.setErrorCode("SCHEDULER_REJECTED");
        submittedRun.setErrorSummary("任务未能提交到本地执行器");
        submittedRun.setFinishedAt(LocalDateTime.now(ApiPullTime.DATABASE_ZONE));
        mapper.updateRun(submittedRun);
      }
      if (cp != null && cp.getFencingToken() != null)
        mapper.releaseLease(c.getCollectTaskId(), properties.getInstanceId(), cp.getFencingToken());
      releaseLocal(c.getCollectTaskId(), hostPermit);
      throw e;
    }
  }

  /**
   * 调度扫描遇到禁止 API 的设备时推进 next_run_at，避免同一到期任务按 5 秒扫描周期空转。
   */
  private boolean requireApiCollectionMode(DeviceCollectApiConfig c, boolean failRequest) {
    DeviceContextResolver.Resolution resolution = deviceResolver.resolveDeviceId(c.getDeviceId());
    if (resolution.isRegistered()) return true;
    // 非归属节点不能修改任务调度时间；设备归属可能在扫描后、执行前发生变化。
    // 手工触发失败同样不应改变正常调度计划。
    if (!failRequest && resolution.getType() != DeviceContextResolver.Type.NOT_ASSIGNED) {
      mapper.scheduleNext(c.getCollectTaskId(), c.getConfigVersion(), c.getPollIntervalSeconds());
    }
    String reason = resolution.getReason() == null ? "设备不允许 API_PULL" : resolution.getReason();
    log.warn("设备 API 拉取已被设备策略跳过: taskId={}, deviceId={}, type={}, reason={}",
        c.getCollectTaskId(), c.getDeviceId(), resolution.getType(), reason);
    if (failRequest) {
      String code = resolution.getType() == DeviceContextResolver.Type.COLLECTION_MODE_DENIED
          ? "COLLECTION_MODE_DENIED" : "DEVICE_CONTEXT_DENIED";
      throw new ApiPullException(code, reason, org.springframework.http.HttpStatus.CONFLICT);
    }
    return false;
  }

  private void releaseLocal(Integer taskId, Semaphore host) {
    if (host != null) host.release();
    global.release();
    localRunning.remove(taskId);
  }

  @PreDestroy
  public void shutdown() {
    // 等待在途窗口正常收尾；超时中断时，未提交的 checkpoint 会使该窗口下次重新拉取。
    if (executor == null) return;
    executor.shutdown();
    try {
      if (!executor.awaitTermination(30, TimeUnit.SECONDS)) executor.shutdownNow();
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      executor.shutdownNow();
    }
  }
}