ai-security-xdr/haobang-security-dm/syslog-serve/src/main/java/com/haobang/dispatch/SyslogWalDispatcher.java

package com.haobang.dispatch;

import com.haobang.config.AppProperties;
import com.haobang.wal.SyslogWalRecord;
import com.haobang.wal.SyslogWalService;
import com.kafka.kafkaProducer;
import com.netty.SyslogMessageHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.Semaphore;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.LongAdder;

/**
 * 从 WAL 有界并发异步投递 Kafka，并只按连续 ACK 推进 WAL 游标。
 *
 * <p>v2/v3 保持原有按 source_ip 识别设备和单时间参数的生产路径；v4 改为按可信 device_id
 * 识别设备，并分别传递 receive_time/event_time。兼容分支留在 Dispatcher，Kafka 消息外层结构
 * 不因接入方式改变。
 */
@Component
public class SyslogWalDispatcher {
    private static final Logger log = LoggerFactory.getLogger(SyslogWalDispatcher.class);
    private static final DateTimeFormatter WAL_TIME = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSS");
    private static final DateTimeFormatter HEADER_TIME = DateTimeFormatter.ofPattern("yyyyMMddHHmmssSSS");

    private final SyslogWalService walService;
    private final DeviceContextResolver deviceResolver;
    private final kafkaProducer producer;
    private final AppProperties properties;
    private final AtomicBoolean running = new AtomicBoolean();
    private final Object completionMonitor = new Object();
    private final Map<Long, SyslogWalRecord> completedOutOfOrder = new HashMap<>();
    private final LongAdder kafkaAcked = new LongAdder();
    private final LongAdder terminalFiltered = new LongAdder();
    private final LongAdder terminalDropped = new LongAdder();
    private final LongAdder retryCount = new LongAdder();
    private final LongAdder largeDispatched = new LongAdder();
    private final AtomicLong inflight = new AtomicLong();

    private Semaphore permits;
    private ByteBudget byteBudget;
    private ExecutorService workers;
    private ScheduledExecutorService scheduler;
    private Thread readerThread;
    private SyslogWalService.WalCursor cursor;
    private long nextCommitSequence;
    private long nextReadSequence;

    public SyslogWalDispatcher(SyslogWalService walService,
                               DeviceContextResolver deviceResolver,
                               kafkaProducer producer,
                               AppProperties properties) {
        this.walService = walService;
        this.deviceResolver = deviceResolver;
        this.producer = producer;
        this.properties = properties;
    }

    @PostConstruct
    public void start() throws Exception {
        validateConfiguration();
        permits = new Semaphore(properties.getSyslogDispatchMaxInflight());
        byteBudget = new ByteBudget(properties.getSyslogDispatchMaxInflightBytes());
        workers = Executors.newFixedThreadPool(properties.getSyslogDispatchWorkerThreads(),
                namedFactory("syslog-dispatch-worker"));
        scheduler = Executors.newScheduledThreadPool(2, namedFactory("syslog-dispatch-scheduler"));
        cursor = walService.openCursor();
        nextCommitSequence = walService.getCheckpointSequence() + 1L;
        nextReadSequence = nextCommitSequence;
        running.set(true);
        readerThread = new Thread(this::readLoop, "syslog-wal-dispatch-reader");
        readerThread.setDaemon(true);
        readerThread.start();
        scheduler.scheduleWithFixedDelay(this::logStatus,
                properties.getSyslogStatusLogIntervalMs(), properties.getSyslogStatusLogIntervalMs(),
                TimeUnit.MILLISECONDS);
        log.info("Syslog WAL Dispatcher 已启动: workerThreads={}, maxInflight={}, maxInflightBytes={}, "
                        + "terminalRecordEnabled={}, checkpointSequence={}",
                properties.getSyslogDispatchWorkerThreads(), properties.getSyslogDispatchMaxInflight(),
                properties.getSyslogDispatchMaxInflightBytes(),
                properties.isSyslogDispatchTerminalRecordEnabled(),
                walService.getCheckpointSequence());
        if (!properties.isSyslogDispatchTerminalRecordEnabled()) {
            log.warn("无效设备日志直接丢弃已启用: 新 Syslog 正文不进 WAL，历史记录不创建 Terminal 文件");
        }
    }

