risingwave_storage/hummock/vector/
monitor.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::cell::RefCell;
16use std::collections::HashMap;
17
18use risingwave_common::catalog::TableId;
19use risingwave_common::metrics::{LabelGuardedHistogram, LabelGuardedIntCounter};
20
21use crate::monitor::HummockStateStoreMetrics;
22use crate::vector::hnsw::HnswStats;
23
24macro_rules! for_all_cache_counter {
25    ($($name:ident),+) => {
26        struct VectorStoreCacheMetrics {
27            $(
28                $name: LabelGuardedIntCounter,
29            )+
30        }
31
32        impl VectorStoreCacheMetrics {
33
34            fn new(metrics: &HummockStateStoreMetrics, table_id: TableId, mode: &str) -> Self {
35                let table_id_label = format!("{}", table_id);
36                $(
37                    let $name = metrics
38                        .vector_object_request_counts
39                        .with_guarded_label_values(&[table_id_label.as_str(), stringify!($name), mode]);
40                )+
41
42                Self {
43                    $($name,)+
44                }
45            }
46
47            fn report(&self, stat: VectorStoreCacheStats) {
48                $(
49                    self.$name.inc_by(stat.$name);
50                )+
51            }
52        }
53
54        #[derive(Default)]
55        pub struct VectorStoreCacheStats {
56            $(pub $name: u64,)+
57        }
58    };
59    () => {
60        for_all_cache_counter! {
61            file_block_total, file_block_miss, file_meta_total, file_meta_miss, hnsw_graph_total, hnsw_graph_miss
62        }
63    }
64}
65
66for_all_cache_counter!();
67
68impl VectorStoreCacheStats {
69    pub fn report(self, table_id: TableId, mode: &'static str, metrics: &HummockStateStoreMetrics) {
70        type StatsType = HashMap<(TableId, &'static str), VectorStoreCacheMetrics>;
71        thread_local! {
72            static THREAD_STATS: RefCell<StatsType> = RefCell::new(HashMap::new());
73        }
74
75        THREAD_STATS.with_borrow_mut(move |global| {
76            let metrics = global
77                .entry((table_id, mode))
78                .or_insert_with(|| VectorStoreCacheMetrics::new(metrics, table_id, mode));
79            metrics.report(self)
80        });
81    }
82}
83
84struct HnswMetrics {
85    distances_computed: LabelGuardedHistogram,
86    n_hops: LabelGuardedHistogram,
87}
88
89impl HnswMetrics {
90    fn new(
91        metrics: &HummockStateStoreMetrics,
92        table_id: TableId,
93        mode: &'static str,
94        top_n: usize,
95        ef: usize,
96    ) -> Self {
97        let table_id_label = format!("{}", table_id);
98        let top_n_label = format!("{}", top_n);
99        let ef_label = format!("{}", ef);
100        let distances_computed = metrics.vector_request_stats.with_guarded_label_values(&[
101            table_id_label.as_str(),
102            "distances_computed",
103            mode,
104            top_n_label.as_str(),
105            ef_label.as_str(),
106        ]);
107        let n_hops = metrics.vector_request_stats.with_guarded_label_values(&[
108            table_id_label.as_str(),
109            "n_hops",
110            mode,
111            top_n_label.as_str(),
112            ef_label.as_str(),
113        ]);
114        Self {
115            distances_computed,
116            n_hops,
117        }
118    }
119}
120
121pub fn report_hnsw_stat(
122    metrics: &HummockStateStoreMetrics,
123    table_id: TableId,
124    mode: &'static str,
125    top_n: usize,
126    ef: usize,
127    stats: impl IntoIterator<Item = HnswStats>,
128) {
129    type StatsType = HashMap<(TableId, &'static str, usize, usize), HnswMetrics>;
130    thread_local! {
131        static THREAD_STATS: RefCell<StatsType> = RefCell::new(HashMap::new());
132    }
133
134    THREAD_STATS.with_borrow_mut(move |global| {
135        let metrics = global
136            .entry((table_id, mode, top_n, ef))
137            .or_insert_with(|| HnswMetrics::new(metrics, table_id, mode, top_n, ef));
138        let distance_computed = metrics.distances_computed.local();
139        let n_hop = metrics.n_hops.local();
140        for stat in stats {
141            distance_computed.observe(stat.distances_computed as _);
142            n_hop.observe(stat.n_hops as _);
143        }
144    });
145}