object_store/
memory.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! An in-memory object store implementation
19use std::collections::{BTreeMap, BTreeSet, HashMap};
20use std::ops::Range;
21use std::sync::Arc;
22
23use async_trait::async_trait;
24use bytes::Bytes;
25use chrono::{DateTime, Utc};
26use futures_util::{StreamExt, stream::BoxStream};
27use parking_lot::RwLock;
28
29use crate::multipart::{MultipartStore, PartId};
30use crate::util::InvalidGetRange;
31use crate::{
32    Attributes, GetRange, GetResult, GetResultPayload, ListResult, MultipartId, MultipartUpload,
33    ObjectMeta, ObjectStore, PutMode, PutMultipartOptions, PutOptions, PutResult, Result,
34    UpdateVersion, UploadPart, path::Path,
35};
36use crate::{CopyMode, CopyOptions, GetOptions, PutPayload};
37
38/// A specialized `Error` for in-memory object store-related errors
39#[derive(Debug, thiserror::Error)]
40enum Error {
41    #[error("No data in memory found. Location: {path}")]
42    NoDataInMemory { path: String },
43
44    #[error("Invalid range: {source}")]
45    Range { source: InvalidGetRange },
46
47    #[error("Object already exists at that location: {path}")]
48    AlreadyExists { path: String },
49
50    #[error("ETag required for conditional update")]
51    MissingETag,
52
53    #[error("MultipartUpload not found: {id}")]
54    UploadNotFound { id: String },
55
56    #[error("Missing part at index: {part}")]
57    MissingPart { part: usize },
58}
59
60impl From<Error> for super::Error {
61    fn from(source: Error) -> Self {
62        match source {
63            Error::NoDataInMemory { ref path } => Self::NotFound {
64                path: path.into(),
65                source: source.into(),
66            },
67            Error::AlreadyExists { ref path } => Self::AlreadyExists {
68                path: path.into(),
69                source: source.into(),
70            },
71            _ => Self::Generic {
72                store: "InMemory",
73                source: Box::new(source),
74            },
75        }
76    }
77}
78
79/// In-memory storage suitable for testing or for opting out of using a cloud
80/// storage provider.
81#[derive(Debug, Default, Clone)]
82pub struct InMemory {
83    storage: SharedStorage,
84}
85
86#[derive(Debug, Clone)]
87struct Entry {
88    data: Bytes,
89    last_modified: DateTime<Utc>,
90    attributes: Attributes,
91    e_tag: usize,
92}
93
94impl Entry {
95    fn new(
96        data: Bytes,
97        last_modified: DateTime<Utc>,
98        e_tag: usize,
99        attributes: Attributes,
100    ) -> Self {
101        Self {
102            data,
103            last_modified,
104            e_tag,
105            attributes,
106        }
107    }
108}
109
110#[derive(Debug, Default, Clone)]
111struct Storage {
112    next_etag: usize,
113    map: BTreeMap<Path, Entry>,
114    uploads: HashMap<usize, PartStorage>,
115}
116
117#[derive(Debug, Default, Clone)]
118struct PartStorage {
119    attributes: Attributes,
120    parts: Vec<Option<Bytes>>,
121}
122
123type SharedStorage = Arc<RwLock<Storage>>;
124
125impl Storage {
126    fn insert(&mut self, location: &Path, bytes: Bytes, attributes: Attributes) -> usize {
127        let etag = self.next_etag;
128        self.next_etag += 1;
129        let entry = Entry::new(bytes, Utc::now(), etag, attributes);
130        self.overwrite(location, entry);
131        etag
132    }
133
134    fn overwrite(&mut self, location: &Path, entry: Entry) {
135        self.map.insert(location.clone(), entry);
136    }
137
138    fn create(&mut self, location: &Path, entry: Entry) -> Result<()> {
139        use std::collections::btree_map;
140        match self.map.entry(location.clone()) {
141            btree_map::Entry::Occupied(_) => Err(Error::AlreadyExists {
142                path: location.to_string(),
143            }
144            .into()),
145            btree_map::Entry::Vacant(v) => {
146                v.insert(entry);
147                Ok(())
148            }
149        }
150    }
151
152    fn update(&mut self, location: &Path, v: UpdateVersion, entry: Entry) -> Result<()> {
153        match self.map.get_mut(location) {
154            // Return Precondition instead of NotFound for consistency with stores
155            None => Err(crate::Error::Precondition {
156                path: location.to_string(),
157                source: format!("Object at location {location} not found").into(),
158            }),
159            Some(e) => {
160                let existing = e.e_tag.to_string();
161                let expected = v.e_tag.ok_or(Error::MissingETag)?;
162                if existing == expected {
163                    *e = entry;
164                    Ok(())
165                } else {
166                    Err(crate::Error::Precondition {
167                        path: location.to_string(),
168                        source: format!("{existing} does not match {expected}").into(),
169                    })
170                }
171            }
172        }
173    }
174
175    fn upload_mut(&mut self, id: &MultipartId) -> Result<&mut PartStorage> {
176        let parts = id
177            .parse()
178            .ok()
179            .and_then(|x| self.uploads.get_mut(&x))
180            .ok_or_else(|| Error::UploadNotFound { id: id.into() })?;
181        Ok(parts)
182    }
183
184    fn remove_upload(&mut self, id: &MultipartId) -> Result<PartStorage> {
185        let parts = id
186            .parse()
187            .ok()
188            .and_then(|x| self.uploads.remove(&x))
189            .ok_or_else(|| Error::UploadNotFound { id: id.into() })?;
190        Ok(parts)
191    }
192}
193
194impl std::fmt::Display for InMemory {
195    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
196        write!(f, "InMemory")
197    }
198}
199
200#[async_trait]
201impl ObjectStore for InMemory {
202    async fn put_opts(
203        &self,
204        location: &Path,
205        payload: PutPayload,
206        opts: PutOptions,
207    ) -> Result<PutResult> {
208        let mut storage = self.storage.write();
209        let etag = storage.next_etag;
210        let entry = Entry::new(payload.into(), Utc::now(), etag, opts.attributes);
211
212        match opts.mode {
213            PutMode::Overwrite => storage.overwrite(location, entry),
214            PutMode::Create => storage.create(location, entry)?,
215            PutMode::Update(v) => storage.update(location, v, entry)?,
216        }
217        storage.next_etag += 1;
218
219        Ok(PutResult {
220            e_tag: Some(etag.to_string()),
221            version: None,
222            extensions: Default::default(),
223        })
224    }
225
226    async fn put_multipart_opts(
227        &self,
228        location: &Path,
229        opts: PutMultipartOptions,
230    ) -> Result<Box<dyn MultipartUpload>> {
231        Ok(Box::new(InMemoryUpload {
232            location: location.clone(),
233            attributes: opts.attributes,
234            parts: vec![],
235            storage: Arc::clone(&self.storage),
236        }))
237    }
238
239    async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
240        let entry = self.entry(location)?;
241        let e_tag = entry.e_tag.to_string();
242
243        let meta = ObjectMeta {
244            location: location.clone(),
245            last_modified: entry.last_modified,
246            size: entry.data.len() as u64,
247            e_tag: Some(e_tag),
248            version: None,
249        };
250        options.check_preconditions(&meta)?;
251
252        let (range, data) = match options.range {
253            Some(range) => {
254                let r = range
255                    .as_range(entry.data.len() as u64)
256                    .map_err(|source| Error::Range { source })?;
257                (
258                    r.clone(),
259                    entry.data.slice(r.start as usize..r.end as usize),
260                )
261            }
262            None => (0..entry.data.len() as u64, entry.data),
263        };
264        let stream = futures_util::stream::once(futures_util::future::ready(Ok(data)));
265
266        Ok(GetResult {
267            payload: GetResultPayload::Stream(stream.boxed()),
268            attributes: entry.attributes,
269            meta,
270            range,
271            extensions: Default::default(),
272        })
273    }
274
275    async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
276        let entry = self.entry(location)?;
277        ranges
278            .iter()
279            .map(|range| {
280                let r = GetRange::Bounded(range.clone())
281                    .as_range(entry.data.len() as u64)
282                    .map_err(|source| Error::Range { source })?;
283                let r_end = usize::try_from(r.end).map_err(|_e| Error::Range {
284                    source: InvalidGetRange::TooLarge {
285                        requested: r.end,
286                        max: usize::MAX as u64,
287                    },
288                })?;
289                let r_start = usize::try_from(r.start).map_err(|_e| Error::Range {
290                    source: InvalidGetRange::TooLarge {
291                        requested: r.start,
292                        max: usize::MAX as u64,
293                    },
294                })?;
295                Ok(entry.data.slice(r_start..r_end))
296            })
297            .collect()
298    }
299
300    fn delete_stream(
301        &self,
302        locations: BoxStream<'static, Result<Path>>,
303    ) -> BoxStream<'static, Result<Path>> {
304        let storage = Arc::clone(&self.storage);
305        locations
306            .map(move |location| {
307                let location = location?;
308                storage.write().map.remove(&location);
309                Ok(location)
310            })
311            .boxed()
312    }
313
314    fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
315        let root = Path::default();
316        let prefix = prefix.unwrap_or(&root);
317
318        let storage = self.storage.read();
319        let values: Vec<_> = storage
320            .map
321            .range((prefix)..)
322            .take_while(|(key, _)| key.as_ref().starts_with(prefix.as_ref()))
323            .filter(|(key, _)| {
324                // Don't return for exact prefix match
325                key.prefix_match(prefix)
326                    .map(|mut x| x.next().is_some())
327                    .unwrap_or(false)
328            })
329            .map(|(key, value)| {
330                Ok(ObjectMeta {
331                    location: key.clone(),
332                    last_modified: value.last_modified,
333                    size: value.data.len() as u64,
334                    e_tag: Some(value.e_tag.to_string()),
335                    version: None,
336                })
337            })
338            .collect();
339
340        futures_util::stream::iter(values).boxed()
341    }
342
343    /// The memory implementation returns all results, as opposed to the cloud
344    /// versions which limit their results to 1k or more because of API
345    /// limitations.
346    async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
347        let root = Path::default();
348        let prefix = prefix.unwrap_or(&root);
349
350        let mut common_prefixes = BTreeSet::new();
351
352        // Only objects in this base level should be returned in the
353        // response. Otherwise, we just collect the common prefixes.
354        let mut objects = vec![];
355        for (k, v) in self.storage.read().map.range((prefix)..) {
356            if !k.as_ref().starts_with(prefix.as_ref()) {
357                break;
358            }
359
360            let mut parts = match k.prefix_match(prefix) {
361                Some(parts) => parts,
362                None => continue,
363            };
364
365            // Pop first element
366            let common_prefix = match parts.next() {
367                Some(p) => p,
368                // Should only return children of the prefix
369                None => continue,
370            };
371
372            if parts.next().is_some() {
373                common_prefixes.insert(prefix.clone().join(common_prefix));
374            } else {
375                let object = ObjectMeta {
376                    location: k.clone(),
377                    last_modified: v.last_modified,
378                    size: v.data.len() as u64,
379                    e_tag: Some(v.e_tag.to_string()),
380                    version: None,
381                };
382                objects.push(object);
383            }
384        }
385
386        Ok(ListResult {
387            objects,
388            common_prefixes: common_prefixes.into_iter().collect(),
389            extensions: Default::default(),
390        })
391    }
392
393    async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> {
394        let CopyOptions {
395            mode,
396            extensions: _,
397        } = options;
398
399        let entry = self.entry(from)?;
400        let mut storage = self.storage.write();
401
402        match mode {
403            CopyMode::Overwrite => {
404                storage.insert(to, entry.data, entry.attributes);
405            }
406            CopyMode::Create => {
407                if storage.map.contains_key(to) {
408                    return Err(Error::AlreadyExists {
409                        path: to.to_string(),
410                    }
411                    .into());
412                }
413                storage.insert(to, entry.data, entry.attributes);
414            }
415        }
416
417        Ok(())
418    }
419}
420
421#[async_trait]
422impl MultipartStore for InMemory {
423    async fn create_multipart(&self, _path: &Path) -> Result<MultipartId> {
424        self.create_multipart_opts(_path, PutMultipartOptions::default())
425            .await
426    }
427
428    async fn create_multipart_opts(
429        &self,
430        _path: &Path,
431        opts: PutMultipartOptions,
432    ) -> Result<MultipartId> {
433        let mut storage = self.storage.write();
434        let etag = storage.next_etag;
435        storage.next_etag += 1;
436        storage.uploads.insert(
437            etag,
438            PartStorage {
439                attributes: opts.attributes,
440                parts: Default::default(),
441            },
442        );
443        Ok(etag.to_string())
444    }
445
446    async fn put_part(
447        &self,
448        _path: &Path,
449        id: &MultipartId,
450        part_idx: usize,
451        payload: PutPayload,
452    ) -> Result<PartId> {
453        let mut storage = self.storage.write();
454        let upload = storage.upload_mut(id)?;
455        if part_idx >= upload.parts.len() {
456            upload.parts.resize(part_idx + 1, None);
457        }
458        upload.parts[part_idx] = Some(payload.into());
459        Ok(PartId {
460            content_id: Default::default(),
461        })
462    }
463
464    async fn complete_multipart(
465        &self,
466        path: &Path,
467        id: &MultipartId,
468        _parts: Vec<PartId>,
469    ) -> Result<PutResult> {
470        let mut storage = self.storage.write();
471        let upload = storage.remove_upload(id)?;
472
473        let mut cap = 0;
474        for (part, x) in upload.parts.iter().enumerate() {
475            cap += x.as_ref().ok_or(Error::MissingPart { part })?.len();
476        }
477        let mut buf = Vec::with_capacity(cap);
478        for x in &upload.parts {
479            buf.extend_from_slice(x.as_ref().unwrap())
480        }
481        let etag = storage.insert(path, buf.into(), upload.attributes);
482        Ok(PutResult {
483            e_tag: Some(etag.to_string()),
484            version: None,
485            extensions: Default::default(),
486        })
487    }
488
489    async fn abort_multipart(&self, _path: &Path, id: &MultipartId) -> Result<()> {
490        self.storage.write().remove_upload(id)?;
491        Ok(())
492    }
493}
494
495impl InMemory {
496    /// Create new in-memory storage.
497    pub fn new() -> Self {
498        Self::default()
499    }
500
501    /// Creates a fork of the store, with the current content copied into the
502    /// new store.
503    pub fn fork(&self) -> Self {
504        let storage = self.storage.read();
505        let storage = Arc::new(RwLock::new(storage.clone()));
506        Self { storage }
507    }
508
509    fn entry(&self, location: &Path) -> Result<Entry> {
510        let storage = self.storage.read();
511        let value = storage
512            .map
513            .get(location)
514            .cloned()
515            .ok_or_else(|| Error::NoDataInMemory {
516                path: location.to_string(),
517            })?;
518
519        Ok(value)
520    }
521}
522
523#[derive(Debug)]
524struct InMemoryUpload {
525    location: Path,
526    attributes: Attributes,
527    parts: Vec<PutPayload>,
528    storage: Arc<RwLock<Storage>>,
529}
530
531#[async_trait]
532impl MultipartUpload for InMemoryUpload {
533    fn put_part(&mut self, payload: PutPayload) -> UploadPart {
534        self.parts.push(payload);
535        Box::pin(futures_util::future::ready(Ok(())))
536    }
537
538    async fn complete(&mut self) -> Result<PutResult> {
539        let cap = self.parts.iter().map(|x| x.content_length()).sum();
540        let mut buf = Vec::with_capacity(cap);
541        let parts = self.parts.iter().flatten();
542        parts.for_each(|x| buf.extend_from_slice(x));
543        let etag = self.storage.write().insert(
544            &self.location,
545            buf.into(),
546            std::mem::take(&mut self.attributes),
547        );
548
549        Ok(PutResult {
550            e_tag: Some(etag.to_string()),
551            version: None,
552            extensions: Default::default(),
553        })
554    }
555
556    async fn abort(&mut self) -> Result<()> {
557        Ok(())
558    }
559}
560
561#[cfg(test)]
562mod tests {
563    use crate::{ObjectStoreExt, integration::*};
564
565    use super::*;
566
567    #[tokio::test]
568    async fn in_memory_test() {
569        let integration = InMemory::new();
570
571        put_get_delete_list(&integration).await;
572        list_with_offset_exclusivity(&integration).await;
573        get_opts(&integration).await;
574        list_uses_directories_correctly(&integration).await;
575        list_with_delimiter(&integration).await;
576        rename_and_copy(&integration).await;
577        copy_if_not_exists(&integration).await;
578        stream_get(&integration).await;
579        put_opts(&integration, true).await;
580        multipart(&integration, &integration).await;
581        multipart_with_opts(&integration, &integration).await;
582        put_get_attributes(&integration).await;
583        multipart_put_part_out_of_order(&integration, &integration).await;
584    }
585
586    #[tokio::test]
587    async fn box_test() {
588        let integration: Box<dyn ObjectStore> = Box::new(InMemory::new());
589
590        put_get_delete_list(&integration).await;
591        list_with_offset_exclusivity(&integration).await;
592        get_opts(&integration).await;
593        list_uses_directories_correctly(&integration).await;
594        list_with_delimiter(&integration).await;
595        rename_and_copy(&integration).await;
596        copy_if_not_exists(&integration).await;
597        stream_get(&integration).await;
598    }
599
600    #[tokio::test]
601    async fn arc_test() {
602        let integration: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
603
604        put_get_delete_list(&integration).await;
605        list_with_offset_exclusivity(&integration).await;
606        get_opts(&integration).await;
607        list_uses_directories_correctly(&integration).await;
608        list_with_delimiter(&integration).await;
609        rename_and_copy(&integration).await;
610        copy_if_not_exists(&integration).await;
611        stream_get(&integration).await;
612    }
613
614    #[tokio::test]
615    async fn unknown_length() {
616        let integration = InMemory::new();
617
618        let location = Path::from("some_file");
619
620        let data = Bytes::from("arbitrary data");
621
622        integration
623            .put(&location, data.clone().into())
624            .await
625            .unwrap();
626
627        let read_data = integration
628            .get(&location)
629            .await
630            .unwrap()
631            .bytes()
632            .await
633            .unwrap();
634        assert_eq!(&*read_data, data);
635    }
636
637    const NON_EXISTENT_NAME: &str = "nonexistentname";
638
639    #[tokio::test]
640    async fn nonexistent_location() {
641        let integration = InMemory::new();
642
643        let location = Path::from(NON_EXISTENT_NAME);
644
645        let err = get_nonexistent_object(&integration, Some(location))
646            .await
647            .unwrap_err();
648        if let crate::Error::NotFound { path, source } = err {
649            let source_variant = source.downcast_ref::<Error>();
650            assert!(
651                matches!(source_variant, Some(Error::NoDataInMemory { .. }),),
652                "got: {source_variant:?}"
653            );
654            assert_eq!(path, NON_EXISTENT_NAME);
655        } else {
656            panic!("unexpected error type: {err:?}");
657        }
658    }
659}
object_store/memory.rs

object_store/
memory.rs
