risingwave_storage/monitor/
monitored_store.rs

1// Copyright 2022 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, HashSet};
17use std::marker::PhantomData;
18use std::sync::Arc;
19
20use await_tree::{InstrumentAwait, SpanExt};
21use bytes::Bytes;
22use futures::future::BoxFuture;
23use futures::{Future, FutureExt, TryFutureExt};
24use risingwave_common::bitmap::Bitmap;
25use risingwave_common::catalog::TableId;
26use risingwave_common::hash::VirtualNode;
27use risingwave_common::metrics::LabelGuardedHistogram;
28use risingwave_common::types::VectorRef;
29use risingwave_hummock_sdk::key::{TableKey, TableKeyRange};
30use risingwave_hummock_sdk::{HummockEpoch, HummockReadEpoch, SyncResult};
31use thiserror_ext::AsReport;
32use tokio::time::Instant;
33use tracing::{Instrument, error};
34
35use super::{MonitoredStateStoreGetStats, MonitoredStateStoreIterStats, MonitoredStorageMetrics};
36use crate::error::StorageResult;
37use crate::hummock::sstable_store::SstableStoreRef;
38use crate::hummock::{HummockStorage, ObjectIdManagerRef};
39use crate::monitor::monitored_storage_metrics::StateStoreIterStats;
40use crate::monitor::{StateStoreIterLogStats, StateStoreIterStatsTrait};
41use crate::store::*;
42use crate::store_impl::AsHummock;
43
44/// A state store wrapper for monitoring metrics.
45#[derive(Clone)]
46pub struct MonitoredStateStore<S, E = ()> {
47    inner: Box<S>,
48    storage_metrics: Arc<MonitoredStorageMetrics>,
49    extra: E,
50}
51
52type MonitoredTableStateStore<S> = MonitoredStateStore<S, TableId>;
53
54impl<S> MonitoredStateStore<S> {
55    pub fn new(inner: S, storage_metrics: Arc<MonitoredStorageMetrics>) -> Self {
56        Self {
57            inner: Box::new(inner),
58            storage_metrics,
59            extra: (),
60        }
61    }
62
63    pub fn inner(&self) -> &S {
64        &self.inner
65    }
66}
67
68impl<S> MonitoredTableStateStore<S> {
69    fn new(inner: S, storage_metrics: Arc<MonitoredStorageMetrics>, table_id: TableId) -> Self {
70        Self {
71            inner: Box::new(inner),
72            storage_metrics,
73            extra: table_id,
74        }
75    }
76
77    fn table_id(&self) -> TableId {
78        self.extra
79    }
80}
81
82// Note: it is important to define the `MonitoredStateStoreIter` type alias, as it marks that
83// the return type of `monitored_iter` only captures the lifetime `'s` and has nothing to do with
84// `'a`. If we simply use `impl StateStoreIter + 's`, the rust compiler will also capture
85// the lifetime `'a` in the scope defined in the scope.
86impl<S, E> MonitoredStateStore<S, E> {
87    async fn monitored_iter<
88        'a,
89        Item: IterItem,
90        I: StateStoreIter<Item> + 'a,
91        Stat: StateStoreIterStatsTrait<Item = Item>,
92    >(
93        &'a self,
94        table_id: TableId,
95        iter_stream_future: impl Future<Output = StorageResult<I>> + 'a,
96    ) -> StorageResult<MonitoredStateStoreIter<Item, I, Stat>> {
97        // start time takes iterator build time into account
98        // wait for iterator creation (e.g. seek)
99        let start_time = Instant::now();
100        let iter_stream = iter_stream_future
101            .await
102            .inspect_err(|e| error!(error = %e.as_report(), "Failed in iter"))?;
103        let iter_init_duration = start_time.elapsed();
104
105        // create a monitored iterator to collect metrics
106        let monitored = MonitoredStateStoreIter {
107            inner: iter_stream,
108            stats: MonitoredStateStoreIterStats {
109                inner: Stat::new(table_id, &self.storage_metrics, iter_init_duration),
110                table_id,
111                metrics: self.storage_metrics.clone(),
112            },
113            _phantom: PhantomData,
114        };
115        Ok(monitored)
116    }
117
118    async fn monitored_on_key_value<O>(
119        &self,
120        on_key_value_future: impl Future<Output = StorageResult<Option<(O, usize)>>>,
121        table_id: TableId,
122        key_len: usize,
123    ) -> StorageResult<Option<O>> {
124        let mut stats = MonitoredStateStoreGetStats::new(table_id, self.storage_metrics.clone());
125
126        let value = on_key_value_future
127            .instrument_await("store_on_key_value".verbose())
128            .instrument(tracing::trace_span!("store_on_key_value"))
129            .await
130            .inspect_err(|e| error!(error = %e.as_report(), "Failed in get"))?;
131
132        stats.get_key_size = key_len;
133        let value = value.map(|(value, value_len)| {
134            stats.get_value_size = value_len;
135            value
136        });
137        stats.report();
138
139        Ok(value)
140    }
141}
142
143impl<S: StateStoreGet> StateStoreGet for MonitoredTableStateStore<S> {
144    fn on_key_value<'a, O: Send + 'a>(
145        &'a self,
146        key: TableKey<Bytes>,
147        read_options: ReadOptions,
148        on_key_value_fn: impl KeyValueFn<'a, O>,
149    ) -> impl StorageFuture<'a, Option<O>> {
150        let table_id = self.table_id();
151        let key_len = key.len();
152        self.monitored_on_key_value(
153            self.inner
154                .on_key_value(key, read_options, move |key, value| {
155                    let result = on_key_value_fn(key, value);
156                    result.map(|output| (output, value.len()))
157                }),
158            table_id,
159            key_len,
160        )
161    }
162}
163
164impl<S: StateStoreRead> StateStoreRead for MonitoredTableStateStore<S> {
165    type Iter = impl StateStoreReadIter;
166    type RevIter = impl StateStoreReadIter;
167
168    fn read_epoch(&self) -> Option<u64> {
169        self.inner.read_epoch()
170    }
171
172    fn iter(
173        &self,
174        key_range: TableKeyRange,
175        read_options: ReadOptions,
176    ) -> impl Future<Output = StorageResult<Self::Iter>> + '_ {
177        self.monitored_iter::<'_, _, _, StateStoreIterStats>(
178            self.table_id(),
179            self.inner.iter(key_range, read_options),
180        )
181    }
182
183    fn rev_iter(
184        &self,
185        key_range: TableKeyRange,
186        read_options: ReadOptions,
187    ) -> impl Future<Output = StorageResult<Self::RevIter>> + '_ {
188        self.monitored_iter::<'_, _, _, StateStoreIterStats>(
189            self.table_id(),
190            self.inner.rev_iter(key_range, read_options),
191        )
192    }
193}
194
195impl<S: StateStoreReadLog> StateStoreReadLog for MonitoredStateStore<S> {
196    type ChangeLogIter = impl StateStoreReadChangeLogIter;
197
198    fn next_epoch(&self, epoch: u64, options: NextEpochOptions) -> impl StorageFuture<'_, u64> {
199        self.inner.next_epoch(epoch, options)
200    }
201
202    fn iter_log(
203        &self,
204        epoch_range: (u64, u64),
205        key_range: TableKeyRange,
206        options: ReadLogOptions,
207    ) -> impl Future<Output = StorageResult<Self::ChangeLogIter>> + Send + '_ {
208        self.monitored_iter::<'_, _, _, StateStoreIterLogStats>(
209            options.table_id,
210            self.inner.iter_log(epoch_range, key_range, options),
211        )
212    }
213}
214
215impl<S: StateStoreReadVector> StateStoreReadVector for MonitoredTableStateStore<S> {
216    fn nearest<'a, O: Send + 'a>(
217        &'a self,
218        vec: VectorRef<'a>,
219        options: VectorNearestOptions,
220        on_nearest_item_fn: impl OnNearestItemFn<'a, O>,
221    ) -> impl StorageFuture<'a, Vec<O>> {
222        thread_local! {
223            static THREAD_HISTOGRAM_VEC: RefCell<HashMap<(TableId, usize, usize), LabelGuardedHistogram>> = RefCell::new(HashMap::new());
224        }
225        let start_time = Instant::now();
226        let metric_key = (self.table_id(), options.top_n, options.hnsw_ef_search);
227
228        self.inner
229            .nearest(vec, options, on_nearest_item_fn)
230            .inspect_ok(move |_| {
231                THREAD_HISTOGRAM_VEC.with_borrow_mut(|map| {
232                    map.entry(metric_key)
233                        .or_insert_with(|| {
234                            let (table_id, top_n, ef) = metric_key;
235                            let labels = [table_id.to_string(), top_n.to_string(), ef.to_string()];
236                            self.storage_metrics
237                                .vector_nearest_duration
238                                .with_guarded_label_values(&labels.each_ref().map(|s| s.as_str()))
239                        })
240                        .observe(start_time.elapsed().as_secs_f64());
241                });
242            })
243    }
244}
245
246impl<S: LocalStateStore> LocalStateStore for MonitoredTableStateStore<S> {
247    type FlushedSnapshotReader = MonitoredTableStateStore<S::FlushedSnapshotReader>;
248
249    type Iter<'a> = impl StateStoreIter + 'a;
250    type RevIter<'a> = impl StateStoreIter + 'a;
251
252    fn iter(
253        &self,
254        key_range: TableKeyRange,
255        read_options: ReadOptions,
256    ) -> impl Future<Output = StorageResult<Self::Iter<'_>>> + Send + '_ {
257        self.monitored_iter::<'_, _, _, StateStoreIterStats>(
258            self.table_id(),
259            self.inner.iter(key_range, read_options),
260        )
261    }
262
263    fn rev_iter(
264        &self,
265        key_range: TableKeyRange,
266        read_options: ReadOptions,
267    ) -> impl Future<Output = StorageResult<Self::RevIter<'_>>> + Send + '_ {
268        self.monitored_iter::<'_, _, _, StateStoreIterStats>(
269            self.table_id(),
270            self.inner.rev_iter(key_range, read_options),
271        )
272    }
273
274    fn insert(
275        &mut self,
276        key: TableKey<Bytes>,
277        new_val: Bytes,
278        old_val: Option<Bytes>,
279    ) -> StorageResult<()> {
280        // TODO: collect metrics
281        self.inner.insert(key, new_val, old_val)
282    }
283
284    fn delete(&mut self, key: TableKey<Bytes>, old_val: Bytes) -> StorageResult<()> {
285        // TODO: collect metrics
286        self.inner.delete(key, old_val)
287    }
288
289    fn get_table_watermark(&self, vnode: VirtualNode) -> Option<Bytes> {
290        self.inner.get_table_watermark(vnode)
291    }
292
293    fn new_flushed_snapshot_reader(&self) -> Self::FlushedSnapshotReader {
294        MonitoredTableStateStore::new(
295            self.inner.new_flushed_snapshot_reader(),
296            self.storage_metrics.clone(),
297            self.table_id(),
298        )
299    }
300
301    async fn update_vnode_bitmap(&mut self, vnodes: Arc<Bitmap>) -> StorageResult<Arc<Bitmap>> {
302        self.inner.update_vnode_bitmap(vnodes).await
303    }
304}
305
306impl<S: StateStoreWriteEpochControl> StateStoreWriteEpochControl for MonitoredTableStateStore<S> {
307    fn flush(&mut self) -> impl Future<Output = StorageResult<usize>> + Send + '_ {
308        self.inner.flush().instrument_await("store_flush".verbose())
309    }
310
311    async fn init(&mut self, options: InitOptions) -> StorageResult<()> {
312        self.inner.init(options).await
313    }
314
315    fn seal_current_epoch(&mut self, next_epoch: u64, opts: SealCurrentEpochOptions) {
316        // TODO: may collect metrics
317        self.inner.seal_current_epoch(next_epoch, opts)
318    }
319
320    fn try_flush(&mut self) -> impl Future<Output = StorageResult<()>> + Send + '_ {
321        self.inner
322            .try_flush()
323            .instrument_await("store_try_flush".verbose())
324    }
325}
326
327impl<S: StateStoreWriteVector> StateStoreWriteVector for MonitoredTableStateStore<S> {
328    fn insert(&mut self, vec: VectorRef<'_>, info: Bytes) -> StorageResult<()> {
329        // TODO: monitor
330        self.inner.insert(vec, info)
331    }
332}
333
334impl<S: StateStore> StateStore for MonitoredStateStore<S> {
335    type Local = MonitoredTableStateStore<S::Local>;
336    type ReadSnapshot = MonitoredTableStateStore<S::ReadSnapshot>;
337    type VectorWriter = MonitoredTableStateStore<S::VectorWriter>;
338
339    fn try_wait_epoch(
340        &self,
341        epoch: HummockReadEpoch,
342        options: TryWaitEpochOptions,
343    ) -> impl Future<Output = StorageResult<()>> + Send + '_ {
344        self.inner
345            .try_wait_epoch(epoch, options)
346            .instrument_await("store_wait_epoch".verbose())
347            .inspect_err(|e| error!(error = %e.as_report(), "Failed in wait_epoch"))
348    }
349
350    fn monitored(
351        self,
352        _storage_metrics: Arc<MonitoredStorageMetrics>,
353    ) -> MonitoredStateStore<Self> {
354        panic!("the state store is already monitored")
355    }
356
357    async fn new_local(&self, option: NewLocalOptions) -> Self::Local {
358        let table_id = option.table_id;
359        MonitoredTableStateStore::new(
360            self.inner
361                .new_local(option)
362                .instrument_await("store_new_local")
363                .await,
364            self.storage_metrics.clone(),
365            table_id,
366        )
367    }
368
369    async fn new_read_snapshot(
370        &self,
371        epoch: HummockReadEpoch,
372        options: NewReadSnapshotOptions,
373    ) -> StorageResult<Self::ReadSnapshot> {
374        Ok(MonitoredTableStateStore::new(
375            self.inner.new_read_snapshot(epoch, options).await?,
376            self.storage_metrics.clone(),
377            options.table_id,
378        ))
379    }
380
381    async fn new_vector_writer(&self, options: NewVectorWriterOptions) -> Self::VectorWriter {
382        let table_id = options.table_id;
383        MonitoredTableStateStore::new(
384            self.inner.new_vector_writer(options).await,
385            self.storage_metrics.clone(),
386            table_id,
387        )
388    }
389}
390
391impl MonitoredStateStore<HummockStorage> {
392    pub fn sstable_store(&self) -> SstableStoreRef {
393        self.inner.sstable_store()
394    }
395
396    pub fn object_id_manager(&self) -> ObjectIdManagerRef {
397        self.inner.object_id_manager().clone()
398    }
399}
400
401impl<S: AsHummock> AsHummock for MonitoredStateStore<S> {
402    fn as_hummock(&self) -> Option<&HummockStorage> {
403        self.inner.as_hummock()
404    }
405
406    fn sync(
407        &self,
408        sync_table_epochs: Vec<(HummockEpoch, HashSet<TableId>)>,
409    ) -> BoxFuture<'_, StorageResult<SyncResult>> {
410        async move {
411            let future = self
412                .inner
413                .sync(sync_table_epochs)
414                .instrument_await("store_sync");
415            let timer = self.storage_metrics.sync_duration.start_timer();
416            let sync_size = self.storage_metrics.sync_size.clone();
417            let sync_result = future
418                .await
419                .inspect_err(|e| error!(error = %e.as_report(), "Failed in sync"))?;
420            timer.observe_duration();
421            if sync_result.sync_size != 0 {
422                sync_size.observe(sync_result.sync_size as _);
423            }
424            Ok(sync_result)
425        }
426        .boxed()
427    }
428}
429
430/// A state store iterator wrapper for monitoring metrics.
431pub(crate) struct MonitoredStateStoreIter<
432    Item: IterItem,
433    I,
434    S: StateStoreIterStatsTrait<Item = Item>,
435> {
436    inner: I,
437    stats: MonitoredStateStoreIterStats<S>,
438    _phantom: PhantomData<Item>,
439}
440
441impl<Item: IterItem, I: StateStoreIter<Item>, S: StateStoreIterStatsTrait<Item = Item>>
442    StateStoreIter<Item> for MonitoredStateStoreIter<Item, I, S>
443{
444    async fn try_next(&mut self) -> StorageResult<Option<Item::ItemRef<'_>>> {
445        if let Some(item) = self
446            .inner
447            .try_next()
448            .instrument(tracing::trace_span!("store_iter_try_next"))
449            .await
450            .inspect_err(|e| error!(error = %e.as_report(), "Failed in next"))?
451        {
452            self.stats.inner.observe(item);
453            Ok(Some(item))
454        } else {
455            Ok(None)
456        }
457    }
458}
risingwave_storage/monitor/monitored_store.rs

risingwave_storage/monitor/
monitored_store.rs
