risingwave_storage/monitor/
hitmap.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::sync::Arc;
16use std::sync::atomic::{AtomicU64, Ordering};
17
18use risingwave_common::util::iter_util::ZipEqFast;
19
20#[derive(Debug, Clone)]
21pub struct Hitmap<const N: usize> {
22    /// For [`Hitmap`] is rarely access in multi-thread pattern,
23    /// the cons of false-sharing can be ignored.
24    data: Arc<[AtomicU64; N]>,
25}
26
27impl<const N: usize> Default for Hitmap<N> {
28    fn default() -> Self {
29        let data = [(); N].map(|_| AtomicU64::default());
30        let data = Arc::new(data);
31        Self { data }
32    }
33}
34
35impl<const N: usize> Hitmap<N> {
36    pub const fn bits() -> usize {
37        N * u64::BITS as usize
38    }
39
40    pub const fn bytes() -> usize {
41        N * 8
42    }
43
44    pub fn report(&self, local: &mut LocalHitmap<N>) {
45        for (global, local) in self.data.iter().zip_eq_fast(local.data.iter()) {
46            global.fetch_or(*local, Ordering::Relaxed);
47        }
48        local.reset();
49    }
50
51    pub fn ones(&self) -> usize {
52        let mut res = 0;
53        for elem in &*self.data {
54            res += elem.load(Ordering::Relaxed).count_ones() as usize;
55        }
56        res
57    }
58
59    pub fn zeros(&self) -> usize {
60        Self::bits() - self.ones()
61    }
62
63    pub fn ratio(&self) -> f64 {
64        self.ones() as f64 / Self::bits() as f64
65    }
66
67    #[cfg(test)]
68    pub fn to_hex_vec(&self) -> Vec<String> {
69        use itertools::Itertools;
70        self.data
71            .iter()
72            .map(|elem| elem.load(Ordering::Relaxed))
73            .map(|v| format!("{v:016x}"))
74            .collect_vec()
75    }
76}
77
78#[derive(Debug)]
79pub struct LocalHitmap<const N: usize> {
80    data: Box<[u64; N]>,
81}
82
83impl<const N: usize> Default for LocalHitmap<N> {
84    fn default() -> Self {
85        Self::new()
86    }
87}
88
89impl<const N: usize> LocalHitmap<N> {
90    pub const fn bits() -> usize {
91        N * u64::BITS as usize
92    }
93
94    pub const fn bytes() -> usize {
95        N * 8
96    }
97
98    pub fn new() -> Self {
99        Self {
100            data: Box::new([0; N]),
101        }
102    }
103
104    pub fn reset(&mut self) {
105        for elem in &mut *self.data {
106            *elem = 0;
107        }
108    }
109
110    pub fn merge(&mut self, other: &mut Self) {
111        for (elem, e) in self.data.iter_mut().zip_eq_fast(other.data.iter()) {
112            *elem |= *e;
113        }
114        other.reset();
115    }
116
117    pub fn fill(&mut self, start_bit: usize, end_bit: usize) {
118        const MASK: usize = (1 << 6) - 1;
119
120        let end_bit = end_bit.clamp(start_bit + 1, Self::bits());
121
122        let head_bits = start_bit & MASK;
123        let tail_bits_rev = end_bit & MASK;
124
125        let head_elem = start_bit >> 6;
126        let tail_elem = end_bit >> 6;
127
128        for i in head_elem..=std::cmp::min(tail_elem, N - 1) {
129            let elem = &mut self.data[i];
130            let mut umask = 0u64;
131            if i == head_elem {
132                umask |= (1u64 << head_bits) - 1;
133            }
134            if i == tail_elem {
135                umask |= !((1u64 << tail_bits_rev) - 1);
136            }
137            *elem |= !umask;
138        }
139    }
140
141    pub fn fill_with_range(&mut self, start: usize, end: usize, len: usize) {
142        let start_bit = Self::bits() * start / len;
143        let end_bit = Self::bits() * end / len;
144        self.fill(start_bit, end_bit)
145    }
146
147    pub fn ones(&self) -> usize {
148        let mut res = 0;
149        for elem in &*self.data {
150            res += elem.count_ones() as usize;
151        }
152        res
153    }
154
155    pub fn zeros(&self) -> usize {
156        Self::bits() - self.ones()
157    }
158
159    pub fn ratio(&self) -> f64 {
160        self.ones() as f64 / Self::bits() as f64
161    }
162
163    #[cfg(test)]
164    pub fn to_hex_vec(&self) -> Vec<String> {
165        use itertools::Itertools;
166        self.data.iter().map(|v| format!("{v:016x}")).collect_vec()
167    }
168}
169
170#[cfg(debug_assertions)]
171impl<const N: usize> Drop for LocalHitmap<N> {
172    fn drop(&mut self) {
173        if self.ones() > 0 {
174            panic!("LocalHitmap is not reported!");
175        }
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182
183    #[test]
184    fn test_hitmap() {
185        // hex: high <== low
186        let g = Hitmap::<4>::default();
187
188        let mut h = LocalHitmap::new();
189        assert_eq!(
190            h.to_hex_vec(),
191            vec![
192                "0000000000000000",
193                "0000000000000000",
194                "0000000000000000",
195                "0000000000000000",
196            ]
197        );
198        assert_eq!(h.ones(), 0);
199        h.fill(16, 24);
200        assert_eq!(
201            h.to_hex_vec(),
202            vec![
203                "0000000000ff0000",
204                "0000000000000000",
205                "0000000000000000",
206                "0000000000000000",
207            ]
208        );
209        assert_eq!(h.ones(), 8);
210        h.fill(32, 64);
211        assert_eq!(
212            h.to_hex_vec(),
213            vec![
214                "ffffffff00ff0000",
215                "0000000000000000",
216                "0000000000000000",
217                "0000000000000000",
218            ]
219        );
220        assert_eq!(h.ones(), 40);
221        h.fill(96, 224);
222        assert_eq!(
223            h.to_hex_vec(),
224            vec![
225                "ffffffff00ff0000",
226                "ffffffff00000000",
227                "ffffffffffffffff",
228                "00000000ffffffff",
229            ]
230        );
231        assert_eq!(h.ones(), 168);
232        h.fill(0, 256);
233        assert_eq!(
234            h.to_hex_vec(),
235            vec![
236                "ffffffffffffffff",
237                "ffffffffffffffff",
238                "ffffffffffffffff",
239                "ffffffffffffffff",
240            ]
241        );
242        assert_eq!(h.ones(), 256);
243        g.report(&mut h);
244        assert_eq!(
245            g.to_hex_vec(),
246            vec![
247                "ffffffffffffffff",
248                "ffffffffffffffff",
249                "ffffffffffffffff",
250                "ffffffffffffffff",
251            ]
252        );
253        assert_eq!(g.ones(), 256);
254    }
255}
risingwave_storage/monitor/hitmap.rs

risingwave_storage/monitor/
hitmap.rs
