risingwave_storage/hummock/
utils.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::backtrace::Backtrace;
16use std::cmp::Ordering;
17use std::collections::VecDeque;
18use std::fmt::{Debug, Formatter};
19use std::ops::Bound::{Excluded, Included, Unbounded};
20use std::ops::{Bound, RangeBounds};
21use std::sync::Arc;
22use std::sync::atomic::{AtomicBool, AtomicU64, Ordering as AtomicOrdering};
23use std::time::{Duration, Instant};
24
25use bytes::Bytes;
26use foyer::Hint;
27use futures::{Stream, StreamExt, pin_mut};
28use parking_lot::Mutex;
29use risingwave_common::catalog::TableId;
30use risingwave_common::config::StorageMemoryConfig;
31use risingwave_expr::codegen::try_stream;
32use risingwave_hummock_sdk::can_concat;
33use risingwave_hummock_sdk::compaction_group::StateTableId;
34use risingwave_hummock_sdk::key::{
35    EmptySliceRef, FullKey, TableKey, UserKey, bound_table_key_range,
36};
37use risingwave_hummock_sdk::sstable_info::SstableInfo;
38use tokio::sync::oneshot::{Receiver, Sender, channel};
39
40use super::{HummockError, HummockResult, SstableStoreRef};
41use crate::error::{StorageError, StorageResult};
42use crate::hummock::CachePolicy;
43use crate::hummock::local_version::pinned_version::PinnedVersion;
44use crate::mem_table::{KeyOp, MemTableError};
45use crate::monitor::MemoryCollector;
46use crate::store::{
47    OpConsistencyLevel, ReadOptions, StateStoreGet, StateStoreKeyedRow, StateStoreRead,
48};
49
50pub fn range_overlap<R, B>(
51    search_key_range: &R,
52    inclusive_start_key: &B,
53    end_key: Bound<&B>,
54) -> bool
55where
56    R: RangeBounds<B>,
57    B: Ord,
58{
59    let (start_bound, end_bound) = (search_key_range.start_bound(), search_key_range.end_bound());
60
61    //        RANGE
62    // TABLE
63    let too_left = match (start_bound, end_key) {
64        (Included(range_start), Included(inclusive_end_key)) => range_start > inclusive_end_key,
65        (Included(range_start), Excluded(end_key))
66        | (Excluded(range_start), Included(end_key))
67        | (Excluded(range_start), Excluded(end_key)) => range_start >= end_key,
68        (Unbounded, _) | (_, Unbounded) => false,
69    };
70    // RANGE
71    //        TABLE
72    let too_right = match end_bound {
73        Included(range_end) => range_end < inclusive_start_key,
74        Excluded(range_end) => range_end <= inclusive_start_key,
75        Unbounded => false,
76    };
77
78    !too_left && !too_right
79}
80
81pub fn filter_single_sst<R, B>(info: &SstableInfo, table_id: TableId, table_key_range: &R) -> bool
82where
83    R: RangeBounds<TableKey<B>>,
84    B: AsRef<[u8]> + EmptySliceRef,
85{
86    debug_assert!(info.table_ids.is_sorted());
87    let table_range = &info.key_range;
88    let table_start = FullKey::decode(table_range.left.as_ref()).user_key;
89    let table_end = FullKey::decode(table_range.right.as_ref()).user_key;
90    let (left, right) = bound_table_key_range(table_id, table_key_range);
91    let left: Bound<UserKey<&[u8]>> = left.as_ref().map(|key| key.as_ref());
92    let right: Bound<UserKey<&[u8]>> = right.as_ref().map(|key| key.as_ref());
93
94    info.table_ids.binary_search(&table_id).is_ok()
95        && range_overlap(
96            &(left, right),
97            &table_start,
98            if table_range.right_exclusive {
99                Bound::Excluded(&table_end)
100            } else {
101                Bound::Included(&table_end)
102            },
103        )
104}
105
106/// Search the SST containing the specified key within a level, using binary search.
107pub(crate) fn search_sst_idx(ssts: &[SstableInfo], key: UserKey<&[u8]>) -> usize {
108    ssts.partition_point(|table| {
109        let ord = FullKey::decode(&table.key_range.left).user_key.cmp(&key);
110        ord == Ordering::Less || ord == Ordering::Equal
111    })
112}
113
114/// Prune overlapping SSTs that does not overlap with a specific key range or does not overlap with
115/// a specific table id. Returns the sst ids after pruning.
116pub fn prune_overlapping_ssts<'a, R, B>(
117    ssts: &'a [SstableInfo],
118    table_id: TableId,
119    table_key_range: &'a R,
120) -> impl DoubleEndedIterator<Item = &'a SstableInfo>
121where
122    R: RangeBounds<TableKey<B>>,
123    B: AsRef<[u8]> + EmptySliceRef,
124{
125    ssts.iter()
126        .filter(move |info| filter_single_sst(info, table_id, table_key_range))
127}
128
129/// Prune non-overlapping SSTs that does not overlap with a specific key range or does not overlap
130/// with a specific table id. Returns the sst ids after pruning.
131#[allow(clippy::type_complexity)]
132pub fn prune_nonoverlapping_ssts<'a>(
133    ssts: &'a [SstableInfo],
134    user_key_range: (Bound<UserKey<&'a [u8]>>, Bound<UserKey<&'a [u8]>>),
135    table_id: StateTableId,
136) -> impl DoubleEndedIterator<Item = &'a SstableInfo> {
137    debug_assert!(can_concat(ssts));
138    let start_table_idx = match user_key_range.0 {
139        Included(key) | Excluded(key) => search_sst_idx(ssts, key).saturating_sub(1),
140        _ => 0,
141    };
142    let end_table_idx = match user_key_range.1 {
143        Included(key) | Excluded(key) => search_sst_idx(ssts, key).saturating_sub(1),
144        _ => ssts.len().saturating_sub(1),
145    };
146    ssts[start_table_idx..=end_table_idx]
147        .iter()
148        .filter(move |sst| sst.table_ids.binary_search(&table_id).is_ok())
149}
150
151type RequestQueue = VecDeque<(Sender<MemoryTracker>, u64)>;
152enum MemoryRequest {
153    Ready(MemoryTracker),
154    Pending(Receiver<MemoryTracker>),
155}
156
157struct MemoryLimiterInner {
158    total_size: AtomicU64,
159    controller: Mutex<RequestQueue>,
160    has_waiter: AtomicBool,
161    quota: u64,
162}
163
164impl MemoryLimiterInner {
165    fn release_quota(&self, quota: u64) {
166        self.total_size.fetch_sub(quota, AtomicOrdering::SeqCst);
167    }
168
169    fn add_memory(&self, quota: u64) {
170        self.total_size.fetch_add(quota, AtomicOrdering::SeqCst);
171    }
172
173    fn may_notify_waiters(self: &Arc<Self>) {
174        // check `has_waiter` to avoid access lock every times drop `MemoryTracker`.
175        if !self.has_waiter.load(AtomicOrdering::Acquire) {
176            return;
177        }
178        let mut notify_waiters = vec![];
179        {
180            let mut waiters = self.controller.lock();
181            while let Some((_, quota)) = waiters.front() {
182                if !self.try_require_memory(*quota) {
183                    break;
184                }
185                let (tx, quota) = waiters.pop_front().unwrap();
186                notify_waiters.push((tx, quota));
187            }
188
189            if waiters.is_empty() {
190                self.has_waiter.store(false, AtomicOrdering::Release);
191            }
192        }
193
194        for (tx, quota) in notify_waiters {
195            let _ = tx.send(MemoryTracker::new(self.clone(), quota));
196        }
197    }
198
199    fn try_require_memory(&self, quota: u64) -> bool {
200        let mut current_quota = self.total_size.load(AtomicOrdering::Acquire);
201        while self.permit_quota(current_quota, quota) {
202            match self.total_size.compare_exchange(
203                current_quota,
204                current_quota + quota,
205                AtomicOrdering::SeqCst,
206                AtomicOrdering::SeqCst,
207            ) {
208                Ok(_) => {
209                    return true;
210                }
211                Err(old_quota) => {
212                    current_quota = old_quota;
213                }
214            }
215        }
216        false
217    }
218
219    fn require_memory(self: &Arc<Self>, quota: u64) -> MemoryRequest {
220        let mut waiters = self.controller.lock();
221        let first_req = waiters.is_empty();
222        if first_req {
223            // When this request is the first waiter but the previous `MemoryTracker` is just release a large quota, it may skip notifying this waiter because it has checked `has_waiter` and found it was false. So we must set it one and retry `try_require_memory` again to avoid deadlock.
224            self.has_waiter.store(true, AtomicOrdering::Release);
225        }
226        // We must require again with lock because some other MemoryTracker may drop just after this thread gets mutex but before it enters queue.
227        if self.try_require_memory(quota) {
228            if first_req {
229                self.has_waiter.store(false, AtomicOrdering::Release);
230            }
231            return MemoryRequest::Ready(MemoryTracker::new(self.clone(), quota));
232        }
233        let (tx, rx) = channel();
234        waiters.push_back((tx, quota));
235        MemoryRequest::Pending(rx)
236    }
237
238    fn permit_quota(&self, current_quota: u64, _request_quota: u64) -> bool {
239        current_quota <= self.quota
240    }
241}
242
243pub struct MemoryLimiter {
244    inner: Arc<MemoryLimiterInner>,
245}
246
247impl Debug for MemoryLimiter {
248    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
249        f.debug_struct("MemoryLimiter")
250            .field("quota", &self.inner.quota)
251            .field("usage", &self.inner.total_size)
252            .finish()
253    }
254}
255
256pub struct MemoryTracker {
257    limiter: Arc<MemoryLimiterInner>,
258    quota: Option<u64>,
259}
260impl MemoryTracker {
261    fn new(limiter: Arc<MemoryLimiterInner>, quota: u64) -> Self {
262        Self {
263            limiter,
264            quota: Some(quota),
265        }
266    }
267}
268
269impl Debug for MemoryTracker {
270    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
271        f.debug_struct("MemoryTracker")
272            .field("quota", &self.quota)
273            .finish()
274    }
275}
276
277impl MemoryLimiter {
278    pub fn unlimit() -> Arc<Self> {
279        Arc::new(Self::new(u64::MAX))
280    }
281
282    pub fn new(quota: u64) -> Self {
283        Self {
284            inner: Arc::new(MemoryLimiterInner {
285                total_size: AtomicU64::new(0),
286                controller: Mutex::new(VecDeque::default()),
287                has_waiter: AtomicBool::new(false),
288                quota,
289            }),
290        }
291    }
292
293    pub fn try_require_memory(&self, quota: u64) -> Option<MemoryTracker> {
294        if self.inner.try_require_memory(quota) {
295            Some(MemoryTracker::new(self.inner.clone(), quota))
296        } else {
297            None
298        }
299    }
300
301    pub fn get_memory_usage(&self) -> u64 {
302        self.inner.total_size.load(AtomicOrdering::Acquire)
303    }
304
305    pub fn quota(&self) -> u64 {
306        self.inner.quota
307    }
308
309    pub fn must_require_memory(&self, quota: u64) -> MemoryTracker {
310        if !self.inner.try_require_memory(quota) {
311            self.inner.add_memory(quota);
312        }
313
314        MemoryTracker::new(self.inner.clone(), quota)
315    }
316}
317
318impl MemoryLimiter {
319    pub async fn require_memory(&self, quota: u64) -> MemoryTracker {
320        match self.inner.require_memory(quota) {
321            MemoryRequest::Ready(tracker) => tracker,
322            MemoryRequest::Pending(rx) => rx.await.unwrap(),
323        }
324    }
325}
326
327impl MemoryTracker {
328    pub fn try_increase_memory(&mut self, target: u64) -> bool {
329        let quota = self.quota.unwrap();
330        if quota >= target {
331            return true;
332        }
333        if self.limiter.try_require_memory(target - quota) {
334            self.quota = Some(target);
335            true
336        } else {
337            false
338        }
339    }
340}
341
342// We must notify waiters outside `MemoryTracker` to avoid dead-lock and loop-owner.
343impl Drop for MemoryTracker {
344    fn drop(&mut self) {
345        if let Some(quota) = self.quota.take() {
346            self.limiter.release_quota(quota);
347            self.limiter.may_notify_waiters();
348        }
349    }
350}
351
352/// Check whether the items in `sub_iter` is a subset of the items in `full_iter`, and meanwhile
353/// preserve the order.
354pub fn check_subset_preserve_order<T: Eq>(
355    sub_iter: impl Iterator<Item = T>,
356    mut full_iter: impl Iterator<Item = T>,
357) -> bool {
358    for sub_iter_item in sub_iter {
359        let mut found = false;
360        for full_iter_item in full_iter.by_ref() {
361            if sub_iter_item == full_iter_item {
362                found = true;
363                break;
364            }
365        }
366        if !found {
367            return false;
368        }
369    }
370    true
371}
372
373static SANITY_CHECK_ENABLED: AtomicBool = AtomicBool::new(cfg!(debug_assertions));
374
375/// This function is intended to be called during compute node initialization if the storage
376/// sanity check is not desired. This controls a global flag so only need to be called once
377/// if need to disable the sanity check.
378pub fn disable_sanity_check() {
379    SANITY_CHECK_ENABLED.store(false, AtomicOrdering::Release);
380}
381
382pub(crate) fn sanity_check_enabled() -> bool {
383    SANITY_CHECK_ENABLED.load(AtomicOrdering::Acquire)
384}
385
386async fn get_from_state_store(
387    state_store: &impl StateStoreGet,
388    key: TableKey<Bytes>,
389    read_options: ReadOptions,
390) -> StorageResult<Option<Bytes>> {
391    state_store
392        .on_key_value(key, read_options, |_, value| {
393            Ok(Bytes::copy_from_slice(value))
394        })
395        .await
396}
397
398/// Make sure the key to insert should not exist in storage.
399pub(crate) async fn do_insert_sanity_check(
400    table_id: TableId,
401    key: &TableKey<Bytes>,
402    value: &Bytes,
403    inner: &impl StateStoreRead,
404    op_consistency_level: &OpConsistencyLevel,
405) -> StorageResult<()> {
406    if let OpConsistencyLevel::Inconsistent = op_consistency_level {
407        return Ok(());
408    }
409    let read_options = ReadOptions {
410        cache_policy: CachePolicy::Fill(Hint::Normal),
411        ..Default::default()
412    };
413    let stored_value = get_from_state_store(inner, key.clone(), read_options).await?;
414
415    if let Some(stored_value) = stored_value {
416        return Err(Box::new(MemTableError::InconsistentOperation {
417            table_id,
418            key: key.clone(),
419            prev: KeyOp::Insert(stored_value),
420            new: KeyOp::Insert(value.clone()),
421        })
422        .into());
423    }
424    Ok(())
425}
426
427/// Make sure that the key to delete should exist in storage and the value should be matched.
428pub(crate) async fn do_delete_sanity_check(
429    table_id: TableId,
430    key: &TableKey<Bytes>,
431    old_value: &Bytes,
432    inner: &impl StateStoreRead,
433    op_consistency_level: &OpConsistencyLevel,
434) -> StorageResult<()> {
435    let OpConsistencyLevel::ConsistentOldValue {
436        check_old_value: old_value_checker,
437        ..
438    } = op_consistency_level
439    else {
440        return Ok(());
441    };
442    let read_options = ReadOptions {
443        cache_policy: CachePolicy::Fill(Hint::Normal),
444        ..Default::default()
445    };
446    match get_from_state_store(inner, key.clone(), read_options).await? {
447        None => Err(Box::new(MemTableError::InconsistentOperation {
448            table_id,
449            key: key.clone(),
450            prev: KeyOp::Delete(Bytes::default()),
451            new: KeyOp::Delete(old_value.clone()),
452        })
453        .into()),
454        Some(stored_value) => {
455            if !old_value_checker(&stored_value, old_value) {
456                Err(Box::new(MemTableError::InconsistentOperation {
457                    table_id,
458                    key: key.clone(),
459                    prev: KeyOp::Insert(stored_value),
460                    new: KeyOp::Delete(old_value.clone()),
461                })
462                .into())
463            } else {
464                Ok(())
465            }
466        }
467    }
468}
469
470/// Make sure that the key to update should exist in storage and the value should be matched
471pub(crate) async fn do_update_sanity_check(
472    table_id: TableId,
473    key: &TableKey<Bytes>,
474    old_value: &Bytes,
475    new_value: &Bytes,
476    inner: &impl StateStoreRead,
477    op_consistency_level: &OpConsistencyLevel,
478) -> StorageResult<()> {
479    let OpConsistencyLevel::ConsistentOldValue {
480        check_old_value: old_value_checker,
481        ..
482    } = op_consistency_level
483    else {
484        return Ok(());
485    };
486    let read_options = ReadOptions {
487        cache_policy: CachePolicy::Fill(Hint::Normal),
488        ..Default::default()
489    };
490
491    match get_from_state_store(inner, key.clone(), read_options).await? {
492        None => Err(Box::new(MemTableError::InconsistentOperation {
493            table_id,
494            key: key.clone(),
495            prev: KeyOp::Delete(Bytes::default()),
496            new: KeyOp::Update((old_value.clone(), new_value.clone())),
497        })
498        .into()),
499        Some(stored_value) => {
500            if !old_value_checker(&stored_value, old_value) {
501                Err(Box::new(MemTableError::InconsistentOperation {
502                    table_id,
503                    key: key.clone(),
504                    prev: KeyOp::Insert(stored_value),
505                    new: KeyOp::Update((old_value.clone(), new_value.clone())),
506                })
507                .into())
508            } else {
509                Ok(())
510            }
511        }
512    }
513}
514
515pub fn cmp_delete_range_left_bounds(a: Bound<&Bytes>, b: Bound<&Bytes>) -> Ordering {
516    match (a, b) {
517        // only right bound of delete range can be `Unbounded`
518        (Unbounded, _) | (_, Unbounded) => unreachable!(),
519        (Included(x), Included(y)) | (Excluded(x), Excluded(y)) => x.cmp(y),
520        (Included(x), Excluded(y)) => x.cmp(y).then(Ordering::Less),
521        (Excluded(x), Included(y)) => x.cmp(y).then(Ordering::Greater),
522    }
523}
524
525pub(crate) fn validate_delete_range(left: &Bound<Bytes>, right: &Bound<Bytes>) -> bool {
526    match (left, right) {
527        // only right bound of delete range can be `Unbounded`
528        (Unbounded, _) => unreachable!(),
529        (_, Unbounded) => true,
530        (Included(x), Included(y)) => x <= y,
531        (Included(x), Excluded(y)) | (Excluded(x), Included(y)) | (Excluded(x), Excluded(y)) => {
532            x < y
533        }
534    }
535}
536
537#[expect(dead_code)]
538pub(crate) fn filter_with_delete_range<'a>(
539    kv_iter: impl Iterator<Item = (TableKey<Bytes>, KeyOp)> + 'a,
540    mut delete_ranges_iter: impl Iterator<Item = &'a (Bound<Bytes>, Bound<Bytes>)> + 'a,
541) -> impl Iterator<Item = (TableKey<Bytes>, KeyOp)> + 'a {
542    let mut range = delete_ranges_iter.next();
543    if let Some((range_start, range_end)) = range {
544        assert!(
545            validate_delete_range(range_start, range_end),
546            "range_end {:?} smaller than range_start {:?}",
547            range_start,
548            range_end
549        );
550    }
551    kv_iter.filter(move |(key, _)| {
552        if let Some(range_bound) = range {
553            if cmp_delete_range_left_bounds(Included(&key.0), range_bound.0.as_ref())
554                == Ordering::Less
555            {
556                true
557            } else if range_bound.contains(key.as_ref()) {
558                false
559            } else {
560                // Key has exceeded the current key range. Advance to the next range.
561                loop {
562                    range = delete_ranges_iter.next();
563                    if let Some(range_bound) = range {
564                        assert!(
565                            validate_delete_range(&range_bound.0, &range_bound.1),
566                            "range_end {:?} smaller than range_start {:?}",
567                            range_bound.0,
568                            range_bound.1
569                        );
570                        if cmp_delete_range_left_bounds(Included(key), range_bound.0.as_ref())
571                            == Ordering::Less
572                        {
573                            // Not fall in the next delete range
574                            break true;
575                        } else if range_bound.contains(key.as_ref()) {
576                            // Fall in the next delete range
577                            break false;
578                        } else {
579                            // Exceed the next delete range. Go to the next delete range if there is
580                            // any in the next loop
581                            continue;
582                        }
583                    } else {
584                        // No more delete range.
585                        break true;
586                    }
587                }
588            }
589        } else {
590            true
591        }
592    })
593}
594
595/// Wait for the `committed_epoch` of `table_id` to reach `wait_epoch`.
596///
597/// When the `table_id` does not exist in the latest version, we assume that
598/// the table is not created yet, and will wait until the table is created.
599pub(crate) async fn wait_for_epoch(
600    notifier: &tokio::sync::watch::Sender<PinnedVersion>,
601    wait_epoch: u64,
602    table_id: TableId,
603) -> StorageResult<PinnedVersion> {
604    let mut prev_committed_epoch = None;
605    let prev_committed_epoch = &mut prev_committed_epoch;
606    let version = wait_for_update(
607        notifier,
608        |version| {
609            let committed_epoch = version.table_committed_epoch(table_id);
610            let ret = if let Some(committed_epoch) = committed_epoch {
611                if committed_epoch >= wait_epoch {
612                    Ok(true)
613                } else {
614                    Ok(false)
615                }
616            } else if prev_committed_epoch.is_none() {
617                Ok(false)
618            } else {
619                Err(HummockError::wait_epoch(format!(
620                    "table {} has been dropped",
621                    table_id
622                )))
623            };
624            *prev_committed_epoch = committed_epoch;
625            ret
626        },
627        || {
628            format!(
629                "wait_for_epoch: epoch: {}, table_id: {}",
630                wait_epoch, table_id
631            )
632        },
633    )
634    .await?;
635    Ok(version)
636}
637
638pub(crate) async fn wait_for_update(
639    notifier: &tokio::sync::watch::Sender<PinnedVersion>,
640    mut inspect_fn: impl FnMut(&PinnedVersion) -> HummockResult<bool>,
641    mut periodic_debug_info: impl FnMut() -> String,
642) -> HummockResult<PinnedVersion> {
643    let mut receiver = notifier.subscribe();
644    {
645        let version = receiver.borrow_and_update();
646        if inspect_fn(&version)? {
647            return Ok(version.clone());
648        }
649    }
650    let start_time = Instant::now();
651    loop {
652        match tokio::time::timeout(Duration::from_secs(30), receiver.changed()).await {
653            Err(_) => {
654                // Provide backtrace iff in debug mode for observability.
655                let backtrace = cfg!(debug_assertions)
656                    .then(Backtrace::capture)
657                    .map(tracing::field::display);
658
659                // The reason that we need to retry here is batch scan in
660                // chain/rearrange_chain is waiting for an
661                // uncommitted epoch carried by the CreateMV barrier, which
662                // can take unbounded time to become committed and propagate
663                // to the CN. We should consider removing the retry as well as wait_epoch
664                // for chain/rearrange_chain if we enforce
665                // chain/rearrange_chain to be scheduled on the same
666                // CN with the same distribution as the upstream MV.
667                // See #3845 for more details.
668                tracing::warn!(
669                    info = periodic_debug_info(),
670                    elapsed = ?start_time.elapsed(),
671                    backtrace,
672                    "timeout when waiting for version update",
673                );
674                continue;
675            }
676            Ok(Err(_)) => {
677                return Err(HummockError::wait_epoch("tx dropped"));
678            }
679            Ok(Ok(_)) => {
680                let version = receiver.borrow_and_update();
681                if inspect_fn(&version)? {
682                    return Ok(version.clone());
683                }
684            }
685        }
686    }
687}
688
689pub struct HummockMemoryCollector {
690    sstable_store: SstableStoreRef,
691    limiter: Arc<MemoryLimiter>,
692    storage_memory_config: StorageMemoryConfig,
693}
694
695impl HummockMemoryCollector {
696    pub fn new(
697        sstable_store: SstableStoreRef,
698        limiter: Arc<MemoryLimiter>,
699        storage_memory_config: StorageMemoryConfig,
700    ) -> Self {
701        Self {
702            sstable_store,
703            limiter,
704            storage_memory_config,
705        }
706    }
707}
708
709impl MemoryCollector for HummockMemoryCollector {
710    fn get_meta_memory_usage(&self) -> u64 {
711        self.sstable_store.meta_cache().memory().usage() as _
712    }
713
714    fn get_data_memory_usage(&self) -> u64 {
715        self.sstable_store.block_cache().memory().usage() as _
716    }
717
718    fn get_vector_meta_memory_usage(&self) -> u64 {
719        self.sstable_store.vector_meta_cache.usage() as _
720    }
721
722    fn get_vector_data_memory_usage(&self) -> u64 {
723        self.sstable_store.vector_block_cache.usage() as _
724    }
725
726    fn get_uploading_memory_usage(&self) -> u64 {
727        self.limiter.get_memory_usage()
728    }
729
730    fn get_prefetch_memory_usage(&self) -> usize {
731        self.sstable_store.get_prefetch_memory_usage()
732    }
733
734    fn get_meta_cache_memory_usage_ratio(&self) -> f64 {
735        self.sstable_store.meta_cache().memory().usage() as f64
736            / self.sstable_store.meta_cache().memory().capacity() as f64
737    }
738
739    fn get_block_cache_memory_usage_ratio(&self) -> f64 {
740        self.sstable_store.block_cache().memory().usage() as f64
741            / self.sstable_store.block_cache().memory().capacity() as f64
742    }
743
744    fn get_vector_meta_cache_memory_usage_ratio(&self) -> f64 {
745        self.sstable_store.vector_meta_cache.usage() as f64
746            / self.sstable_store.vector_meta_cache.capacity() as f64
747    }
748
749    fn get_vector_data_cache_memory_usage_ratio(&self) -> f64 {
750        self.sstable_store.vector_block_cache.usage() as f64
751            / self.sstable_store.vector_block_cache.capacity() as f64
752    }
753
754    fn get_shared_buffer_usage_ratio(&self) -> f64 {
755        self.limiter.get_memory_usage() as f64
756            / (self.storage_memory_config.shared_buffer_capacity_mb * 1024 * 1024) as f64
757    }
758}
759
760#[try_stream(ok = StateStoreKeyedRow, error = StorageError)]
761pub(crate) async fn merge_stream<'a>(
762    mem_table_iter: impl Iterator<Item = (&'a TableKey<Bytes>, &'a KeyOp)> + 'a,
763    inner_stream: impl Stream<Item = StorageResult<StateStoreKeyedRow>> + 'static,
764    table_id: TableId,
765    epoch: u64,
766    rev: bool,
767) {
768    let inner_stream = inner_stream.peekable();
769    pin_mut!(inner_stream);
770
771    let mut mem_table_iter = mem_table_iter.fuse().peekable();
772
773    loop {
774        match (inner_stream.as_mut().peek().await, mem_table_iter.peek()) {
775            (None, None) => break,
776            // The mem table side has come to an end, return data from the shared storage.
777            (Some(_), None) => {
778                let (key, value) = inner_stream.next().await.unwrap()?;
779                yield (key, value)
780            }
781            // The stream side has come to an end, return data from the mem table.
782            (None, Some(_)) => {
783                let (key, key_op) = mem_table_iter.next().unwrap();
784                match key_op {
785                    KeyOp::Insert(value) | KeyOp::Update((_, value)) => {
786                        yield (FullKey::new(table_id, key.clone(), epoch), value.clone())
787                    }
788                    _ => {}
789                }
790            }
791            (Some(Ok((inner_key, _))), Some((mem_table_key, _))) => {
792                debug_assert_eq!(inner_key.user_key.table_id, table_id);
793                let mut ret = inner_key.user_key.table_key.cmp(mem_table_key);
794                if rev {
795                    ret = ret.reverse();
796                }
797                match ret {
798                    Ordering::Less => {
799                        // yield data from storage
800                        let (key, value) = inner_stream.next().await.unwrap()?;
801                        yield (key, value);
802                    }
803                    Ordering::Equal => {
804                        // both memtable and storage contain the key, so we advance both
805                        // iterators and return the data in memory.
806
807                        let (_, key_op) = mem_table_iter.next().unwrap();
808                        let (key, old_value_in_inner) = inner_stream.next().await.unwrap()?;
809                        match key_op {
810                            KeyOp::Insert(value) => {
811                                yield (key.clone(), value.clone());
812                            }
813                            KeyOp::Delete(_) => {}
814                            KeyOp::Update((old_value, new_value)) => {
815                                debug_assert!(old_value == &old_value_in_inner);
816
817                                yield (key, new_value.clone());
818                            }
819                        }
820                    }
821                    Ordering::Greater => {
822                        // yield data from mem table
823                        let (key, key_op) = mem_table_iter.next().unwrap();
824
825                        match key_op {
826                            KeyOp::Insert(value) => {
827                                yield (FullKey::new(table_id, key.clone(), epoch), value.clone());
828                            }
829                            KeyOp::Delete(_) => {}
830                            KeyOp::Update(_) => unreachable!(
831                                "memtable update should always be paired with a storage key"
832                            ),
833                        }
834                    }
835                }
836            }
837            (Some(Err(_)), Some(_)) => {
838                // Throw the error.
839                return Err(inner_stream.next().await.unwrap().unwrap_err());
840            }
841        }
842    }
843}
844
845#[cfg(test)]
846mod tests {
847    use std::future::{Future, poll_fn};
848    use std::sync::Arc;
849    use std::task::Poll;
850
851    use futures::FutureExt;
852    use futures::future::join_all;
853    use rand::random_range;
854
855    use crate::hummock::utils::MemoryLimiter;
856
857    async fn assert_pending(future: &mut (impl Future + Unpin)) {
858        for _ in 0..10 {
859            assert!(
860                poll_fn(|cx| Poll::Ready(future.poll_unpin(cx)))
861                    .await
862                    .is_pending()
863            );
864        }
865    }
866
867    #[tokio::test]
868    async fn test_loose_memory_limiter() {
869        let quota = 5;
870        let memory_limiter = MemoryLimiter::new(quota);
871        drop(memory_limiter.require_memory(6).await);
872        let tracker1 = memory_limiter.require_memory(3).await;
873        assert_eq!(3, memory_limiter.get_memory_usage());
874        let tracker2 = memory_limiter.require_memory(4).await;
875        assert_eq!(7, memory_limiter.get_memory_usage());
876        let mut future = memory_limiter.require_memory(5).boxed();
877        assert_pending(&mut future).await;
878        assert_eq!(7, memory_limiter.get_memory_usage());
879        drop(tracker1);
880        let tracker3 = future.await;
881        assert_eq!(9, memory_limiter.get_memory_usage());
882        drop(tracker2);
883        assert_eq!(5, memory_limiter.get_memory_usage());
884        drop(tracker3);
885        assert_eq!(0, memory_limiter.get_memory_usage());
886    }
887
888    #[tokio::test(flavor = "multi_thread", worker_threads = 8)]
889    async fn test_multi_thread_acquire_memory() {
890        const QUOTA: u64 = 10;
891        let memory_limiter = Arc::new(MemoryLimiter::new(200));
892        let mut handles = vec![];
893        for _ in 0..40 {
894            let limiter = memory_limiter.clone();
895            let h = tokio::spawn(async move {
896                let mut buffers = vec![];
897                let mut current_buffer_usage = random_range(2..=9);
898                for _ in 0..1000 {
899                    if buffers.len() < current_buffer_usage
900                        && let Some(tracker) = limiter.try_require_memory(QUOTA)
901                    {
902                        buffers.push(tracker);
903                    } else {
904                        buffers.clear();
905                        current_buffer_usage = random_range(2..=9);
906                        let req = limiter.require_memory(QUOTA);
907                        match tokio::time::timeout(std::time::Duration::from_millis(1), req).await {
908                            Ok(tracker) => {
909                                buffers.push(tracker);
910                            }
911                            Err(_) => {
912                                continue;
913                            }
914                        }
915                    }
916                    let sleep_time = random_range(1..=3);
917                    tokio::time::sleep(std::time::Duration::from_millis(sleep_time)).await;
918                }
919            });
920            handles.push(h);
921        }
922        let h = join_all(handles);
923        let _ = h.await;
924    }
925}
risingwave_storage/hummock/utils.rs

risingwave_storage/hummock/
utils.rs
