ai-security-xdr/haobang-security-dm/syslog-consumer-rule/src/main/java/com/common/service/IngressStatsProjectionWriter.java

package com.common.service;

import com.common.entity.IngressStatsEvent;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.support.TransactionTemplate;

import java.sql.Date;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.ZoneId;
import java.util.Comparator;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 把 Kafka 入口元数据按设备/探针批内聚合后幂等写入小型统计投影。
 * 不写逐事件接收日志；失败时抛出，让调用方禁止提交 Kafka offset。
 */
@Service
public class IngressStatsProjectionWriter {
    private final JdbcTemplate jdbc;
    private final PlatformTransactionManager transactionManager;
    private final StatsWriteModeProvider modeProvider;
    private final StatsMaterializationPublisher materializationPublisher;

    @Value("${stats.ingress-metadata.enabled:false}")
    private boolean enabled;
    @Value("${stats.incremental.read-enabled:false}")
    private boolean readEnabled;
    @Value("${stats.incremental.transaction-timeout-seconds:45}")
    private int transactionTimeoutSeconds;
    @Value("${stats.direct.shadow-device-table:device_daily_stats_direct_shadow}")
    private String shadowDeviceTable;
    @Value("${stats.business-zone:Asia/Shanghai}")
    private String businessZone;
    private final AtomicLong successfulBatches = new AtomicLong();
    private final AtomicLong receivedEvents = new AtomicLong();
    private final AtomicLong updatedGroups = new AtomicLong();
    private final AtomicLong failedBatches = new AtomicLong();
    private volatile LocalDateTime lastSuccessTime;
    private volatile LocalDateTime lastFailureTime;

    public IngressStatsProjectionWriter(JdbcTemplate jdbc,
                                        PlatformTransactionManager transactionManager,
                                        StatsWriteModeProvider modeProvider,
                                        StatsMaterializationPublisher materializationPublisher) {
        this.jdbc = jdbc;
        this.transactionManager = transactionManager;
        this.modeProvider = modeProvider;
        this.materializationPublisher = materializationPublisher;
    }

    public int writeBatch(List<IngressStatsEvent> events) {
        if (!enabled || events == null || events.isEmpty()) return 0;
        final Aggregation aggregation = aggregate(events);
        if (aggregation.devices.isEmpty() && aggregation.collects.isEmpty()) return 0;

        try {
            int groups;
            try {
                groups = writeAggregation(aggregation);
            } catch (RuntimeException concurrentWrite) {
                if (!retryableConcurrentWrite(concurrentWrite)) throw concurrentWrite;
                // 新设备/探针在多节点首次出现时，两个 MERGE 可能同时走 NOT MATCHED。
                // 锁冲突或失败事务已整体回滚；重试后的 MIN/MAX 更新仍保持幂等。
                groups = writeAggregation(aggregation);
            }
            successfulBatches.incrementAndGet();
            receivedEvents.addAndGet(events.size());
            updatedGroups.addAndGet(groups);
            lastSuccessTime = LocalDateTime.now();
            return groups;
        } catch (RuntimeException failure) {
            failedBatches.incrementAndGet();
            lastFailureTime = LocalDateTime.now();
            throw failure;
        }
    }

    public Map<String, Object> status() {
        Map<String, Object> status = new LinkedHashMap<>();
        status.put("enabled", enabled);
        status.put("read_enabled", readEnabled);
        status.put("target_device_table", modeProvider.writesShadowProjection()
                ? "device_daily_stats+" + shadowDeviceTable : "device_daily_stats");
        status.put("business_materialization_mode", asynchronousMaterialization()
                ? "KAFKA_DIRTY_KEY" : "SYNCHRONOUS");
        status.put("successful_batches", successfulBatches.get());
        status.put("received_events", receivedEvents.get());
        status.put("updated_groups", updatedGroups.get());
        status.put("failed_batches", failedBatches.get());
        status.put("last_success_time", lastSuccessTime);
        status.put("last_failure_time", lastFailureTime);
        return status;
    }