    private void readLoop() {
        try {
            while (running.get()) {
                permits.acquire();
                SyslogWalRecord record = cursor.next(1000L);
                if (record == null) {
                    permits.release();
                    continue;
                }
                record.setDispatchPayloadBytes(Math.max(1L,
                        Math.max(record.getRawLength(), utf8Length(record.getRawMessage()))));
                byteBudget.acquire(record.getDispatchPayloadBytes());
                if (record.getSequence() != nextReadSequence) {
                    byteBudget.release(record.getDispatchPayloadBytes());
                    permits.release();
                    throw new IllegalStateException("WAL sequence 不连续，拒绝越过: expected="
                            + nextReadSequence + ", actual=" + record.getSequence());
                }
                nextReadSequence++;
                inflight.incrementAndGet();
                workers.execute(() -> dispatch(record, 1));
            }
        } catch (InterruptedException interrupted) {
            Thread.currentThread().interrupt();
        } catch (Throwable failure) {
            running.set(false);
            log.error("WAL Dispatcher 读取失败，停止推进游标以保护数据", failure);
        }
    }

    private void dispatch(SyslogWalRecord record, int attempt) {
        if (!running.get()) return;
        if (record.isQuarantinedInvalid()) {
            log.error("损坏 WAL 已隔离并推进游标: sequence={}, eventId={}, reason={}",
                    record.getSequence(), record.getEventId(), record.getValidationError());
            markComplete(record);
            return;
        }
        final DeviceContextResolver.Resolution resolution;
        try {
            // API 的 HTTP 对端可能是代理或网关；v4 必须使用任务写入 WAL 的 device_id。
            if (record.getSchemaVersion() >= 4 && record.getDeviceId() != null) {
                resolution = deviceResolver.resolveDeviceId(record.getDeviceId());
            } else {
                resolution = deviceResolver.resolve(record.getSourceIp(), record.getProtocol());
            }
        } catch (Throwable failure) {
            scheduleRetry(record, attempt, "DEVICE_LOOKUP", failure);
            return;
        }

        if (!resolution.isRegistered()) {
            try {
                if (properties.isSyslogDispatchTerminalRecordEnabled()) {
                    walService.recordTerminal(record, resolution.getType().name(), resolution.getReason());
                } else {
                    terminalDropped.increment();
                }
                terminalFiltered.increment();
                markComplete(record);
            } catch (Throwable failure) {
                scheduleRetry(record, attempt, "TERMINAL_LEDGER", failure);
            }
            return;
        }

        final String receiveTime;
        final String eventTime;
        try {
            receiveTime = LocalDateTime.parse(record.getReceiveTime(), WAL_TIME).format(HEADER_TIME);
            String walEventTime = record.getEventTime() == null
                    ? record.getReceiveTime()
                    : record.getEventTime();
            eventTime = LocalDateTime.parse(walEventTime, WAL_TIME).format(HEADER_TIME);

            // v2/v3 继续调用旧重载，确保升级前已落盘记录产生的 Kafka 消息完全不变。
            boolean large = properties.getSyslogLargeMessageThresholdBytes() > 0L
                    && record.getDispatchPayloadBytes() >= properties.getSyslogLargeMessageThresholdBytes();
            java.util.concurrent.CompletableFuture<kafkaProducer.DeliveryResult> deliveryFuture = large
                    ? producer.messagePushLarge(
                            resolution.getDevice(), record.getRawMessage(), receiveTime, eventTime,
                            record.getEventId(), record.getPointTimeNs())
                    : record.getSchemaVersion() >= 4
                            ? producer.messagePush(
                                    resolution.getDevice(),
                                    record.getRawMessage(),
                                    receiveTime,
                                    eventTime,
                                    record.getEventId(),
                                    record.getPointTimeNs())
                            : producer.messagePush(
                                    resolution.getDevice(),
                                    record.getRawMessage(),
                                    receiveTime,
                                    record.getEventId(),
                                    record.getPointTimeNs());
            deliveryFuture
                    .whenComplete((delivery, failure) -> {
                        if (failure == null) {
                            kafkaAcked.increment();
                            if (large) largeDispatched.increment();
                            markComplete(record);
                        } else {
                            scheduleRetry(record, attempt, "KAFKA_ACK", unwrap(failure));
                        }
                    });
        } catch (Throwable failure) {
            // 加密、配置或 Producer 本地错误都不能越过 WAL；修复依赖后沿用同一 event_id 重试。
            scheduleRetry(record, attempt, "KAFKA_PREPARE", failure);
        }
    }

