WorkloadRuntimeStatusMgr.java

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package org.apache.doris.resource.workloadschedpolicy;

import org.apache.doris.catalog.Env;
import org.apache.doris.common.Config;
import org.apache.doris.common.Pair;
import org.apache.doris.common.util.DebugUtil;
import org.apache.doris.common.util.MasterDaemon;
import org.apache.doris.plugin.AuditEvent;
import org.apache.doris.qe.QeProcessorImpl;
import org.apache.doris.system.Backend;
import org.apache.doris.thrift.TQueryStatistics;
import org.apache.doris.thrift.TQueryStatisticsResult;
import org.apache.doris.thrift.TReportWorkloadRuntimeStatusParams;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.IdentityHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLongArray;
import java.util.concurrent.locks.ReentrantLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class WorkloadRuntimeStatusMgr extends MasterDaemon {

    private static final Logger LOG = LogManager.getLogger(WorkloadRuntimeStatusMgr.class);
    // backend process incarnation --> {query id --> (query last report time, query stats)}
    private final ConcurrentMap<BackendIncarnation, BeReportInfo> beToQueryStatsMap = Maps.newConcurrentMap();
    private final ConcurrentMap<Long, Long> lastAcceptedBackendStartTimes = Maps.newConcurrentMap();
    private final ConcurrentMap<RetainedQuery, Integer> pendingQueryReferences
            = Maps.newConcurrentMap();
    // Publish an immutable snapshot for synchronous proc/REST readers.
    private volatile Map<String, TQueryStatistics> queryStatisticsSnapshot = ImmutableMap.of();
    private final ReentrantLock queryAuditEventLock = new ReentrantLock();
    private final Set<PendingAuditEvent> queryAuditEventList = new LinkedHashSet<>();
    private final ReentrantLock pendingAuditBindingLock = new ReentrantLock();
    private final Map<BackendQuery, List<PendingAuditBinding>> unboundAuditParticipants = new HashMap<>();
    private final Map<String, Integer> inFlightAuditQueryReferences = new HashMap<>();
    private volatile long lastWarnTime;

    private static class BackendIncarnation {
        final long backendId;
        final long backendStartTime;

        BackendIncarnation(long backendId, long backendStartTime) {
            this.backendId = backendId;
            this.backendStartTime = backendStartTime;
        }

        @Override
        public boolean equals(Object other) {
            if (!(other instanceof BackendIncarnation)) {
                return false;
            }
            BackendIncarnation that = (BackendIncarnation) other;
            return backendId == that.backendId && backendStartTime == that.backendStartTime;
        }

        @Override
        public int hashCode() {
            return Long.hashCode(backendId) * 31 + Long.hashCode(backendStartTime);
        }
    }

    private static class BeReportInfo {
        final long backendStartTime;
        final ReentrantReadWriteLock lifecycleLock = new ReentrantReadWriteLock();
        // query id --> (query last report time, query stats)
        final ConcurrentMap<String, Pair<Long, TQueryStatisticsResult>> queryStatsMap
                = Maps.newConcurrentMap();

        BeReportInfo(long backendStartTime) {
            this.backendStartTime = backendStartTime;
        }
    }

    private static class BackendQuery {
        final long backendId;
        final String queryId;

        BackendQuery(long backendId, String queryId) {
            this.backendId = backendId;
            this.queryId = queryId;
        }

        @Override
        public boolean equals(Object other) {
            if (!(other instanceof BackendQuery)) {
                return false;
            }
            BackendQuery that = (BackendQuery) other;
            return backendId == that.backendId && queryId.equals(that.queryId);
        }

        @Override
        public int hashCode() {
            return Long.hashCode(backendId) * 31 + queryId.hashCode();
        }
    }

    private static class RetainedQuery {
        final BackendIncarnation incarnation;
        final String queryId;

        RetainedQuery(BackendIncarnation incarnation, String queryId) {
            this.incarnation = incarnation;
            this.queryId = queryId;
        }

        @Override
        public boolean equals(Object other) {
            if (!(other instanceof RetainedQuery)) {
                return false;
            }
            RetainedQuery that = (RetainedQuery) other;
            return incarnation.equals(that.incarnation) && queryId.equals(that.queryId);
        }

        @Override
        public int hashCode() {
            return incarnation.hashCode() * 31 + queryId.hashCode();
        }
    }

    private static class PendingAuditEvent {
        final AuditEvent event;
        final String statisticsQueryId;
        final long[] expectedBackendIds;
        final AtomicLongArray expectedBackendStartTimes;
        final AtomicBoolean active = new AtomicBoolean(true);

        PendingAuditEvent(AuditEvent event, Set<Long> expectedBackendIds, String statisticsQueryId) {
            this.event = event;
            this.statisticsQueryId = statisticsQueryId;
            // The audit queue can contain hundreds of thousands of events. Keep backend IDs in a
            // compact primitive array instead of retaining one boxed hash set per event.
            this.expectedBackendIds = expectedBackendIds.stream()
                    .mapToLong(Long::longValue).toArray();
            this.expectedBackendStartTimes = new AtomicLongArray(this.expectedBackendIds.length);
        }
    }

    private static class PendingAuditBinding {
        final PendingAuditEvent pending;
        final int participantIndex;

        PendingAuditBinding(PendingAuditEvent pending, int participantIndex) {
            this.pending = pending;
            this.participantIndex = participantIndex;
        }
    }

    private static class RuntimeStatisticsSnapshot {
        final Map<String, TQueryStatistics> queryStatistics = new HashMap<>();
    }

    public WorkloadRuntimeStatusMgr() {
        super("workload-runtime-stats-thread", Config.workload_runtime_status_thread_interval_ms);
    }

    @Override
    protected void runAfterCatalogReady() {
        try {
            List<AuditEvent> auditEventList = getQueryNeedAudit();
            int missedLogCount = 0;
            int succLogCount = 0;
            for (AuditEvent auditEvent : auditEventList) {
                boolean ret = Env.getCurrentAuditEventProcessor().handleAuditEvent(auditEvent);
                if (!ret) {
                    missedLogCount++;
                } else {
                    succLogCount++;
                }
            }
            if (missedLogCount > 0) {
                LOG.warn("discard audit event because of log queue is full, discard num : {}, succ num : {}",
                        missedLogCount, succLogCount);
            }
        } catch (Throwable t) {
            LOG.warn("exception happens when handleAuditEvent, ", t);
        }

        clearReportTimeoutBeStatistics();
    }

    // After the query or insert finished, FE will not audit immediately, it will send an audit
    // event to this queue. And the worker thread will handle it. If the queue is full, the event
    // will be handled immediately and may miss some statistic info. So the statistic info of audit
    // event may be not accurate, but it can avoid the case that FE OOM because of too many audit
    // events in queue when QPS is high. The event will be logged directly if the queue is full.
    // And the worker thread will get an event from the queue and get the statistic info for this
    // event from queryStatisticsMap.
    public void submitFinishQueryToAudit(AuditEvent event) {
        submitFinishQueryToAudit(event, ImmutableSet.of());
    }

    public void submitFinishQueryToAudit(AuditEvent event, Set<Long> expectedBackendIds) {
        submitFinishQueryToAudit(event, expectedBackendIds, event.queryId, false);
    }

    public void submitFinishQueryToAudit(AuditEvent event, Set<Long> expectedBackendIds,
            String statisticsQueryId, boolean transfersInFlightOwnership) {
        queryAuditEventLogWriteLock();
        try {
            if (queryAuditEventList.size() > Config.audit_event_log_queue_size) {
                long now = System.currentTimeMillis();

                if (now - lastWarnTime >= 1000) {
                    lastWarnTime = now;
                    // if queryAuditEventList is full, we don't put the event to queryAuditEventList.
                    // so that the statistic info of this audit event will be ignored,
                    // and event will be logged directly.
                    LOG.warn("audit log event queue size {} is full, this may cause audit log missing "
                            + "statistics. you can check whether qps is too high or set "
                            + "audit_event_log_queue_size to a larger value in fe.conf. query id: {}",
                            queryAuditEventList.size(), event.queryId);
                }
                Env.getCurrentAuditEventProcessor().handleAuditEvent(event);
                if (transfersInFlightOwnership) {
                    releaseAuditStatisticsOwnership(statisticsQueryId);
                }
            } else {
                // put the event to queryAuditEventList and let the worker thread to handle it.
                // the worker thread will try best to wait for the statistic info before logging this event.
                event.pushToAuditLogQueueTime = System.currentTimeMillis();
                PendingAuditEvent pending = new PendingAuditEvent(
                        event, expectedBackendIds, statisticsQueryId);
                // Replace the execution owner with participant-scoped owners under one lock so an
                // accepted final snapshot cannot disappear between coordinator teardown and audit.
                registerPendingAuditEvent(pending, transfersInFlightOwnership);
                queryAuditEventList.add(pending);
            }
        } finally {
            queryAuditEventLogWriteUnlock();
        }
    }

    @VisibleForTesting
    List<AuditEvent> getQueryNeedAudit() {
        RuntimeStatisticsSnapshot runtimeSnapshot = buildRuntimeStatisticsSnapshot();
        queryStatisticsSnapshot = ImmutableMap.copyOf(runtimeSnapshot.queryStatistics);

        long currentTime = System.currentTimeMillis();
        int queryAuditLogTimeout = Config.query_audit_log_timeout_ms;
        long maximumWaitMs = Math.max(queryAuditLogTimeout,
                Config.be_report_query_statistics_timeout_ms);
        List<PendingAuditEvent> dueEvents = new ArrayList<>();
        // Keep only identity/timestamp traversal under the producer-facing queue lock; participant
        // lookup and statistics merging may touch every expected backend and must run lock-free.
        queryAuditEventLogWriteLock();
        try {
            for (PendingAuditEvent pending : queryAuditEventList) {
                long elapsed = currentTime - pending.event.pushToAuditLogQueueTime;
                if (elapsed <= queryAuditLogTimeout) {
                    // Wall-clock corrections can make insertion order disagree with deadlines.
                    continue;
                }
                dueEvents.add(pending);
            }
        } finally {
            queryAuditEventLogWriteUnlock();
        }

        Map<PendingAuditEvent, TQueryStatistics> readyEvents = new IdentityHashMap<>();
        for (PendingAuditEvent pending : dueEvents) {
            beforeHydrateAuditEvent(pending.event.queryId);
            boolean allFinalSnapshotsReceived = true;
            TQueryStatistics auditStatistics = pending.expectedBackendIds.length == 0
                    ? runtimeSnapshot.queryStatistics.get(pending.statisticsQueryId)
                    : new TQueryStatistics();
            for (int i = 0; i < pending.expectedBackendIds.length; i++) {
                bindParticipantFromExistingReport(pending, i);
                long backendStartTime = pending.expectedBackendStartTimes.get(i);
                TQueryStatisticsResult statistics = findStatisticsForBackend(
                        pending.statisticsQueryId, pending.expectedBackendIds[i], backendStartTime);
                if (statistics == null) {
                    allFinalSnapshotsReceived = false;
                    continue;
                }
                mergeQueryStatistics(auditStatistics, statistics);
                if (!statistics.isSetQueryFinished() || !statistics.isQueryFinished()) {
                    allFinalSnapshotsReceived = false;
                }
            }
            long elapsed = currentTime - pending.event.pushToAuditLogQueueTime;
            if (allFinalSnapshotsReceived || elapsed > maximumWaitMs) {
                readyEvents.put(pending, auditStatistics);
            }
        }

        List<PendingAuditEvent> removedEvents = new ArrayList<>();
        queryAuditEventLogWriteLock();
        try {
            for (PendingAuditEvent pending : dueEvents) {
                if (readyEvents.containsKey(pending) && queryAuditEventList.remove(pending)) {
                    removedEvents.add(pending);
                }
            }
        } finally {
            queryAuditEventLogWriteUnlock();
        }

        List<AuditEvent> ret = new ArrayList<>(removedEvents.size());
        for (PendingAuditEvent pending : removedEvents) {
            // A DML audit must not pass the queue until every scheduled BE has published its
            // query-level final snapshot. The upper bound only protects the queue if a BE dies.
            applyQueryStatisticsToAuditEvent(pending.event, readyEvents.get(pending));
            ret.add(pending.event);
            unregisterPendingAuditEvent(pending);
        }
        return ret;
    }

    public boolean updateBeQueryStats(TReportWorkloadRuntimeStatusParams params) {
        if (!params.isSetBackendId()) {
            LOG.warn("be report workload runtime status but without beid");
            return false;
        }
        if (!params.isSetQueryStatisticsResultMap()) {
            LOG.warn("be report workload runtime status but without query stats map");
            return false;
        }
        if (!params.isSetBackendStartTime()) {
            LOG.warn("be report workload runtime status without backend start time");
            return false;
        }
        long beId = params.backend_id;
        long backendStartTime = params.backend_start_time;
        long currentBackendStartTime = getBackendStartTime(beId);
        if (currentBackendStartTime > 0 && currentBackendStartTime != backendStartTime) {
            LOG.info("ignore stale workload runtime status from backend {}, report start time {}, current {}",
                    beId, backendStartTime, currentBackendStartTime);
            return false;
        }
        // Newly replayed backends use -1 as well as 0 for an unknown heartbeat epoch. In either
        // state the first accepted report must latch the only process allowed to publish data.
        if (currentBackendStartTime <= 0) {
            Long acceptedStartTime = lastAcceptedBackendStartTimes.putIfAbsent(beId, backendStartTime);
            if (acceptedStartTime != null && acceptedStartTime != backendStartTime) {
                return false;
            }
        } else {
            lastAcceptedBackendStartTimes.put(beId, backendStartTime);
        }
        BackendIncarnation incarnation = new BackendIncarnation(beId, backendStartTime);
        beToQueryStatsMap.computeIfAbsent(incarnation, ignored -> new BeReportInfo(backendStartTime));
        long currentTime = System.currentTimeMillis();
        for (Map.Entry<String, TQueryStatisticsResult> entry
                : params.query_statistics_result_map.entrySet()) {
            beforeUpdateQueryStatistics(entry.getKey());
            while (true) {
                BeReportInfo beReportInfo = beToQueryStatsMap.computeIfAbsent(incarnation,
                        ignored -> new BeReportInfo(backendStartTime));
                beReportInfo.lifecycleLock.readLock().lock();
                try {
                    // Cleanup removes only an empty shell under the write lock. If this insertion
                    // raced with removal, retry the single entry against the replacement shell.
                    if (beToQueryStatsMap.get(incarnation) != beReportInfo) {
                        continue;
                    }
                    Pair<Long, TQueryStatisticsResult> incoming =
                            Pair.of(currentTime, entry.getValue());
                    beReportInfo.queryStatsMap.compute(entry.getKey(), (queryId, previous) ->
                            previous == null
                                    || shouldReplaceQueryStatistics(previous.second, incoming.second)
                                            ? incoming : previous);
                    break;
                } finally {
                    beReportInfo.lifecycleLock.readLock().unlock();
                }
            }
        }
        // Heartbeat can change while a large report is being merged. Refuse the acknowledgement
        // so the BE retries against the FE state that is authoritative after the merge.
        long latestBackendStartTime = getBackendStartTime(beId);
        boolean accepted = latestBackendStartTime <= 0 || latestBackendStartTime == backendStartTime;
        if (accepted) {
            beforeBindPendingAuditEvents();
            bindPendingAuditEvents(beId, backendStartTime,
                    params.query_statistics_result_map.keySet());
        }
        return accepted;
    }

    private TQueryStatisticsResult findStatisticsForBackend(String queryId,
            long backendId, long backendStartTime) {
        if (backendStartTime <= 0) {
            return null;
        }
        BeReportInfo reportInfo = beToQueryStatsMap.get(
                new BackendIncarnation(backendId, backendStartTime));
        if (reportInfo != null) {
            Pair<Long, TQueryStatisticsResult> pair = reportInfo.queryStatsMap.get(queryId);
            return pair == null ? null : pair.second;
        }
        return null;
    }

    public void beginAuditStatisticsOwnership(String queryId) {
        if (queryId == null) {
            return;
        }
        pendingAuditBindingLock.lock();
        try {
            inFlightAuditQueryReferences.compute(queryId,
                    (ignored, count) -> count == null ? 1 : count + 1);
        } finally {
            pendingAuditBindingLock.unlock();
        }
    }

    public void transferAuditStatisticsOwnership(String previousQueryId, String nextQueryId) {
        if (previousQueryId == null || nextQueryId == null || previousQueryId.equals(nextQueryId)) {
            return;
        }
        pendingAuditBindingLock.lock();
        try {
            releaseAuditStatisticsOwnershipLocked(previousQueryId);
            inFlightAuditQueryReferences.compute(nextQueryId,
                    (ignored, count) -> count == null ? 1 : count + 1);
        } finally {
            pendingAuditBindingLock.unlock();
        }
    }

    public void releaseAuditStatisticsOwnership(String queryId) {
        if (queryId == null) {
            return;
        }
        pendingAuditBindingLock.lock();
        try {
            releaseAuditStatisticsOwnershipLocked(queryId);
        } finally {
            pendingAuditBindingLock.unlock();
        }
    }

    private void releaseAuditStatisticsOwnershipLocked(String queryId) {
        inFlightAuditQueryReferences.computeIfPresent(queryId,
                (ignored, count) -> count == 1 ? null : count - 1);
    }

    private void registerPendingAuditEvent(PendingAuditEvent pending,
            boolean transfersInFlightOwnership) {
        pendingAuditBindingLock.lock();
        try {
            if (!pending.active.get()) {
                return;
            }
            for (int i = 0; i < pending.expectedBackendIds.length; i++) {
                if (bindParticipantFromExistingReportLocked(pending, i)) {
                    continue;
                }
                BackendQuery key = new BackendQuery(
                        pending.expectedBackendIds[i], pending.statisticsQueryId);
                unboundAuditParticipants.computeIfAbsent(key, ignored -> new ArrayList<>())
                        .add(new PendingAuditBinding(pending, i));
            }
            if (transfersInFlightOwnership) {
                releaseAuditStatisticsOwnershipLocked(pending.statisticsQueryId);
            }
        } finally {
            pendingAuditBindingLock.unlock();
        }
    }

    private void bindPendingAuditEvents(long backendId, long backendStartTime,
            Set<String> reportedQueryIds) {
        for (String queryId : reportedQueryIds) {
            pendingAuditBindingLock.lock();
            try {
                BackendQuery key = new BackendQuery(backendId, queryId);
                List<PendingAuditBinding> bindings = unboundAuditParticipants.remove(key);
                if (bindings == null) {
                    continue;
                }
                for (PendingAuditBinding binding : bindings) {
                    bindParticipantLocked(binding.pending, binding.participantIndex,
                            backendStartTime);
                }
            } finally {
                pendingAuditBindingLock.unlock();
            }
        }
    }

    private void bindParticipantFromExistingReport(PendingAuditEvent pending, int participantIndex) {
        if (pending.expectedBackendStartTimes.get(participantIndex) > 0) {
            return;
        }
        pendingAuditBindingLock.lock();
        try {
            if (bindParticipantFromExistingReportLocked(pending, participantIndex)) {
                removeUnboundParticipantLocked(pending, participantIndex);
            }
        } finally {
            pendingAuditBindingLock.unlock();
        }
    }

    private boolean bindParticipantFromExistingReportLocked(
            PendingAuditEvent pending, int participantIndex) {
        if (!pending.active.get()) {
            return false;
        }
        if (pending.expectedBackendStartTimes.get(participantIndex) > 0) {
            return true;
        }
        long backendId = pending.expectedBackendIds[participantIndex];
        long currentBackendStartTime = getBackendStartTime(backendId);
        long reportStartTime = currentBackendStartTime > 0
                ? currentBackendStartTime
                : lastAcceptedBackendStartTimes.getOrDefault(backendId, 0L);
        if (findStatisticsForBackend(pending.statisticsQueryId, backendId, reportStartTime) == null) {
            return false;
        }
        bindParticipantLocked(pending, participantIndex, reportStartTime);
        return true;
    }

    private void bindParticipantLocked(PendingAuditEvent pending, int participantIndex,
            long backendStartTime) {
        if (!pending.active.get()
                || !pending.expectedBackendStartTimes.compareAndSet(
                        participantIndex, 0, backendStartTime)) {
            return;
        }
        // Retention is query-scoped so one queued audit cannot pin unrelated completed queries
        // from the same long-lived backend process.
        RetainedQuery retainedQuery = new RetainedQuery(
                new BackendIncarnation(pending.expectedBackendIds[participantIndex],
                        backendStartTime), pending.statisticsQueryId);
        pendingQueryReferences.compute(retainedQuery,
                (ignored, count) -> count == null ? 1 : count + 1);
    }

    private void unregisterPendingAuditEvent(PendingAuditEvent pending) {
        pendingAuditBindingLock.lock();
        try {
            if (!pending.active.compareAndSet(true, false)) {
                return;
            }
            for (int i = 0; i < pending.expectedBackendIds.length; i++) {
                long backendStartTime = pending.expectedBackendStartTimes.get(i);
                if (backendStartTime > 0) {
                    RetainedQuery retainedQuery = new RetainedQuery(
                            new BackendIncarnation(pending.expectedBackendIds[i], backendStartTime),
                            pending.statisticsQueryId);
                    pendingQueryReferences.computeIfPresent(retainedQuery, (ignored, count) ->
                            count == 1 ? null : count - 1);
                } else {
                    removeUnboundParticipantLocked(pending, i);
                }
            }
        } finally {
            pendingAuditBindingLock.unlock();
        }
    }

    private void removeUnboundParticipantLocked(PendingAuditEvent pending, int participantIndex) {
        BackendQuery key = new BackendQuery(
                pending.expectedBackendIds[participantIndex], pending.statisticsQueryId);
        List<PendingAuditBinding> bindings = unboundAuditParticipants.get(key);
        if (bindings == null) {
            return;
        }
        bindings.removeIf(binding -> binding.pending == pending
                && binding.participantIndex == participantIndex);
        if (bindings.isEmpty()) {
            unboundAuditParticipants.remove(key);
        }
    }

    @VisibleForTesting
    void beforeHydrateAuditEvent(String queryId) {
    }

    @VisibleForTesting
    void beforeBindPendingAuditEvents() {
    }

    private boolean isQueryRetainedForAudit(BackendIncarnation incarnation, String queryId) {
        RetainedQuery retainedQuery = new RetainedQuery(incarnation, queryId);
        if (pendingQueryReferences.containsKey(retainedQuery)) {
            return true;
        }
        pendingAuditBindingLock.lock();
        try {
            // Serialize cleanup with first-report binding. Until the participant has an epoch,
            // retaining only this backend/query candidate closes the insert-to-bind race.
            return pendingQueryReferences.containsKey(retainedQuery)
                    || inFlightAuditQueryReferences.containsKey(queryId)
                    || unboundAuditParticipants.containsKey(
                            new BackendQuery(incarnation.backendId, queryId));
        } finally {
            pendingAuditBindingLock.unlock();
        }
    }

    @VisibleForTesting
    long getBackendStartTime(long backendId) {
        Backend backend = Env.getCurrentSystemInfo().getBackend(backendId);
        return backend == null ? 0 : backend.getLastStartTime();
    }

    @VisibleForTesting
    void beforeUpdateQueryStatistics(String queryId) {
    }

    private boolean shouldReplaceQueryStatistics(TQueryStatisticsResult previous,
            TQueryStatisticsResult incoming) {
        long previousGeneration = previous.isSetQueryStatisticsGeneration()
                ? previous.getQueryStatisticsGeneration() : 0;
        long incomingGeneration = incoming.isSetQueryStatisticsGeneration()
                ? incoming.getQueryStatisticsGeneration() : 0;
        if (previousGeneration != incomingGeneration) {
            return incomingGeneration > previousGeneration;
        }

        long previousSequence = previous.isSetQueryStatisticsSequence()
                ? previous.getQueryStatisticsSequence() : 0;
        long incomingSequence = incoming.isSetQueryStatisticsSequence()
                ? incoming.getQueryStatisticsSequence() : 0;
        if (previousSequence != incomingSequence) {
            return incomingSequence > previousSequence;
        }

        return false;
    }

    void clearReportTimeoutBeStatistics() {
        Set<BackendIncarnation> currentBeIdSet = beToQueryStatsMap.keySet();
        Long currentTime = System.currentTimeMillis();
        for (BackendIncarnation incarnation : currentBeIdSet) {
            BeReportInfo beReportInfo = beToQueryStatsMap.get(incarnation);
            if (beReportInfo == null) {
                continue;
            }
            Set<String> queryIdSet = beReportInfo.queryStatsMap.keySet();
            for (String queryId : queryIdSet) {
                beforeExpireQueryStatistics(queryId);
                beReportInfo.queryStatsMap.computeIfPresent(queryId, (ignoredQueryId, pair) -> {
                    long queryLastReportTime = pair.first;
                    boolean timeout = currentTime - queryLastReportTime
                            > Config.be_report_query_statistics_timeout_ms;
                    // Conditional removal and updates serialize only for the same query. Work on
                    // a high-cardinality report cannot delay another query's terminal update.
                    return timeout && !isQueryRetainedForAudit(incarnation, queryId)
                            && isQueryNotExistInFe(queryId) ? null : pair;
                });
            }
            beReportInfo.lifecycleLock.writeLock().lock();
            try {
                if (beReportInfo.queryStatsMap.isEmpty()) {
                    beToQueryStatsMap.remove(incarnation, beReportInfo);
                }
            } finally {
                beReportInfo.lifecycleLock.writeLock().unlock();
            }
        }
    }

    @VisibleForTesting
    void beforeExpireQueryStatistics(String queryId) {
    }

    boolean isQueryNotExistInFe(String queryId) {
        try {
            return QeProcessorImpl.INSTANCE.getCoordinator(DebugUtil.parseTUniqueIdFromString(queryId)) == null;
        } catch (NumberFormatException e) {
            return true;
        }
    }

    // Rebuild query statistics from concurrent runtime maps and publish an immutable snapshot.
    // This method is intentionally called by daemon thread and unit tests only.
    void rebuildQueryStatisticsSnapshot() {
        queryStatisticsSnapshot = ImmutableMap.copyOf(buildRuntimeStatisticsSnapshot().queryStatistics);
    }

    // Return the latest published snapshot for synchronous readers such as proc/REST paths.
    public Map<String, TQueryStatistics> getQueryStatisticsMap() {
        return queryStatisticsSnapshot;
    }

    // Build a merged map by traversing concurrent runtime structures.
    private RuntimeStatisticsSnapshot buildRuntimeStatisticsSnapshot() {
        RuntimeStatisticsSnapshot snapshot = new RuntimeStatisticsSnapshot();
        for (Map.Entry<BackendIncarnation, BeReportInfo> beEntry : beToQueryStatsMap.entrySet()) {
            BackendIncarnation incarnation = beEntry.getKey();
            long beId = incarnation.backendId;
            BeReportInfo beReportInfo = beEntry.getValue();
            if (beReportInfo == null) {
                continue;
            }
            long currentBackendStartTime = getBackendStartTime(beId);
            long authoritativeStartTime = currentBackendStartTime > 0
                    ? currentBackendStartTime
                    : lastAcceptedBackendStartTimes.getOrDefault(beId, 0L);
            boolean isCurrentIncarnation = authoritativeStartTime == 0
                    || authoritativeStartTime == beReportInfo.backendStartTime;
            if (isCurrentIncarnation) {
                beforeBuildBackendStatistics(beId);
            }
            for (Map.Entry<String, Pair<Long, TQueryStatisticsResult>> queryEntry
                    : beReportInfo.queryStatsMap.entrySet()) {
                String queryId = queryEntry.getKey();
                Pair<Long, TQueryStatisticsResult> queryStatsPair = queryEntry.getValue();
                if (queryStatsPair == null || queryStatsPair.second == null) {
                    continue;
                }
                TQueryStatisticsResult curQueryStats = queryStatsPair.second;
                if (!isCurrentIncarnation) {
                    // Old process data stays available to audits already bound to that incarnation,
                    // but it must never enter the current public snapshot.
                    continue;
                }
                TQueryStatistics retQuery = snapshot.queryStatistics.computeIfAbsent(
                        queryId, ignored -> new TQueryStatistics());
                mergeQueryStatistics(retQuery, curQueryStats);
            }
        }
        return snapshot;
    }

    @VisibleForTesting
    void beforeBuildBackendStatistics(long backendId) {
    }

    private void applyQueryStatisticsToAuditEvent(AuditEvent auditEvent, TQueryStatistics queryStats) {
        if (queryStats == null) {
            return;
        }
        auditEvent.scanRows = queryStats.scan_rows;
        auditEvent.scanBytes = queryStats.scan_bytes;
        auditEvent.scanBytesFromLocalStorage = queryStats.scan_bytes_from_local_storage;
        auditEvent.scanBytesFromRemoteStorage = queryStats.scan_bytes_from_remote_storage;
        auditEvent.peakMemoryBytes = queryStats.max_peak_memory_bytes;
        auditEvent.cpuTimeMs = queryStats.cpu_ms;
        auditEvent.shuffleSendBytes = queryStats.shuffle_send_bytes;
        auditEvent.shuffleSendRows = queryStats.shuffle_send_rows;
        auditEvent.spillWriteBytesToLocalStorage = queryStats.spill_write_bytes_to_local_storage;
        auditEvent.spillReadBytesFromLocalStorage = queryStats.spill_read_bytes_from_local_storage;
    }

    public Map<Long, TQueryStatisticsResult> getQueryStatistics(String queryId) {
        Map<Long, TQueryStatisticsResult> result = Maps.newHashMap();
        for (Map.Entry<BackendIncarnation, BeReportInfo> entry : beToQueryStatsMap.entrySet()) {
            long currentBackendStartTime = getBackendStartTime(entry.getKey().backendId);
            long authoritativeStartTime = currentBackendStartTime > 0
                    ? currentBackendStartTime
                    : lastAcceptedBackendStartTimes.getOrDefault(entry.getKey().backendId, 0L);
            if (authoritativeStartTime > 0
                    && authoritativeStartTime != entry.getKey().backendStartTime) {
                continue;
            }
            Pair<Long, TQueryStatisticsResult> pair = entry.getValue().queryStatsMap.get(queryId);
            if (pair != null) {
                result.put(entry.getKey().backendId, pair.second);
            }
        }
        return result;
    }


    private void mergeQueryStatistics(TQueryStatistics dst, TQueryStatisticsResult src) {
        TQueryStatistics srcStats = src.getStatistics();
        if (srcStats == null) {
            return;
        }
        dst.setScanRows(dst.scan_rows + srcStats.scan_rows);
        dst.setScanBytes(dst.scan_bytes + srcStats.scan_bytes);
        dst.setScanBytesFromLocalStorage(dst.scan_bytes_from_local_storage
                + srcStats.scan_bytes_from_local_storage);
        dst.setScanBytesFromRemoteStorage(dst.scan_bytes_from_remote_storage
                + srcStats.scan_bytes_from_remote_storage);
        dst.setCpuMs(dst.cpu_ms + srcStats.cpu_ms);
        dst.setShuffleSendBytes(dst.shuffle_send_bytes + srcStats.shuffle_send_bytes);
        dst.setShuffleSendRows(dst.shuffle_send_rows + srcStats.shuffle_send_rows);
        dst.setProcessRows(dst.process_rows + srcStats.process_rows);
        dst.setReturnedRows(dst.returned_rows + srcStats.returned_rows);
        if (srcStats.isSetTotalTasksNum()) {
            dst.setTotalTasksNum(dst.total_tasks_num + srcStats.total_tasks_num);
        }
        if (srcStats.isSetFinishedTasksNum()) {
            dst.setFinishedTasksNum(dst.finished_tasks_num + srcStats.finished_tasks_num);
        }
        if (dst.current_used_memory_bytes < srcStats.current_used_memory_bytes) {
            dst.setCurrentUsedMemoryBytes(srcStats.current_used_memory_bytes);
        }
        if (dst.workload_group_id <= 0 && srcStats.workload_group_id > 0) {
            dst.setWorkloadGroupId(srcStats.workload_group_id);
        }
        if (dst.max_peak_memory_bytes < srcStats.max_peak_memory_bytes) {
            dst.setMaxPeakMemoryBytes(srcStats.max_peak_memory_bytes);
        }
        dst.setSpillWriteBytesToLocalStorage(dst.spill_write_bytes_to_local_storage
                + srcStats.spill_write_bytes_to_local_storage);
        dst.setSpillReadBytesFromLocalStorage(dst.spill_read_bytes_from_local_storage
                + srcStats.spill_read_bytes_from_local_storage);
        dst.setBytesWriteIntoCache(dst.bytes_write_into_cache + srcStats.bytes_write_into_cache);
    }

    private void queryAuditEventLogWriteLock() {
        queryAuditEventLock.lock();
    }

    private void queryAuditEventLogWriteUnlock() {
        queryAuditEventLock.unlock();
    }

}