    private static boolean retryableConcurrentWrite(RuntimeException failure) {
        return DirectStatsProjectionWriter.isConcurrentFailure(failure);
    }

    private int writeAggregation(Aggregation aggregation) {
        TransactionTemplate transaction = new TransactionTemplate(transactionManager);
        transaction.setTimeout(Math.max(1, transactionTimeoutSeconds));
        Integer updated = transaction.execute(status -> {
            int groups = 0;
            boolean asyncMaterialization = asynchronousMaterialization();
            List<StatsMaterializationCommand> commands = asyncMaterialization
                    ? new ArrayList<>(aggregation.devices.size() + aggregation.collects.size())
                    : null;
            LocalDate today = LocalDate.now(ZoneId.of(businessZone));
            for (Map.Entry<DeviceKey, LocalDateTime> entry : aggregation.devices.entrySet()) {
                DeviceKey key = entry.getKey();
                LocalDateTime receiveTime = entry.getValue();
                // 接收时间是 MIN/MAX 幂等字段。SHADOW 期间正式表仍需对外服务，因此同时写正式和影子；
                // 与可累加的 parse/non-log 不同，重复 MAX 不会造成双计数。
                mergeDeviceReceive("device_daily_stats", key, receiveTime);
                if (modeProvider.writesShadowProjection()) {
                    mergeDeviceReceive(safeTable(shadowDeviceTable), key, receiveTime);
                }
                if (readEnabled && today.equals(key.date)) {
                    if (asyncMaterialization) {
                        commands.add(StatsMaterializationCommand.device(key.date, key.deviceId));
                    } else {
                        jdbc.update("UPDATE device_device SET last_receive_time=CASE " +
                                        "WHEN last_receive_time IS NULL OR last_receive_time<? THEN ? ELSE last_receive_time END," +
                                        "updated_at=SYSDATE WHERE id=?",
                                receiveTime, receiveTime, key.deviceId);
                    }
                }
                groups++;
            }
            for (Map.Entry<Long, CollectTimes> entry : aggregation.collects.entrySet()) {
                CollectTimes times = entry.getValue();
                jdbc.update("MERGE INTO device_collect_stats t USING " +
                                "(SELECT ? device_collect_id,? first_success_time,? last_success_time FROM DUAL) s " +
                                "ON (t.device_collect_id=s.device_collect_id) " +
                                "WHEN MATCHED THEN UPDATE SET t.first_success_time=CASE " +
                                "WHEN t.first_success_time IS NULL OR t.first_success_time>s.first_success_time " +
                                "THEN s.first_success_time ELSE t.first_success_time END," +
                                "t.last_success_time=CASE WHEN t.last_success_time IS NULL " +
                                "OR t.last_success_time<s.last_success_time THEN s.last_success_time " +
                                "ELSE t.last_success_time END,t.updated_at=SYSDATE " +
                                "WHEN NOT MATCHED THEN INSERT (device_collect_id,first_success_time,last_success_time," +
                                "last_fail_time,projection_seq,updated_at) VALUES " +
                                "(s.device_collect_id,s.first_success_time,s.last_success_time,NULL,0,SYSDATE)",
                        entry.getKey(), times.min, times.max);
                if (readEnabled) {
                    if (asyncMaterialization) {
                        commands.add(StatsMaterializationCommand.collect(entry.getKey()));
                    } else {
                        jdbc.update("UPDATE device_collect_task SET " +
                                        "first_time=CASE WHEN first_time IS NULL OR first_time>? THEN ? ELSE first_time END," +
                                        "last_success_time=CASE WHEN last_success_time IS NULL OR last_success_time<? " +
                                        "THEN ? ELSE last_success_time END,updated_at=SYSDATE WHERE id=?",
                                times.min, times.min, times.max, times.max, entry.getKey());
                    }
                }
                groups++;
            }
            if (asyncMaterialization) {
                materializationPublisher.enqueueAfterCommit(commands);
            }
            return groups;
        });
        return updated == null ? 0 : updated;
    }