    private void scheduleRetry(SyslogWalRecord record, int attempt, String stage, Throwable failure) {
        retryCount.increment();
        long retries = retryCount.sum();
        if (attempt == 1 || retries % 100 == 0) {
            log.warn("WAL 事件投递失败，将保留原 event_id 重试: eventId={}, sequence={}, stage={}, "
                            + "attempt={}, retryTotal={}, error={}",
                    record.getEventId(), record.getSequence(), stage, attempt, retries,
                    failure == null ? "unknown" : failure.toString());
        }
        if (!running.get()) return;
        scheduler.schedule(() -> {
            if (running.get()) workers.execute(() -> dispatch(record, attempt + 1));
        }, properties.getSyslogDispatchRetryBackoffMs(), TimeUnit.MILLISECONDS);
    }

    private void markComplete(SyslogWalRecord record) {
        synchronized (completionMonitor) {
            if (record.getSequence() < nextCommitSequence) return;
            if (completedOutOfOrder.putIfAbsent(record.getSequence(), checkpointOnly(record)) != null) return;
            // 已 ACK 的正文不再占字节预算；条数许可仍要等连续 checkpoint，限制乱序完成表大小。
            byteBudget.release(record.getDispatchPayloadBytes());
            SyslogWalRecord latest = null;
            int released = 0;
            while (true) {
                SyslogWalRecord contiguous = completedOutOfOrder.remove(nextCommitSequence);
                if (contiguous == null) break;
                latest = contiguous;
                nextCommitSequence++;
                released++;
            }
            if (latest != null) walService.markCheckpoint(latest);
            if (released > 0) {
                inflight.addAndGet(-released);
                permits.release(released);
            }
        }
    }

    private SyslogWalRecord checkpointOnly(SyslogWalRecord source) {
        SyslogWalRecord value = new SyslogWalRecord();
        value.setSequence(source.getSequence());
        value.setSegment(source.getSegment());
        value.setStartOffset(source.getStartOffset());
        value.setEndOffset(source.getEndOffset());
        return value;
    }

    private void logStatus() {
        if (!running.get()) return;
        log.info("Syslog WAL状态: networkReceived={}, appended={}, admissionOversized={}, walPendingApprox={}, walPendingBytes={}, walQueue={}/{}, durableSequence={}, "
                        + "checkpointSequence={}, kafkaInflight={}, kafkaInflightBytes={}, kafkaAcked={}, largeDispatched={}, "
                        + "terminalRecordEnabled={}, terminalFiltered={}, terminalDropped={}, "
                        + "retry={}, tcpPauseTotal={}, udpDropped={}, deviceCacheSize={}, deviceCacheHits={}, "
                        + "deviceCacheMisses={}",
                SyslogMessageHandler.getNetworkReceivedCount(), walService.getAppendedCount(),
                walService.getOversizedCount(), walService.getPendingRecordsApprox(), walService.getPendingAppendBytes(),
                walService.getQueueSize(), walService.getQueueCapacity(), walService.getDurableSequence(),
                walService.getCheckpointSequence(), inflight.get(), byteBudget == null ? 0L : byteBudget.used(),
                kafkaAcked.sum(), largeDispatched.sum(), properties.isSyslogDispatchTerminalRecordEnabled(),
                terminalFiltered.sum(), terminalDropped.sum(),
                retryCount.sum(), SyslogMessageHandler.getTcpPauseCount(),
                SyslogMessageHandler.getUdpDroppedCount(), deviceResolver.getCacheSize(),
                deviceResolver.getHitCount(), deviceResolver.getMissCount());
    }

    private Throwable unwrap(Throwable failure) {
        if (failure instanceof java.util.concurrent.CompletionException && failure.getCause() != null) {
            return failure.getCause();
        }
        return failure;
    }

    /** 不分配临时 byte[] 的 UTF-8 长度估算；非法孤立代理项按 3 字节保守计算。 */
    static long utf8Length(String value) {
        if (value == null || value.isEmpty()) return 0L;
        long bytes = 0L;
        for (int index = 0; index < value.length(); index++) {
            char current = value.charAt(index);
            if (current <= 0x7f) bytes++;
            else if (current <= 0x7ff) bytes += 2L;
            else if (Character.isHighSurrogate(current) && index + 1 < value.length()
                    && Character.isLowSurrogate(value.charAt(index + 1))) {
                bytes += 4L;
                index++;
            } else bytes += 3L;
        }
        return bytes;
    }

