risingwave_storage/hummock/sstable/
writer.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::sync::Arc;
16
17use bytes::Bytes;
18use fail::fail_point;
19use foyer::HybridCacheProperties;
20use risingwave_hummock_sdk::HummockSstableObjectId;
21use risingwave_object_store::object::ObjectStreamingUploader;
22use tokio::task::JoinHandle;
23use zstd::zstd_safe::WriteBuf;
24
25use super::multi_builder::UploadJoinHandle;
26use super::{Block, BlockMeta};
27use crate::hummock::utils::MemoryTracker;
28use crate::hummock::{
29    CachePolicy, HummockResult, RecentFilterTrait, SstableBlockIndex, SstableBuilderOptions,
30    SstableMeta, SstableStore, SstableStoreRef,
31};
32
33/// A consumer of SST data.
34#[async_trait::async_trait]
35pub trait SstableWriter: Send {
36    type Output;
37
38    /// Write an SST block to the writer.
39    async fn write_block(&mut self, block: &[u8], meta: &BlockMeta) -> HummockResult<()>;
40
41    async fn write_block_bytes(&mut self, block: Bytes, meta: &BlockMeta) -> HummockResult<()>;
42
43    /// Finish writing the SST.
44    async fn finish(self, meta: SstableMeta) -> HummockResult<Self::Output>;
45
46    /// Get the length of data that has already been written.
47    fn data_len(&self) -> usize;
48}
49
50/// Append SST data to a buffer. Used for tests and benchmarks.
51pub struct InMemWriter {
52    buf: Vec<u8>,
53}
54
55impl InMemWriter {
56    pub fn new(capacity: usize) -> Self {
57        Self {
58            buf: Vec::with_capacity(capacity),
59        }
60    }
61}
62
63impl From<&SstableBuilderOptions> for InMemWriter {
64    fn from(options: &SstableBuilderOptions) -> Self {
65        Self::new(options.capacity + options.block_capacity)
66    }
67}
68
69#[async_trait::async_trait]
70impl SstableWriter for InMemWriter {
71    type Output = (Bytes, SstableMeta);
72
73    async fn write_block(&mut self, block: &[u8], _meta: &BlockMeta) -> HummockResult<()> {
74        self.buf.extend_from_slice(block);
75        Ok(())
76    }
77
78    async fn write_block_bytes(&mut self, block: Bytes, _meta: &BlockMeta) -> HummockResult<()> {
79        self.buf.extend_from_slice(&block);
80        Ok(())
81    }
82
83    async fn finish(mut self, meta: SstableMeta) -> HummockResult<Self::Output> {
84        meta.encode_to(&mut self.buf);
85        Ok((Bytes::from(self.buf), meta))
86    }
87
88    fn data_len(&self) -> usize {
89        self.buf.len()
90    }
91}
92
93pub struct SstableWriterOptions {
94    /// Total length of SST data.
95    pub capacity_hint: Option<usize>,
96    pub tracker: Option<MemoryTracker>,
97    pub policy: CachePolicy,
98}
99
100impl Default for SstableWriterOptions {
101    fn default() -> Self {
102        Self {
103            capacity_hint: None,
104            tracker: None,
105            policy: CachePolicy::NotFill,
106        }
107    }
108}
109#[async_trait::async_trait]
110pub trait SstableWriterFactory: Send {
111    type Writer: SstableWriter<Output = UploadJoinHandle>;
112
113    async fn create_sst_writer(
114        &mut self,
115        object_id: impl Into<HummockSstableObjectId> + Send,
116        options: SstableWriterOptions,
117    ) -> HummockResult<Self::Writer>;
118}
119
120pub struct BatchSstableWriterFactory {
121    sstable_store: SstableStoreRef,
122}
123
124impl BatchSstableWriterFactory {
125    pub fn new(sstable_store: SstableStoreRef) -> Self {
126        BatchSstableWriterFactory { sstable_store }
127    }
128}
129
130#[async_trait::async_trait]
131impl SstableWriterFactory for BatchSstableWriterFactory {
132    type Writer = BatchUploadWriter;
133
134    async fn create_sst_writer(
135        &mut self,
136        object_id: impl Into<HummockSstableObjectId> + Send,
137        options: SstableWriterOptions,
138    ) -> HummockResult<Self::Writer> {
139        Ok(BatchUploadWriter::new(
140            object_id,
141            self.sstable_store.clone(),
142            options,
143        ))
144    }
145}
146
147/// Buffer SST data and upload it as a whole on `finish`.
148/// The upload is finished when the returned `JoinHandle` is joined.
149pub struct BatchUploadWriter {
150    object_id: HummockSstableObjectId,
151    sstable_store: SstableStoreRef,
152    policy: CachePolicy,
153    buf: Vec<u8>,
154    block_info: Vec<Block>,
155    tracker: Option<MemoryTracker>,
156}
157
158impl BatchUploadWriter {
159    pub fn new(
160        object_id: impl Into<HummockSstableObjectId>,
161        sstable_store: Arc<SstableStore>,
162        options: SstableWriterOptions,
163    ) -> Self {
164        Self {
165            object_id: object_id.into(),
166            sstable_store,
167            policy: options.policy,
168            buf: Vec::with_capacity(options.capacity_hint.unwrap_or(0)),
169            block_info: Vec::new(),
170            tracker: options.tracker,
171        }
172    }
173}
174
175#[async_trait::async_trait]
176impl SstableWriter for BatchUploadWriter {
177    type Output = JoinHandle<HummockResult<()>>;
178
179    async fn write_block(&mut self, block: &[u8], meta: &BlockMeta) -> HummockResult<()> {
180        self.buf.extend_from_slice(block);
181        if let CachePolicy::Fill(_) = self.policy {
182            self.block_info.push(Block::decode(
183                Bytes::from(block.to_vec()),
184                meta.uncompressed_size as usize,
185            )?);
186        }
187        Ok(())
188    }
189
190    async fn write_block_bytes(&mut self, block: Bytes, meta: &BlockMeta) -> HummockResult<()> {
191        self.buf.extend_from_slice(&block);
192        if let CachePolicy::Fill(_) = self.policy {
193            self.block_info
194                .push(Block::decode(block, meta.uncompressed_size as usize)?);
195        }
196        Ok(())
197    }
198
199    async fn finish(mut self, meta: SstableMeta) -> HummockResult<Self::Output> {
200        fail_point!("data_upload_err");
201        let join_handle = tokio::spawn(async move {
202            meta.encode_to(&mut self.buf);
203            let data = Bytes::from(self.buf);
204            let _tracker = self.tracker.map(|mut t| {
205                if !t.try_increase_memory(data.capacity() as u64) {
206                    tracing::debug!("failed to allocate increase memory for data file, sst object id: {}, file size: {}",
207                                    self.object_id, data.capacity());
208                }
209                t
210            });
211
212            // Upload data to object store.
213            self.sstable_store
214                .clone()
215                .put_sst_data(self.object_id, data)
216                .await?;
217            self.sstable_store.insert_meta_cache(self.object_id, meta);
218
219            // Only update recent filter with sst obj id is okay here, for l0 is only filter by sst obj id with recent filter.
220            self.sstable_store
221                .recent_filter()
222                .insert((self.object_id, usize::MAX));
223
224            // Add block cache.
225            if let CachePolicy::Fill(hint) = self.policy {
226                // The `block_info` may be empty when there is only range-tombstones, because we
227                //  store them in meta-block.
228                for (block_idx, block) in self.block_info.into_iter().enumerate() {
229                    self.sstable_store.block_cache().insert_with_properties(
230                        SstableBlockIndex {
231                            sst_id: self.object_id,
232                            block_idx: block_idx as _,
233                        },
234                        Box::new(block),
235                        HybridCacheProperties::default().with_hint(hint),
236                    );
237                }
238            }
239            Ok(())
240        });
241        Ok(join_handle)
242    }
243
244    fn data_len(&self) -> usize {
245        self.buf.len()
246    }
247}
248
249pub struct StreamingUploadWriter {
250    object_id: HummockSstableObjectId,
251    sstable_store: SstableStoreRef,
252    policy: CachePolicy,
253    /// Data are uploaded block by block, except for the size footer.
254    object_uploader: ObjectStreamingUploader,
255    /// Compressed blocks to refill block or meta cache. Keep the uncompressed capacity for decode.
256    blocks: Vec<Block>,
257    data_len: usize,
258    tracker: Option<MemoryTracker>,
259}
260
261impl StreamingUploadWriter {
262    pub fn new(
263        object_id: HummockSstableObjectId,
264        sstable_store: SstableStoreRef,
265        object_uploader: ObjectStreamingUploader,
266        options: SstableWriterOptions,
267    ) -> Self {
268        Self {
269            object_id,
270            sstable_store,
271            policy: options.policy,
272            object_uploader,
273            blocks: Vec::new(),
274            data_len: 0,
275            tracker: options.tracker,
276        }
277    }
278}
279
280#[async_trait::async_trait]
281impl SstableWriter for StreamingUploadWriter {
282    type Output = JoinHandle<HummockResult<()>>;
283
284    async fn write_block(&mut self, block_data: &[u8], meta: &BlockMeta) -> HummockResult<()> {
285        self.data_len += block_data.len();
286        let block_data = Bytes::from(block_data.to_vec());
287        if let CachePolicy::Fill(_) = self.policy {
288            let block = Block::decode(block_data.clone(), meta.uncompressed_size as usize)?;
289            self.blocks.push(block);
290        }
291        self.object_uploader
292            .write_bytes(block_data)
293            .await
294            .map_err(Into::into)
295    }
296
297    async fn write_block_bytes(&mut self, block: Bytes, meta: &BlockMeta) -> HummockResult<()> {
298        self.data_len += block.len();
299        if let CachePolicy::Fill(_) = self.policy {
300            let block = Block::decode(block.clone(), meta.uncompressed_size as usize)?;
301            self.blocks.push(block);
302        }
303        self.object_uploader
304            .write_bytes(block)
305            .await
306            .map_err(Into::into)
307    }
308
309    async fn finish(mut self, meta: SstableMeta) -> HummockResult<UploadJoinHandle> {
310        let metadata = Bytes::from(meta.encode_to_bytes());
311
312        self.object_uploader.write_bytes(metadata).await?;
313        let join_handle = tokio::spawn(async move {
314            let uploader_memory_usage = self.object_uploader.get_memory_usage();
315            let _tracker = self.tracker.map(|mut t| {
316                    if !t.try_increase_memory(uploader_memory_usage) {
317                        tracing::debug!("failed to allocate increase memory for data file, sst object id: {}, file size: {}",
318                                        self.object_id, uploader_memory_usage);
319                    }
320                    t
321                });
322
323            assert!(!meta.block_metas.is_empty());
324
325            // Upload data to object store.
326            self.object_uploader.finish().await?;
327            // Add meta cache.
328            self.sstable_store.insert_meta_cache(self.object_id, meta);
329
330            // Add block cache.
331            if let CachePolicy::Fill(hint) = self.policy
332                && !self.blocks.is_empty()
333            {
334                for (block_idx, block) in self.blocks.into_iter().enumerate() {
335                    self.sstable_store.block_cache().insert_with_properties(
336                        SstableBlockIndex {
337                            sst_id: self.object_id,
338                            block_idx: block_idx as _,
339                        },
340                        Box::new(block),
341                        HybridCacheProperties::default().with_hint(hint),
342                    );
343                }
344            }
345            Ok(())
346        });
347        Ok(join_handle)
348    }
349
350    fn data_len(&self) -> usize {
351        self.data_len
352    }
353}
354
355pub struct StreamingSstableWriterFactory {
356    sstable_store: SstableStoreRef,
357}
358
359impl StreamingSstableWriterFactory {
360    pub fn new(sstable_store: SstableStoreRef) -> Self {
361        StreamingSstableWriterFactory { sstable_store }
362    }
363}
364
365#[async_trait::async_trait]
366impl SstableWriterFactory for StreamingSstableWriterFactory {
367    type Writer = StreamingUploadWriter;
368
369    async fn create_sst_writer(
370        &mut self,
371        object_id: impl Into<HummockSstableObjectId> + Send,
372        options: SstableWriterOptions,
373    ) -> HummockResult<Self::Writer> {
374        let object_id = object_id.into();
375        let path = self.sstable_store.get_sst_data_path(object_id);
376        let uploader = self.sstable_store.create_streaming_uploader(&path).await?;
377        Ok(StreamingUploadWriter::new(
378            object_id,
379            self.sstable_store.clone(),
380            uploader,
381            options,
382        ))
383    }
384}
385
386#[cfg(test)]
387mod tests {
388
389    use bytes::Bytes;
390    use rand::{Rng, SeedableRng};
391    use risingwave_common::util::iter_util::ZipEqFast;
392
393    use crate::hummock::sstable::VERSION;
394    use crate::hummock::{BlockMeta, InMemWriter, SstableMeta, SstableWriter};
395
396    fn get_sst() -> (Bytes, Vec<Bytes>, SstableMeta) {
397        let mut rng = rand::rngs::StdRng::seed_from_u64(0);
398        let mut buffer: Vec<u8> = vec![0; 5000];
399        rng.fill(&mut buffer[..]);
400        buffer.extend((5_u32).to_le_bytes());
401        let data = Bytes::from(buffer);
402
403        let mut blocks = Vec::with_capacity(5);
404        let mut block_metas = Vec::with_capacity(5);
405        for i in 0..5 {
406            block_metas.push(BlockMeta {
407                smallest_key: Vec::new(),
408                len: 1000,
409                offset: i * 1000,
410                ..Default::default()
411            });
412            blocks.push(data.slice((i * 1000) as usize..((i + 1) * 1000) as usize));
413        }
414        #[expect(deprecated)]
415        let meta = SstableMeta {
416            block_metas,
417            bloom_filter: vec![],
418            estimated_size: 0,
419            key_count: 0,
420            smallest_key: Vec::new(),
421            largest_key: Vec::new(),
422            meta_offset: data.len() as u64,
423            monotonic_tombstone_events: vec![],
424            version: VERSION,
425        };
426
427        (data, blocks, meta)
428    }
429
430    #[tokio::test]
431    async fn test_in_mem_writer() {
432        let (data, blocks, meta) = get_sst();
433        let mut writer = Box::new(InMemWriter::new(0));
434        for (block, meta) in blocks.iter().zip_eq_fast(meta.block_metas.iter()) {
435            writer.write_block(&block[..], meta).await.unwrap();
436        }
437
438        let meta_offset = meta.meta_offset as usize;
439        let (output_data, _) = writer.finish(meta).await.unwrap();
440        assert_eq!(output_data.slice(0..meta_offset), data);
441    }
442}
risingwave_storage/hummock/sstable/writer.rs

risingwave_storage/hummock/sstable/
writer.rs
