risingwave_object_store/object/
mem.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::collections::{HashMap, VecDeque};
16use std::pin::Pin;
17use std::sync::{Arc, LazyLock};
18use std::task::{Context, Poll};
19use std::time::{SystemTime, UNIX_EPOCH};
20
21use bytes::{BufMut, Bytes, BytesMut};
22use fail::fail_point;
23use futures::Stream;
24use itertools::Itertools;
25use risingwave_common::range::RangeBoundsExt;
26use thiserror::Error;
27use tokio::sync::Mutex;
28
29use super::{
30    ObjectError, ObjectMetadata, ObjectRangeBounds, ObjectResult, ObjectStore, StreamingUploader,
31};
32use crate::object::{ObjectDataStream, ObjectMetadataIter};
33
34#[derive(Error, Debug)]
35pub enum Error {
36    #[error("NotFound error: {0}")]
37    NotFound(String),
38    #[error("Other error: {0}")]
39    Other(String),
40}
41
42impl Error {
43    pub fn is_object_not_found_error(&self) -> bool {
44        matches!(self, Error::NotFound(_))
45    }
46}
47
48impl Error {
49    fn not_found(msg: impl ToString) -> Self {
50        Error::NotFound(msg.to_string())
51    }
52
53    fn other(msg: impl ToString) -> Self {
54        Error::Other(msg.to_string())
55    }
56}
57
58/// Store multiple parts in a map, and concatenate them on finish.
59pub struct InMemStreamingUploader {
60    path: String,
61    buf: BytesMut,
62    objects: Arc<Mutex<HashMap<String, (ObjectMetadata, Bytes)>>>,
63}
64
65impl StreamingUploader for InMemStreamingUploader {
66    async fn write_bytes(&mut self, data: Bytes) -> ObjectResult<()> {
67        fail_point!("mem_write_bytes_err", |_| Err(ObjectError::internal(
68            "mem write bytes error"
69        )));
70        self.buf.put(data);
71        Ok(())
72    }
73
74    async fn finish(self) -> ObjectResult<()> {
75        fail_point!("mem_finish_streaming_upload_err", |_| Err(
76            ObjectError::internal("mem finish streaming upload error")
77        ));
78        let obj = self.buf.freeze();
79        if obj.is_empty() {
80            Err(Error::other("upload empty object").into())
81        } else {
82            let metadata = get_obj_meta(&self.path, &obj)?;
83            self.objects.lock().await.insert(self.path, (metadata, obj));
84            Ok(())
85        }
86    }
87
88    fn get_memory_usage(&self) -> u64 {
89        self.buf.capacity() as u64
90    }
91}
92
93/// In-memory object storage, useful for testing.
94#[derive(Default, Clone)]
95pub struct InMemObjectStore {
96    objects: Arc<Mutex<HashMap<String, (ObjectMetadata, Bytes)>>>,
97}
98
99#[async_trait::async_trait]
100impl ObjectStore for InMemObjectStore {
101    type StreamingUploader = InMemStreamingUploader;
102
103    fn get_object_prefix(&self, _obj_id: u64, _use_new_object_prefix_strategy: bool) -> String {
104        String::default()
105    }
106
107    async fn upload(&self, path: &str, obj: Bytes) -> ObjectResult<()> {
108        fail_point!("mem_upload_err", |_| Err(ObjectError::internal(
109            "mem upload error"
110        )));
111        if obj.is_empty() {
112            Err(Error::other("upload empty object").into())
113        } else {
114            let metadata = get_obj_meta(path, &obj)?;
115            self.objects
116                .lock()
117                .await
118                .insert(path.into(), (metadata, obj));
119            Ok(())
120        }
121    }
122
123    async fn streaming_upload(&self, path: &str) -> ObjectResult<Self::StreamingUploader> {
124        Ok(InMemStreamingUploader {
125            path: path.to_owned(),
126            buf: BytesMut::new(),
127            objects: self.objects.clone(),
128        })
129    }
130
131    async fn read(&self, path: &str, range: impl ObjectRangeBounds) -> ObjectResult<Bytes> {
132        fail_point!("mem_read_err", |_| Err(ObjectError::internal(
133            "mem read error"
134        )));
135        self.get_object(path, range).await
136    }
137
138    /// Returns a stream reading the object specified in `path`. If given, the stream starts at the
139    /// byte with index `start_pos` (0-based). As far as possible, the stream only loads the amount
140    /// of data into memory that is read from the stream.
141    async fn streaming_read(
142        &self,
143        path: &str,
144        read_range: impl ObjectRangeBounds,
145    ) -> ObjectResult<ObjectDataStream> {
146        fail_point!("mem_streaming_read_err", |_| Err(ObjectError::internal(
147            "mem streaming read error"
148        )));
149        let bytes = self.get_object(path, read_range).await?;
150
151        Ok(Box::pin(InMemDataIterator::new(bytes)))
152    }
153
154    async fn metadata(&self, path: &str) -> ObjectResult<ObjectMetadata> {
155        self.objects
156            .lock()
157            .await
158            .get(path)
159            .map(|(metadata, _)| metadata)
160            .cloned()
161            .ok_or_else(|| Error::not_found(format!("no object at path '{}'", path)).into())
162    }
163
164    async fn delete(&self, path: &str) -> ObjectResult<()> {
165        fail_point!("mem_delete_err", |_| Err(ObjectError::internal(
166            "mem delete error"
167        )));
168        self.objects.lock().await.remove(path);
169        Ok(())
170    }
171
172    /// Deletes the objects with the given paths permanently from the storage. If an object
173    /// specified in the request is not found, it will be considered as successfully deleted.
174    async fn delete_objects(&self, paths: &[String]) -> ObjectResult<()> {
175        let mut guard = self.objects.lock().await;
176
177        for path in paths {
178            guard.remove(path);
179        }
180
181        Ok(())
182    }
183
184    async fn list(
185        &self,
186        prefix: &str,
187        start_after: Option<String>,
188        limit: Option<usize>,
189    ) -> ObjectResult<ObjectMetadataIter> {
190        let list_result = self
191            .objects
192            .lock()
193            .await
194            .iter()
195            .filter_map(|(path, (metadata, _))| {
196                if let Some(ref start_after) = start_after
197                    && metadata.key.le(start_after)
198                {
199                    return None;
200                }
201                if path.starts_with(prefix) {
202                    return Some(metadata.clone());
203                }
204                None
205            })
206            .sorted_by(|a, b| Ord::cmp(&a.key, &b.key))
207            .take(limit.unwrap_or(usize::MAX))
208            .collect_vec();
209        Ok(Box::pin(InMemObjectIter::new(list_result)))
210    }
211
212    fn store_media_type(&self) -> &'static str {
213        "mem"
214    }
215}
216
217pub struct InMemDataIterator {
218    data: Bytes,
219    offset: usize,
220}
221
222impl InMemDataIterator {
223    pub fn new(data: Bytes) -> Self {
224        Self { data, offset: 0 }
225    }
226}
227
228impl Stream for InMemDataIterator {
229    type Item = ObjectResult<Bytes>;
230
231    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
232        const MAX_PACKET_SIZE: usize = 128 * 1024;
233        if self.offset >= self.data.len() {
234            return Poll::Ready(None);
235        }
236        let read_len = std::cmp::min(self.data.len() - self.offset, MAX_PACKET_SIZE);
237        let data = self.data.slice(self.offset..(self.offset + read_len));
238        self.offset += read_len;
239        Poll::Ready(Some(Ok(data)))
240    }
241}
242
243static SHARED: LazyLock<spin::Mutex<InMemObjectStore>> =
244    LazyLock::new(|| spin::Mutex::new(InMemObjectStore::new()));
245
246impl InMemObjectStore {
247    fn new() -> Self {
248        Self {
249            objects: Arc::new(Mutex::new(HashMap::new())),
250        }
251    }
252
253    /// Create a new in-memory object store for testing, isolated with others.
254    pub fn for_test() -> Self {
255        Self::new()
256    }
257
258    /// Get a reference to the in-memory object store shared in this process.
259    ///
260    /// Note: Should only be used for `risedev playground`, when there're multiple compute-nodes or
261    /// compactors in the same process.
262    pub(super) fn shared() -> Self {
263        SHARED.lock().clone()
264    }
265
266    /// Reset the shared in-memory object store.
267    pub fn reset_shared() {
268        *SHARED.lock() = InMemObjectStore::new();
269    }
270
271    async fn get_object(&self, path: &str, range: impl ObjectRangeBounds) -> ObjectResult<Bytes> {
272        let objects = self.objects.lock().await;
273
274        let obj = objects
275            .get(path)
276            .map(|(_, obj)| obj)
277            .ok_or_else(|| Error::not_found(format!("no object at path '{}'", path)))?;
278
279        if let Some(end) = range.end()
280            && end > obj.len()
281        {
282            return Err(Error::other("bad block offset and size").into());
283        }
284
285        Ok(obj.slice(range))
286    }
287}
288
289fn get_obj_meta(path: &str, obj: &Bytes) -> ObjectResult<ObjectMetadata> {
290    Ok(ObjectMetadata {
291        key: path.to_owned(),
292        last_modified: SystemTime::now()
293            .duration_since(UNIX_EPOCH)
294            .map_err(ObjectError::internal)?
295            .as_secs_f64(),
296        total_size: obj.len(),
297    })
298}
299
300struct InMemObjectIter {
301    list_result: VecDeque<ObjectMetadata>,
302}
303
304impl InMemObjectIter {
305    fn new(list_result: Vec<ObjectMetadata>) -> Self {
306        Self {
307            list_result: list_result.into(),
308        }
309    }
310}
311
312impl Stream for InMemObjectIter {
313    type Item = ObjectResult<ObjectMetadata>;
314
315    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
316        if let Some(i) = self.list_result.pop_front() {
317            return Poll::Ready(Some(Ok(i)));
318        }
319        Poll::Ready(None)
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use futures::TryStreamExt;
326    use itertools::enumerate;
327
328    use super::*;
329
330    #[tokio::test]
331    async fn test_upload() {
332        let block = Bytes::from("123456");
333
334        let s3 = InMemObjectStore::for_test();
335        s3.upload("/abc", block).await.unwrap();
336
337        // No such object.
338        let err = s3.read("/ab", 0..3).await.unwrap_err();
339        assert!(err.is_object_not_found_error());
340
341        let bytes = s3.read("/abc", 4..6).await.unwrap();
342        assert_eq!(String::from_utf8(bytes.to_vec()).unwrap(), "56".to_owned());
343
344        // Overflow.
345        s3.read("/abc", 4..8).await.unwrap_err();
346
347        s3.delete("/abc").await.unwrap();
348
349        // No such object.
350        s3.read("/abc", 0..3).await.unwrap_err();
351    }
352
353    #[tokio::test]
354    async fn test_streaming_upload() {
355        let blocks = vec![Bytes::from("123"), Bytes::from("456"), Bytes::from("789")];
356        let obj = Bytes::from("123456789");
357
358        let store = InMemObjectStore::for_test();
359        let mut uploader = store.streaming_upload("/abc").await.unwrap();
360
361        for block in blocks {
362            uploader.write_bytes(block).await.unwrap();
363        }
364        uploader.finish().await.unwrap();
365
366        // Read whole object.
367        let read_obj = store.read("/abc", ..).await.unwrap();
368        assert!(read_obj.eq(&obj));
369
370        // Read part of the object.
371        let read_obj = store.read("/abc", 4..6).await.unwrap();
372        assert_eq!(
373            String::from_utf8(read_obj.to_vec()).unwrap(),
374            "56".to_owned()
375        );
376    }
377
378    #[tokio::test]
379    async fn test_metadata() {
380        let block = Bytes::from("123456");
381
382        let obj_store = InMemObjectStore::for_test();
383        obj_store.upload("/abc", block).await.unwrap();
384
385        let err = obj_store.metadata("/not_exist").await.unwrap_err();
386        assert!(err.is_object_not_found_error());
387
388        let metadata = obj_store.metadata("/abc").await.unwrap();
389        assert_eq!(metadata.total_size, 6);
390    }
391
392    async fn list_all(prefix: &str, store: &InMemObjectStore) -> Vec<ObjectMetadata> {
393        store
394            .list(prefix, None, None)
395            .await
396            .unwrap()
397            .try_collect::<Vec<_>>()
398            .await
399            .unwrap()
400    }
401
402    #[tokio::test]
403    async fn test_list() {
404        let payload = Bytes::from("123456");
405        let store = InMemObjectStore::for_test();
406        assert!(list_all("", &store).await.is_empty());
407
408        let paths = vec!["001/002/test.obj", "001/003/test.obj"];
409        for (i, path) in enumerate(paths.clone()) {
410            assert_eq!(list_all("", &store).await.len(), i);
411            store.upload(path, payload.clone()).await.unwrap();
412            assert_eq!(list_all("", &store).await.len(), i + 1);
413        }
414
415        let list_path = list_all("", &store)
416            .await
417            .iter()
418            .map(|p| p.key.clone())
419            .collect_vec();
420        assert_eq!(list_path, paths);
421
422        for i in 0..=5 {
423            assert_eq!(list_all(&paths[0][0..=i], &store).await.len(), 2);
424        }
425        for i in 6..=paths[0].len() - 1 {
426            assert_eq!(list_all(&paths[0][0..=i], &store).await.len(), 1)
427        }
428        assert!(list_all("003", &store).await.is_empty());
429
430        for (i, path) in enumerate(paths.clone()) {
431            assert_eq!(list_all("", &store).await.len(), paths.len() - i);
432            store.delete(path).await.unwrap();
433            assert_eq!(list_all("", &store).await.len(), paths.len() - i - 1);
434        }
435    }
436
437    #[tokio::test]
438    async fn test_delete_objects() {
439        let block1 = Bytes::from("123456");
440        let block2 = Bytes::from("987654");
441
442        let store = InMemObjectStore::for_test();
443        store.upload("/abc", block1).await.unwrap();
444        store.upload("/klm", block2).await.unwrap();
445
446        assert_eq!(list_all("", &store).await.len(), 2);
447
448        let str_list = [
449            String::from("/abc"),
450            String::from("/klm"),
451            String::from("/xyz"),
452        ];
453
454        store.delete_objects(&str_list).await.unwrap();
455
456        assert_eq!(list_all("", &store).await.len(), 0);
457    }
458}
risingwave_object_store/object/mem.rs

risingwave_object_store/object/
mem.rs
