// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 use std::path::PathBuf; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, OnceLock}; use anyhow::Context; use dynamo_kv_router::indexer::{KvIndexerMetrics, METRIC_WARNING_DUPLICATE_STORE}; use tracing::{Event, Level, Subscriber}; use tracing_subscriber::Registry; use tracing_subscriber::layer::{Context as LayerContext, Layer}; use tracing_subscriber::prelude::*; struct WarningCounterLayer { count: Arc, target_prefixes: &'static [&'static str], } impl Layer for WarningCounterLayer where S: Subscriber, { fn on_event(&self, event: &Event<'_>, _ctx: LayerContext<'_, S>) { let metadata = event.metadata(); let target = metadata.target(); if matches!(*metadata.level(), Level::WARN | Level::ERROR) && self .target_prefixes .iter() .any(|prefix| target.starts_with(prefix)) { self.count.fetch_add(1, Ordering::Relaxed); } } } pub fn warning_counter(target_prefixes: &'static [&'static str]) -> Arc { static COUNTER: OnceLock> = OnceLock::new(); COUNTER .get_or_init(|| { let count = Arc::new(AtomicUsize::new(0)); let subscriber = Registry::default().with(WarningCounterLayer { count: Arc::clone(&count), target_prefixes, }); tracing::subscriber::set_global_default(subscriber) .expect("global warning counter subscriber should initialize once"); count }) .clone() } pub fn reset_warning_count(counter: &Arc) { counter.store(0, Ordering::Relaxed); } pub fn fixture_path(file_name: &str) -> anyhow::Result { PathBuf::from(env!("CARGO_MANIFEST_DIR")) .join("testdata") .join(file_name) .to_str() .map(str::to_owned) .context("fixture path is not valid UTF-8") } #[allow(dead_code)] pub fn duplicate_store_warning_count(metrics: &KvIndexerMetrics) -> u64 { metrics .kv_cache_event_warnings .get_metric_with_label_values(&[METRIC_WARNING_DUPLICATE_STORE]) .expect("duplicate_store warning metric should exist") .get() }