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

risingwave_storage/hummock/
utils.rs