    private ThreadFactory namedFactory(String prefix) {
        AtomicLong sequence = new AtomicLong();
        return runnable -> {
            Thread thread = new Thread(runnable, prefix + "-" + sequence.incrementAndGet());
            thread.setDaemon(true);
            return thread;
        };
    }

    private void validateConfiguration() {
        if (properties.getSyslogDispatchWorkerThreads() <= 0
                || properties.getSyslogDispatchMaxInflight() <= 0
                || properties.getSyslogDispatchMaxInflightBytes() < 0L
                || properties.getSyslogDispatchRetryBackoffMs() <= 0
                || properties.getSyslogDispatchShutdownTimeoutMs() <= 0
                || properties.getSyslogDeviceContextCacheTtlMs() <= 0
                || properties.getSyslogStatusLogIntervalMs() <= 0) {
            throw new IllegalArgumentException("Dispatcher 线程、在途、重试、缓存和状态日志参数必须大于 0");
        }
        if (properties.getSyslogMessageMaxBytes() > 0L
                && properties.getSyslogDispatchMaxInflightBytes() > 0L
                && properties.getSyslogMessageMaxBytes() > properties.getSyslogDispatchMaxInflightBytes()) {
            throw new IllegalArgumentException("单条消息上限不能大于 Dispatcher 在途字节预算");
        }
    }

    public long getInflight() { return inflight.get(); }
    public long getKafkaAcked() { return kafkaAcked.sum(); }
    public long getTerminalFiltered() { return terminalFiltered.sum(); }
    public long getTerminalDropped() { return terminalDropped.sum(); }
    public boolean isTerminalRecordEnabled() { return properties.isSyslogDispatchTerminalRecordEnabled(); }
    public long getRetryCount() { return retryCount.sum(); }
    public long getLargeDispatched() { return largeDispatched.sum(); }
    public long getInflightBytes() { return byteBudget == null ? 0L : byteBudget.used(); }

    /** 同时保留条数 Semaphore 和字节预算，避免少量大消息占满堆。 */
    static final class ByteBudget {
        private final long capacity;
        private long used;

        ByteBudget(long capacity) {
            if (capacity < 0L) throw new IllegalArgumentException("capacity 不能小于 0");
            this.capacity = capacity;
        }

        synchronized void acquire(long bytes) throws InterruptedException {
            if (bytes <= 0L) throw new IllegalArgumentException("bytes 必须大于 0");
            if (capacity == 0L) { used += bytes; return; }
            if (bytes > capacity) {
                // 硬上限关闭时仍保证进展：超预算单条只能独占，不允许与其他正文并存。
                while (used != 0L) wait();
                used = bytes;
                return;
            }
            while (used > capacity - bytes) wait();
            used += bytes;
        }

        synchronized void release(long bytes) {
            used -= bytes;
            if (used < 0L) throw new IllegalStateException("Dispatcher 字节预算释放越界");
            notifyAll();
        }

        synchronized long used() { return used; }
    }

    @PreDestroy
    public void shutdown() {
        if (!running.compareAndSet(true, false)) return;
        if (readerThread != null) readerThread.interrupt();
        if (scheduler != null) scheduler.shutdownNow();
        if (workers != null) workers.shutdown();
        long deadline = System.nanoTime()
                + TimeUnit.MILLISECONDS.toNanos(properties.getSyslogDispatchShutdownTimeoutMs());
        try {
            while (inflight.get() > 0 && System.nanoTime() < deadline) Thread.sleep(20L);
            if (workers != null) {
                long remaining = Math.max(1L, deadline - System.nanoTime());
                workers.awaitTermination(remaining, TimeUnit.NANOSECONDS);
            }
        } catch (InterruptedException interrupted) {
            Thread.currentThread().interrupt();
        }
        if (workers != null && !workers.isTerminated()) workers.shutdownNow();
        try { if (cursor != null) cursor.close(); }
        catch (Exception failure) { log.warn("关闭 WAL cursor 失败", failure); }
        log.info("Syslog WAL Dispatcher 已停止: inflight={}, checkpointSequence={}",
                inflight.get(), walService.getCheckpointSequence());
    }
}