    static Aggregation aggregate(List<IngressStatsEvent> events) {
        Aggregation result = new Aggregation();
        for (IngressStatsEvent event : events) {
            if (event == null || event.getReceiveTime() == null) continue;
            if (event.getDeviceId() != null) {
                DeviceKey key = new DeviceKey(event.getReceiveTime().toLocalDate(), event.getDeviceId());
                LocalDateTime current = result.devices.get(key);
                if (current == null || event.getReceiveTime().isAfter(current)) {
                    result.devices.put(key, event.getReceiveTime());
                }
            }
            if (event.getDeviceCollectId() != null) {
                CollectTimes times = result.collects.get(event.getDeviceCollectId());
                if (times == null) result.collects.put(event.getDeviceCollectId(),
                        new CollectTimes(event.getReceiveTime(), event.getReceiveTime()));
                else {
                    if (event.getReceiveTime().isBefore(times.min)) times.min = event.getReceiveTime();
                    if (event.getReceiveTime().isAfter(times.max)) times.max = event.getReceiveTime();
                }
            }
        }
        return result;
    }

    static final class Aggregation {
        // 所有节点使用相同顺序取得投影行锁，避免不同 Kafka 批次顺序造成交叉等待。
        final Map<DeviceKey, LocalDateTime> devices = new TreeMap<>(
                Comparator.comparing((DeviceKey key) -> key.date)
                        .thenComparingLong(key -> key.deviceId));
        final Map<Long, CollectTimes> collects = new TreeMap<>();
    }
    static final class DeviceKey {
        final LocalDate date; final long deviceId;
        DeviceKey(LocalDate date, long deviceId) { this.date = date; this.deviceId = deviceId; }
        @Override public boolean equals(Object other) {
            if (this == other) return true;
            if (!(other instanceof DeviceKey)) return false;
            DeviceKey key = (DeviceKey) other;
            return deviceId == key.deviceId && date.equals(key.date);
        }
        @Override public int hashCode() { return 31 * date.hashCode() + Long.hashCode(deviceId); }
    }
    static final class CollectTimes {
        LocalDateTime min; LocalDateTime max;
        CollectTimes(LocalDateTime min, LocalDateTime max) { this.min = min; this.max = max; }
    }

    private void mergeDeviceReceive(String table, DeviceKey key, LocalDateTime receiveTime) {
        jdbc.update("MERGE INTO " + safeTable(table) + " t USING " +
                        "(SELECT ? stat_date,? device_id,? receive_time FROM DUAL) s " +
                        "ON (t.stat_date=s.stat_date AND t.device_id=s.device_id) " +
                        "WHEN MATCHED THEN UPDATE SET t.last_receive_time=CASE " +
                        "WHEN t.last_receive_time IS NULL OR t.last_receive_time<s.receive_time " +
                        "THEN s.receive_time ELSE t.last_receive_time END,t.updated_at=SYSDATE " +
                        "WHEN NOT MATCHED THEN INSERT (stat_date,device_id,parse_count,non_log_count," +
                        "last_receive_time,projection_seq,updated_at) VALUES " +
                        "(s.stat_date,s.device_id,0,0,s.receive_time,0,SYSDATE)",
                Date.valueOf(key.date), key.deviceId, receiveTime);
    }

    private boolean asynchronousMaterialization() {
        return modeProvider.materializesBusinessFields() && readEnabled
                && materializationPublisher.isEnabled();
    }

    private static String safeTable(String table) {
        if (table == null || !table.matches("[A-Za-z][A-Za-z0-9_]*")) {
            throw new IllegalStateException("非法入口设备统计表名: " + table);
        }
        return table;
    }
}