risingwave_storage/hummock/
block_stream.rs

1// Copyright 2023 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::collections::VecDeque;
16use std::sync::Arc;
17use std::sync::atomic::{AtomicUsize, Ordering};
18
19use bytes::{Bytes, BytesMut};
20use fail::fail_point;
21use risingwave_object_store::object::{MonitoredStreamingReader, ObjectError};
22
23use super::BlockMeta;
24use crate::hummock::{BlockHolder, HummockResult};
25
26pub struct MemoryUsageTracker {
27    total_usage: Arc<AtomicUsize>,
28    usage: usize,
29}
30
31impl MemoryUsageTracker {
32    pub fn new(total_usage: Arc<AtomicUsize>, usage: usize) -> Self {
33        total_usage.fetch_add(usage, Ordering::SeqCst);
34        Self { total_usage, usage }
35    }
36}
37
38impl Drop for MemoryUsageTracker {
39    fn drop(&mut self) {
40        self.total_usage.fetch_sub(self.usage, Ordering::SeqCst);
41    }
42}
43
44/// An iterator that reads the blocks of an SST step by step from a given stream of bytes.
45pub struct BlockDataStream {
46    buf_reader: MonitoredStreamingReader,
47
48    /// The index of the next block. Note that `block_idx` is relative to the start index of the
49    /// stream (and is compatible with `block_sizes`); it is not relative to the corresponding
50    /// SST. That is, if streaming starts at block 2 of a given SST `T`, then `block_idx = 0`
51    /// refers to the third block of `T`.
52    block_idx: usize,
53
54    /// The sizes of each block which the stream reads. The first number states the compressed size
55    /// in the stream. The second number is the block's uncompressed size.  Note that the list does
56    /// not contain the size of blocks which precede the first streamed block. That is, if
57    /// streaming starts at block 2 of a given SST, then the list does not contain information
58    /// about block 0 and block 1.
59    block_sizes: Vec<(u32, u32)>,
60
61    buf: Bytes,
62
63    buff_offset: usize,
64}
65
66impl BlockDataStream {
67    /// Reads the blocks described by `block_metas` from a byte stream positioned at their start.
68    /// The block index is relative to this slice, not to the full SST.
69    /// Only retain the lengths needed to frame and decode blocks, without cloning their keys.
70    pub fn new(
71        // The stream that provides raw data.
72        byte_stream: MonitoredStreamingReader,
73        // Meta data of the SST that is streamed.
74        block_metas: &[BlockMeta],
75    ) -> Self {
76        Self {
77            buf_reader: byte_stream,
78            block_idx: 0,
79            block_sizes: block_metas
80                .iter()
81                .map(|meta| (meta.len, meta.uncompressed_size))
82                .collect(),
83            buf: Bytes::default(),
84            buff_offset: 0,
85        }
86    }
87
88    /// Reads the next block from the stream and returns it. Returns `None` if there are no blocks
89    /// left to read.
90    pub async fn next_block(&mut self) -> HummockResult<Option<(Bytes, usize)>> {
91        if self.block_idx >= self.block_sizes.len() {
92            return Ok(None);
93        }
94
95        let (compressed_size, uncompressed_size) = self.block_sizes[self.block_idx];
96        fail_point!("stream_read_err", |_| Err(ObjectError::internal(
97            "stream read error"
98        )
99        .into()));
100        let uncompressed_size = uncompressed_size as usize;
101        let end = self.buff_offset + compressed_size as usize;
102        let data = if end > self.buf.len() {
103            let current_block = self.read_next_buf(compressed_size as usize).await?;
104            self.buff_offset = 0;
105            current_block
106        } else {
107            let data = self.buf.slice(self.buff_offset..end);
108            self.buff_offset = end;
109            data
110        };
111
112        self.block_idx += 1;
113        Ok(Some((data, uncompressed_size)))
114    }
115
116    async fn read_next_buf(&mut self, read_size: usize) -> HummockResult<Bytes> {
117        let mut read_buf = BytesMut::with_capacity(read_size);
118        let start_pos = if self.buff_offset < self.buf.len() {
119            read_buf.extend_from_slice(&self.buf[self.buff_offset..]);
120            self.buf.len() - self.buff_offset
121        } else {
122            0
123        };
124        let mut rest = read_size - start_pos;
125        while rest > 0 {
126            let next_packet = self
127                .buf_reader
128                .read_bytes()
129                .await
130                .unwrap_or_else(|| Err(ObjectError::internal("read unexpected EOF")))?;
131            let read_len = std::cmp::min(next_packet.len(), rest);
132            read_buf.extend_from_slice(&next_packet[..read_len]);
133            rest -= read_len;
134            if rest == 0 {
135                self.buf = next_packet.slice(read_len..);
136                return Ok(read_buf.freeze());
137            }
138        }
139        self.buf = Bytes::default();
140        Ok(read_buf.freeze())
141    }
142}
143
144/// Consecutive decoded blocks whose I/O has already completed in `SstableStore::prefetch_blocks`.
145/// Consuming them is synchronous and infallible. The tracker is retained until the stream drops.
146pub struct PrefetchBlockStream {
147    blocks: VecDeque<BlockHolder>,
148    /// SST index of the first remaining block, or the end index when exhausted.
149    block_index: usize,
150    _tracker: Option<MemoryUsageTracker>,
151}
152
153pub(super) enum PrefetchLookup {
154    Hit(BlockHolder),
155    /// The target precedes the remaining range. The stream is unchanged.
156    BeforeStart,
157    /// All buffered blocks were consumed without reaching the target.
158    Exhausted,
159}
160
161impl PrefetchBlockStream {
162    pub(super) fn new(
163        blocks: VecDeque<BlockHolder>,
164        block_index: usize,
165        _tracker: Option<MemoryUsageTracker>,
166    ) -> Self {
167        Self {
168            blocks,
169            block_index,
170            _tracker,
171        }
172    }
173
174    /// Takes the block at the given SST index, discarding earlier buffered blocks.
175    /// A backward lookup leaves the stream unchanged; a lookup past the end exhausts it.
176    pub(super) fn take_block(&mut self, target: usize) -> PrefetchLookup {
177        if target < self.block_index {
178            return PrefetchLookup::BeforeStart;
179        }
180        while let Some(block) = self.blocks.pop_front() {
181            let block_index = self.block_index;
182            self.block_index += 1;
183            if block_index == target {
184                return PrefetchLookup::Hit(block);
185            }
186        }
187        PrefetchLookup::Exhausted
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use crate::hummock::test_utils::test_key_of;
195    use crate::hummock::{Block, BlockBuilder, BlockBuilderOptions};
196
197    fn test_blocks() -> Vec<Arc<Block>> {
198        (10..14)
199            .map(|idx| {
200                let mut builder = BlockBuilder::new(BlockBuilderOptions::default());
201                builder.add_for_test(test_key_of(idx).to_ref(), b"value");
202                let capacity = builder.uncompressed_block_size();
203                Arc::new(Block::decode(Bytes::copy_from_slice(builder.build()), capacity).unwrap())
204            })
205            .collect()
206    }
207
208    #[test]
209    fn test_prefetch_take_block_and_tracker_lifetime() {
210        let blocks = test_blocks();
211        let usage = Arc::new(AtomicUsize::new(7));
212        let mut stream = PrefetchBlockStream::new(
213            blocks
214                .iter()
215                .cloned()
216                .map(BlockHolder::from_ref_block)
217                .collect(),
218            10,
219            Some(MemoryUsageTracker::new(usage.clone(), 100)),
220        );
221        assert!(matches!(stream.take_block(9), PrefetchLookup::BeforeStart));
222        let PrefetchLookup::Hit(first) = stream.take_block(10) else {
223            panic!("missing first block");
224        };
225        assert!(std::ptr::eq(&*first, &*blocks[0]));
226        // Neither rereading a consumed block nor seeking backward may consume future blocks.
227        assert!(matches!(stream.take_block(10), PrefetchLookup::BeforeStart));
228        let PrefetchLookup::Hit(skipped_to) = stream.take_block(12) else {
229            panic!("missing block after a forward skip");
230        };
231        assert!(std::ptr::eq(&*skipped_to, &*blocks[2]));
232        assert_eq!(Arc::strong_count(&blocks[1]), 1);
233        assert!(matches!(stream.take_block(11), PrefetchLookup::BeforeStart));
234        let PrefetchLookup::Hit(last) = stream.take_block(13) else {
235            panic!("backward lookup consumed the remaining block");
236        };
237        assert!(std::ptr::eq(&*last, &*blocks[3]));
238        assert!(matches!(stream.take_block(14), PrefetchLookup::Exhausted));
239        assert!(matches!(
240            stream.take_block(usize::MAX),
241            PrefetchLookup::Exhausted
242        ));
243        assert_eq!(usage.load(Ordering::SeqCst), 107);
244        drop(stream);
245        assert_eq!(usage.load(Ordering::SeqCst), 7);
246        // Returned holders continue to own their blocks after the stream and tracker drop.
247        drop(blocks);
248        assert!(!first.data().is_empty());
249        assert!(!last.data().is_empty());
250    }
251
252    #[test]
253    fn test_prefetch_take_block_past_end_and_empty() {
254        let blocks = test_blocks();
255        let mut stream = PrefetchBlockStream::new(
256            blocks
257                .iter()
258                .cloned()
259                .map(BlockHolder::from_ref_block)
260                .collect(),
261            10,
262            None,
263        );
264        assert!(matches!(
265            stream.take_block(usize::MAX),
266            PrefetchLookup::Exhausted
267        ));
268        assert!(blocks.iter().all(|block| Arc::strong_count(block) == 1));
269        assert!(matches!(stream.take_block(14), PrefetchLookup::Exhausted));
270        assert!(matches!(stream.take_block(13), PrefetchLookup::BeforeStart));
271
272        let mut empty = PrefetchBlockStream::new(VecDeque::new(), 10, None);
273        assert!(matches!(empty.take_block(9), PrefetchLookup::BeforeStart));
274        assert!(matches!(empty.take_block(10), PrefetchLookup::Exhausted));
275        assert!(matches!(empty.take_block(11), PrefetchLookup::Exhausted));
276    }
277}
risingwave_storage/hummock/block_stream.rs

risingwave_storage/hummock/
block_stream.rs
