objectstore_service/backend/
in_memory.rs

1//! In-memory backend for tests.
2//!
3//! This provides a [`Backend`](super::common::Backend) backed by a `HashMap`,
4//! removing the need for filesystem tempdir management in unit tests. The
5//! backend is [`Clone`] so tests can hold a handle for direct inspection while
6//! the service owns a boxed copy.
7
8use std::collections::HashMap;
9use std::num::NonZeroU64;
10use std::sync::{Arc, Mutex};
11
12use objectstore_types::range::ByteRange;
13use objectstore_types::time::Timestamp;
14
15use bytes::{Bytes, BytesMut};
16use futures_util::TryStreamExt;
17use objectstore_types::metadata::Metadata;
18
19use crate::backend::common::{
20    self, DeleteResponse, ExpiryUpdate, GetResponse, HighVolumeBackend, PutResponse,
21    SetExpiryResponse, TieredGet, TieredMetadata, TieredUpdate, TieredWrite, Tombstone,
22};
23use crate::change_stream::{ChangeStream, NoopStream, flush_change_stream};
24use crate::error::{ErrorKind, Result};
25use crate::id::ObjectId;
26use crate::resumable::{BackendToken, Session, UploadProgress};
27use crate::stream::ClientStream;
28
29/// An entry in the in-memory store.
30#[derive(Clone, Debug)]
31enum StoreEntry {
32    Object(Metadata, Bytes),
33    Tombstone(Tombstone),
34}
35
36impl StoreEntry {
37    fn is_expired(&self, now: Timestamp) -> bool {
38        match self {
39            StoreEntry::Object(metadata, _) => metadata.is_expired(now),
40            StoreEntry::Tombstone(tombstone) => tombstone.is_expired(now),
41        }
42    }
43
44    /// Number of bytes an entry occupies (payload plus serialized metadata)
45    pub(crate) fn stored_size(&self) -> usize {
46        match self {
47            StoreEntry::Object(metadata, payload) => json_len(metadata) + payload.len(),
48            // A tombstone carries no payload: its bytes are the redirect target plus the
49            // deadline stored alongside it.
50            StoreEntry::Tombstone(tombstone) => {
51                tombstone.target.as_storage_path().to_string().len()
52                    + tombstone.time_expires.map_or(0, |d| json_len(&d))
53            }
54        }
55    }
56}
57
58type Store = HashMap<ObjectId, StoreEntry>;
59
60#[derive(Debug)]
61struct ResumableUpload {
62    metadata: Metadata,
63    data: BytesMut,
64}
65
66// None marks a removed session for operations already waiting on its mutex.
67type ResumableSession = Arc<tokio::sync::Mutex<Option<ResumableUpload>>>;
68type ResumableStore = HashMap<(ObjectId, BackendToken), ResumableSession>;
69
70/// In-memory [`Backend`](super::common::Backend) backed by a `HashMap`.
71///
72/// Removes the need for filesystem tempdir management in unit tests. The
73/// backend is [`Clone`] so tests can hold a handle for direct inspection while
74/// the service owns a boxed copy.
75#[derive(Debug, Clone)]
76pub struct InMemoryBackend {
77    name: &'static str,
78    store: Arc<Mutex<Store>>,
79    resumable_store: Arc<Mutex<ResumableStore>>,
80    upload_markers: Arc<Mutex<HashMap<ObjectId, Timestamp>>>,
81    change_stream: Arc<dyn ChangeStream>,
82}
83
84impl InMemoryBackend {
85    /// Creates a new `InMemoryBackend` with the given diagnostic `name`.
86    pub fn new(name: &'static str) -> Self {
87        Self {
88            name,
89            store: Arc::new(Mutex::new(HashMap::new())),
90            resumable_store: Arc::new(Mutex::new(HashMap::new())),
91            upload_markers: Arc::new(Mutex::new(HashMap::new())),
92            change_stream: Arc::new(NoopStream),
93        }
94    }
95
96    fn upload_session(&self, session: &Session) -> Result<ResumableSession> {
97        self.resumable_store
98            .lock()
99            .unwrap()
100            .get(&(session.object_id.clone(), session.backend_token.clone()))
101            // Clone the session handle so that the `resumable_store` mutex is released immediately.
102            .cloned()
103            .ok_or_else(|| ErrorKind::UnknownUploadSession.into())
104    }
105
106    /// Publishes this backend's changes to `change_stream`.
107    pub fn with_change_stream(mut self, change_stream: Arc<dyn ChangeStream>) -> Self {
108        self.change_stream = change_stream;
109        self
110    }
111
112    /// Applies `extend` to `id`'s entry, calling [`ChangeStream::begin_update`] first.
113    ///
114    /// Like a conditional write, this is rejected if the entry changed in the meantime.
115    async fn update_expiry(
116        &self,
117        id: &ObjectId,
118        access_time: Timestamp,
119        extend: impl Fn(&mut StoreEntry) -> Result<ExpiryOutcome> + Send + Sync,
120    ) -> Result<SetExpiryResponse> {
121        let try_extend = |store: &Store| -> Result<(ExpiryOutcome, Option<StoreEntry>)> {
122            match store.get(id) {
123                Some(entry) if !entry.is_expired(access_time) => {
124                    let mut updated = entry.clone();
125                    let outcome = extend(&mut updated)?;
126                    Ok((outcome, Some(updated)))
127                }
128                _ => Ok((ExpiryOutcome::NotFound, None)),
129            }
130        };
131
132        let (outcome, _) = try_extend(&self.store.lock().unwrap())?;
133        let ExpiryOutcome::Extended(expire_at) = outcome else {
134            return Ok(outcome.response());
135        };
136        self.change_stream.begin_update(id, Some(expire_at)).await?;
137
138        {
139            let mut store = self.store.lock().unwrap();
140            match try_extend(&store)? {
141                (ExpiryOutcome::Extended(t), Some(updated)) if t == expire_at => {
142                    store.insert(id.clone(), updated);
143                }
144                _ => return Ok(SetExpiryResponse::Rejected),
145            }
146        }
147        self.change_stream.commit_update(id, Some(expire_at)).await;
148        Ok(outcome.response())
149    }
150
151    /// Returns the stored entry for `id`, for direct inspection in tests.
152    pub fn get(&self, id: &ObjectId) -> Entry {
153        match self.store.lock().unwrap().get(id).cloned() {
154            None => Entry::NotFound,
155            Some(StoreEntry::Tombstone(tombstone)) => Entry::Tombstone(tombstone),
156            Some(StoreEntry::Object(metadata, bytes)) => Entry::Object(metadata, bytes),
157        }
158    }
159
160    /// Returns `true` if the backend contains an entry for the given id.
161    pub fn contains(&self, id: &ObjectId) -> bool {
162        self.store.lock().unwrap().contains_key(id)
163    }
164
165    /// Returns `true` if the backend has no stored objects.
166    pub fn is_empty(&self) -> bool {
167        self.store.lock().unwrap().is_empty()
168    }
169
170    /// Removes an entry directly, bypassing the `Backend` trait.
171    ///
172    /// Useful for simulating partial failures (e.g. orphan tombstones).
173    pub fn remove(&self, id: &ObjectId) {
174        self.store.lock().unwrap().remove(id);
175    }
176}
177
178#[async_trait::async_trait]
179impl super::common::Backend for InMemoryBackend {
180    fn name(&self) -> &'static str {
181        self.name
182    }
183
184    async fn put_object(
185        &self,
186        id: &ObjectId,
187        metadata: &Metadata,
188        stream: ClientStream,
189        _access_time: Timestamp,
190    ) -> Result<PutResponse> {
191        self.change_stream
192            .begin_write(id, metadata.time_expires)
193            .await?;
194        let bytes: BytesMut = stream.try_collect().await?;
195        let entry = StoreEntry::Object(metadata.clone(), bytes.freeze());
196        let size = entry.stored_size();
197        self.store.lock().unwrap().insert(id.clone(), entry);
198        self.change_stream
199            .commit_write(id, size as u64, metadata.time_expires)
200            .await;
201        Ok(())
202    }
203
204    async fn get_object(
205        &self,
206        id: &ObjectId,
207        access_time: Timestamp,
208        range: Option<ByteRange>,
209    ) -> Result<GetResponse> {
210        let entry = self.store.lock().unwrap().get(id).cloned();
211        match entry {
212            None => Ok(None),
213            Some(entry) if entry.is_expired(access_time) => Ok(None),
214            Some(StoreEntry::Tombstone(_)) => Err(ErrorKind::UnexpectedTombstone.into()),
215            Some(StoreEntry::Object(mut metadata, bytes)) => {
216                let total = bytes.len() as u64;
217                metadata.size = Some(bytes.len());
218                let (content_range, payload) = match range {
219                    Some(range) => {
220                        let content_range = range
221                            .resolve(total)
222                            .ok_or(ErrorKind::RangeNotSatisfiable { total })?;
223                        let sliced =
224                            bytes.slice(content_range.start as usize..=content_range.end as usize);
225                        (Some(content_range), sliced)
226                    }
227                    None => (None, bytes),
228                };
229                Ok(Some((
230                    metadata,
231                    content_range,
232                    crate::stream::single(payload),
233                )))
234            }
235        }
236    }
237
238    async fn set_expiry(
239        &self,
240        id: &ObjectId,
241        target: ExpiryUpdate,
242        access_time: Timestamp,
243    ) -> Result<SetExpiryResponse> {
244        self.update_expiry(id, access_time, |entry| match entry {
245            StoreEntry::Object(metadata, _) => extend_object_expiry(metadata, target, access_time),
246            StoreEntry::Tombstone(_) => Ok(ExpiryOutcome::Rejected),
247        })
248        .await
249    }
250
251    async fn delete_object(
252        &self,
253        id: &ObjectId,
254        _access_time: Timestamp,
255    ) -> Result<DeleteResponse> {
256        if self.store.lock().unwrap().remove(id).is_some() {
257            self.change_stream.commit_delete(id).await;
258        }
259        Ok(())
260    }
261
262    async fn create_upload_session(
263        &self,
264        id: &ObjectId,
265        metadata: &Metadata,
266        _upload_length: NonZeroU64,
267    ) -> Result<Option<BackendToken>> {
268        self.change_stream
269            .begin_write(id, metadata.time_expires)
270            .await?;
271        let token = uuid::Uuid::now_v7().to_string();
272        let upload = ResumableUpload {
273            metadata: metadata.clone(),
274            data: BytesMut::new(),
275        };
276        self.resumable_store.lock().unwrap().insert(
277            (id.clone(), token.clone()),
278            Arc::new(tokio::sync::Mutex::new(Some(upload))),
279        );
280        Ok(Some(token))
281    }
282
283    async fn put_chunk(
284        &self,
285        session: &Session,
286        offset: u64,
287        content_length: u64,
288        mut stream: ClientStream,
289    ) -> Result<UploadProgress> {
290        let upload_session = self.upload_session(session)?;
291        let mut guard = upload_session.lock().await;
292        let upload = guard.as_mut().ok_or(ErrorKind::UnknownUploadSession)?;
293        offset
294            .checked_add(content_length)
295            .filter(|end| *end <= session.upload_length.get())
296            .ok_or(ErrorKind::ChunkExceedsUploadLength {
297                offset,
298                content_length,
299                upload_length: session.upload_length.get(),
300            })?;
301        if content_length != 0 && offset != upload.data.len() as u64 {
302            return Err(ErrorKind::UploadOffsetMismatch {
303                offset: upload.data.len() as u64,
304            }
305            .into());
306        }
307
308        let mut remaining = content_length;
309        while remaining > 0 {
310            let Some(chunk) = stream.try_next().await? else {
311                break;
312            };
313            let count = remaining.min(chunk.len() as u64) as usize;
314            upload.data.extend_from_slice(&chunk[..count]);
315            remaining -= count as u64;
316        }
317
318        let offset = upload.data.len() as u64;
319        if offset != session.upload_length.get() {
320            return Ok(UploadProgress::Incomplete { offset });
321        }
322
323        let upload = guard.take().unwrap();
324        let metadata = upload.metadata;
325        let expires_at = metadata.time_expires;
326        let entry = StoreEntry::Object(metadata, upload.data.freeze());
327        let size = entry.stored_size();
328        self.store
329            .lock()
330            .unwrap()
331            .insert(session.object_id.clone(), entry);
332
333        self.change_stream
334            .commit_write(&session.object_id, size as u64, expires_at)
335            .await;
336        self.resumable_store
337            .lock()
338            .unwrap()
339            .remove(&(session.object_id.clone(), session.backend_token.clone()));
340
341        Ok(UploadProgress::Complete)
342    }
343
344    async fn upload_offset(&self, session: &Session) -> Result<UploadProgress> {
345        let upload_session = self.upload_session(session)?;
346        let guard = upload_session.lock().await;
347        let upload = guard.as_ref().ok_or(ErrorKind::UnknownUploadSession)?;
348        Ok(UploadProgress::Incomplete {
349            offset: upload.data.len() as u64,
350        })
351    }
352
353    async fn cancel_upload(&self, session: &Session) -> Result<()> {
354        let upload_session = self.upload_session(session)?;
355        let mut guard = upload_session.lock().await;
356        guard.take().ok_or(ErrorKind::UnknownUploadSession)?;
357        self.resumable_store
358            .lock()
359            .unwrap()
360            .remove(&(session.object_id.clone(), session.backend_token.clone()));
361        Ok(())
362    }
363
364    async fn join(&self) {
365        flush_change_stream(&self.change_stream).await;
366    }
367}
368
369#[async_trait::async_trait]
370impl HighVolumeBackend for InMemoryBackend {
371    async fn create_upload_marker(
372        &self,
373        revision: &ObjectId,
374        time_expires: Timestamp,
375    ) -> Result<()> {
376        self.change_stream
377            .begin_write(revision, Some(time_expires))
378            .await?;
379        self.upload_markers
380            .lock()
381            .unwrap()
382            .insert(revision.clone(), time_expires);
383        self.change_stream
384            .commit_write(revision, 1, Some(time_expires))
385            .await;
386        Ok(())
387    }
388
389    async fn has_upload_marker(&self, revision: &ObjectId, access_time: Timestamp) -> Result<bool> {
390        Ok(self
391            .upload_markers
392            .lock()
393            .unwrap()
394            .get(revision)
395            .is_some_and(|expiry| *expiry >= access_time))
396    }
397
398    async fn delete_upload_marker(
399        &self,
400        revision: &ObjectId,
401        access_time: Timestamp,
402    ) -> Result<bool> {
403        let deleted = self
404            .upload_markers
405            .lock()
406            .unwrap()
407            .remove(revision)
408            .is_some_and(|expiry| expiry >= access_time);
409        if deleted {
410            self.change_stream.commit_delete(revision).await;
411        }
412        Ok(deleted)
413    }
414
415    async fn put_non_tombstone(
416        &self,
417        id: &ObjectId,
418        metadata: &Metadata,
419        payload: Bytes,
420        access_time: Timestamp,
421    ) -> Result<Option<Tombstone>> {
422        self.change_stream
423            .begin_write(id, metadata.time_expires)
424            .await?;
425        let (size, expires_at) = {
426            let mut store = self.store.lock().unwrap();
427            if let Some(StoreEntry::Tombstone(tombstone)) = store.get(id)
428                && !tombstone.is_expired(access_time)
429            {
430                return Ok(Some(tombstone.clone()));
431            }
432
433            let mut metadata = metadata.clone();
434            metadata.size = Some(payload.len());
435            let expires_at = metadata.time_expires;
436            let entry = StoreEntry::Object(metadata, payload);
437            let size = entry.stored_size();
438            store.insert(id.clone(), entry);
439            (size, expires_at)
440        };
441
442        self.change_stream
443            .commit_write(id, size as u64, expires_at)
444            .await;
445        Ok(None)
446    }
447
448    async fn get_tiered_object(
449        &self,
450        id: &ObjectId,
451        access_time: Timestamp,
452        range: Option<ByteRange>,
453    ) -> Result<TieredGet> {
454        let entry = self.store.lock().unwrap().get(id).cloned();
455        Ok(match entry {
456            None => TieredGet::NotFound,
457            Some(entry) if entry.is_expired(access_time) => TieredGet::NotFound,
458            Some(StoreEntry::Tombstone(tombstone)) => TieredGet::Tombstone(tombstone),
459            Some(StoreEntry::Object(mut metadata, bytes)) => {
460                let total = bytes.len() as u64;
461                metadata.size = Some(bytes.len());
462                let (content_range, payload) = match range {
463                    Some(range) => {
464                        let content_range = range
465                            .resolve(total)
466                            .ok_or(ErrorKind::RangeNotSatisfiable { total })?;
467                        let sliced =
468                            bytes.slice(content_range.start as usize..=content_range.end as usize);
469                        (Some(content_range), sliced)
470                    }
471                    None => (None, bytes),
472                };
473                TieredGet::Object(metadata, content_range, crate::stream::single(payload))
474            }
475        })
476    }
477
478    async fn get_tiered_metadata(
479        &self,
480        id: &ObjectId,
481        access_time: Timestamp,
482    ) -> Result<TieredMetadata> {
483        let entry = self.store.lock().unwrap().get(id).cloned();
484        Ok(match entry {
485            None => TieredMetadata::NotFound,
486            Some(entry) if entry.is_expired(access_time) => TieredMetadata::NotFound,
487            Some(StoreEntry::Tombstone(tombstone)) => TieredMetadata::Tombstone(tombstone),
488            Some(StoreEntry::Object(metadata, _bytes)) => TieredMetadata::Object(metadata),
489        })
490    }
491
492    async fn delete_non_tombstone(
493        &self,
494        id: &ObjectId,
495        access_time: Timestamp,
496    ) -> Result<Option<Tombstone>> {
497        let removed = {
498            let mut store = self.store.lock().unwrap();
499            if let Some(StoreEntry::Tombstone(tombstone)) = store.get(id).cloned()
500                && !tombstone.is_expired(access_time)
501            {
502                return Ok(Some(tombstone));
503            }
504
505            store.remove(id).is_some()
506        };
507
508        if removed {
509            self.change_stream.commit_delete(id).await;
510        }
511        Ok(None)
512    }
513
514    async fn compare_and_update(
515        &self,
516        id: &ObjectId,
517        current: Option<&ObjectId>,
518        update: TieredUpdate,
519        access_time: Timestamp,
520    ) -> Result<SetExpiryResponse> {
521        let TieredUpdate::SetExpiry(expiry_target) = update;
522        self.update_expiry(id, access_time, |entry| match (entry, current) {
523            (StoreEntry::Object(metadata, _), None) => {
524                extend_object_expiry(metadata, expiry_target, access_time)
525            }
526            (StoreEntry::Tombstone(t), Some(target)) if t.target == *target => {
527                extend_expiry(&mut t.time_expires, expiry_target, None, access_time)
528            }
529            _ => Ok(ExpiryOutcome::Rejected),
530        })
531        .await
532    }
533
534    async fn compare_and_write(
535        &self,
536        id: &ObjectId,
537        current: Option<&ObjectId>,
538        write: TieredWrite,
539        access_time: Timestamp,
540    ) -> Result<bool> {
541        match &write {
542            TieredWrite::Tombstone(tombstone) => {
543                self.change_stream
544                    .begin_write(id, tombstone.time_expires)
545                    .await?;
546            }
547            TieredWrite::Object(metadata, _) => {
548                self.change_stream
549                    .begin_write(id, metadata.time_expires)
550                    .await?;
551            }
552            TieredWrite::Delete => {}
553        }
554
555        let mut written = None;
556        let mut deleted = false;
557        let matches = {
558            let mut store = self.store.lock().unwrap();
559
560            let actual = store.get(id);
561            let matches_current = matches_redirect(actual, current, access_time);
562            let matches_next = matches_redirect(actual, write.target(), access_time);
563
564            if matches_current {
565                match write {
566                    TieredWrite::Tombstone(tombstone) => {
567                        let expires_at = tombstone.time_expires;
568                        let entry = StoreEntry::Tombstone(tombstone);
569                        written = Some((entry.stored_size(), expires_at));
570                        store.insert(id.clone(), entry);
571                    }
572                    TieredWrite::Object(metadata, payload) => {
573                        let expires_at = metadata.time_expires;
574                        let entry = StoreEntry::Object(metadata, payload);
575                        written = Some((entry.stored_size(), expires_at));
576                        store.insert(id.clone(), entry);
577                    }
578                    TieredWrite::Delete => deleted = store.remove(id).is_some(),
579                }
580            }
581
582            matches_current || matches_next
583        };
584
585        if let Some((size, expires_at)) = written {
586            self.change_stream
587                .commit_write(id, size as u64, expires_at)
588                .await;
589        }
590        if deleted {
591            self.change_stream.commit_delete(id).await;
592        }
593
594        Ok(matches)
595    }
596}
597
598/// Serialized length of `value`, or `0` if it cannot be serialized.
599fn json_len<T: serde::Serialize>(value: &T) -> usize {
600    serde_json::to_string(value).map_or(0, |json| json.len())
601}
602
603/// Returns `true` if `entry` matches the expected tombstone redirect state.
604///
605/// - `expected = None`: matches any non-tombstone (absent or inline object).
606/// - `expected = Some(target)`: matches a tombstone whose redirect target equals `target`.
607fn matches_redirect(
608    entry: Option<&StoreEntry>,
609    expected: Option<&ObjectId>,
610    now: Timestamp,
611) -> bool {
612    match entry {
613        None | Some(StoreEntry::Object(..)) => expected.is_none(),
614        Some(StoreEntry::Tombstone(tombstone)) => match expected {
615            None => tombstone.is_expired(now),
616            Some(target) => tombstone.target == *target && !tombstone.is_expired(now),
617        },
618    }
619}
620
621/// What [`extend_expiry`] did to an entry's deadline.
622///
623/// Callers need to tell a deadline that actually moved apart from one that already
624/// covered the request, because only the former is a change worth reporting to the
625/// [`ChangeStream`].
626#[derive(Clone, Copy, Debug, PartialEq, Eq)]
627enum ExpiryOutcome {
628    /// The entry is absent or expired.
629    NotFound,
630    /// The entry was not eligible or the target could not be resolved.
631    Rejected,
632    /// The stored deadline already covered `expire_at`. Nothing was written.
633    AlreadySatisfied(Timestamp),
634    /// The stored deadline moved out to `expire_at`.
635    Extended(Timestamp),
636}
637
638impl ExpiryOutcome {
639    /// Converts the internal mutation outcome to the backend response.
640    fn response(self) -> SetExpiryResponse {
641        match self {
642            Self::NotFound => SetExpiryResponse::NotFound,
643            Self::Rejected => SetExpiryResponse::Rejected,
644            Self::AlreadySatisfied(deadline) | Self::Extended(deadline) => {
645                SetExpiryResponse::Satisfied(deadline)
646            }
647        }
648    }
649}
650
651/// Extends an object's deadline and TTL together, leaving metadata intact on error.
652fn extend_object_expiry(
653    metadata: &mut Metadata,
654    target: ExpiryUpdate,
655    access_time: Timestamp,
656) -> Result<ExpiryOutcome> {
657    let Some(original_expires) = metadata.time_expires else {
658        return Ok(ExpiryOutcome::Rejected); // entry without a deadline
659    };
660
661    // Only write `updated_expires` back if `extended_expiration_policy` also succeeds
662    let mut updated_expires = Some(original_expires);
663    let outcome = extend_expiry(
664        &mut updated_expires,
665        target,
666        metadata.time_created,
667        access_time,
668    )?;
669
670    if let ExpiryOutcome::Extended(expire_at) = outcome {
671        let updated_policy = common::extended_expiration_policy(
672            metadata.expiration_policy,
673            metadata.time_created,
674            original_expires,
675            expire_at,
676        )?;
677
678        metadata.time_expires = updated_expires;
679        metadata.expiration_policy = updated_policy;
680    }
681
682    Ok(outcome)
683}
684
685/// Resolves `target` and extends an active expiry time where valid.
686///
687/// Eligibility is checked before target resolution.
688fn extend_expiry(
689    field: &mut Option<Timestamp>,
690    target: ExpiryUpdate,
691    time_created: Option<Timestamp>,
692    access_time: Timestamp,
693) -> Result<ExpiryOutcome> {
694    let Some(time_expires) = *field else {
695        return Ok(ExpiryOutcome::Rejected); // entries without a deadline cannot be extended
696    };
697
698    if time_expires < access_time {
699        return Ok(ExpiryOutcome::NotFound); // already expired
700    }
701
702    let Some(expire_at) = target.resolve(time_created, access_time)? else {
703        return Ok(ExpiryOutcome::Rejected);
704    };
705    if time_expires >= expire_at {
706        Ok(ExpiryOutcome::AlreadySatisfied(expire_at))
707    } else {
708        *field = Some(expire_at);
709        Ok(ExpiryOutcome::Extended(expire_at))
710    }
711}
712
713/// Type returned by [`InMemoryBackend::get`] for direct inspection of stored entries.
714#[derive(Clone, Debug)]
715pub enum Entry {
716    /// No entry exists at this key.
717    NotFound,
718    /// A real object with its metadata and payload bytes.
719    Object(Metadata, Bytes),
720    /// A redirect tombstone indicating the real object lives in the long-term backend.
721    Tombstone(Tombstone),
722}
723
724impl Entry {
725    /// Returns `true` if the entry is [`Entry::NotFound`].
726    pub fn is_not_found(&self) -> bool {
727        matches!(self, Entry::NotFound)
728    }
729
730    /// Returns `true` if the entry is [`Entry::Object`].
731    pub fn is_object(&self) -> bool {
732        matches!(self, Entry::Object(_, _))
733    }
734
735    /// Returns `true` if the entry is [`Entry::Tombstone`].
736    pub fn is_tombstone(&self) -> bool {
737        matches!(self, Entry::Tombstone(_))
738    }
739
740    /// Panics unless the entry is [`Entry::NotFound`].
741    pub fn expect_not_found(&self) {
742        match self {
743            Entry::NotFound => (),
744            _ => panic!("expected not found entry, got {self:?}"),
745        }
746    }
747
748    /// Returns the metadata and payload bytes, panicking if the entry is not [`Entry::Object`].
749    pub fn expect_object(&self) -> (Metadata, Bytes) {
750        match self {
751            Entry::Object(metadata, bytes) => (metadata.clone(), bytes.clone()),
752            _ => panic!("expected object entry, got {self:?}"),
753        }
754    }
755
756    /// Returns the tombstone, panicking if the entry is not [`Entry::Tombstone`].
757    pub fn expect_tombstone(&self) -> Tombstone {
758        match self {
759            Entry::Tombstone(tombstone) => tombstone.clone(),
760            _ => panic!("expected tombstone entry, got {self:?}"),
761        }
762    }
763}
764
765#[cfg(test)]
766mod tests {
767    use futures_util::StreamExt;
768    use std::time::Duration;
769
770    #[cfg(feature = "storage-cogs")]
771    use objectstore_inventory_tracker::{OpType, test_utils::DummyProducer};
772    use objectstore_types::metadata::ExpirationPolicy;
773    use objectstore_types::scope::{Scope, Scopes};
774
775    use super::*;
776    use crate::backend::common::{Backend, ExpiryTarget};
777    use crate::id::ObjectContext;
778    use crate::stream;
779
780    fn make_id() -> ObjectId {
781        ObjectId::random(ObjectContext {
782            usecase: "testing".parse().unwrap(),
783            scopes: Scopes::from_iter([Scope::create("testing", "value").unwrap()]),
784        })
785    }
786
787    async fn create_session(backend: &InMemoryBackend, id: &ObjectId, length: u64) -> Session {
788        let backend_token = backend
789            .create_upload_session(id, &Metadata::default(), NonZeroU64::new(length).unwrap())
790            .await
791            .unwrap()
792            .unwrap();
793        Session {
794            object_id: id.clone(),
795            upload_length: NonZeroU64::new(length).unwrap(),
796            backend_token,
797        }
798    }
799
800    #[tokio::test]
801    async fn resumable_upload() {
802        let backend = InMemoryBackend::new("test");
803        let id = make_id();
804        let metadata = Metadata {
805            custom: [("preserved".into(), "yes".into())].into(),
806            ..Default::default()
807        };
808
809        // Upload an object and create a session for the same key.
810        backend
811            .put_object(
812                &id,
813                &Metadata::default(),
814                stream::single("old"),
815                Timestamp::now(),
816            )
817            .await
818            .unwrap();
819        let upload_length = NonZeroU64::new(3).unwrap();
820        let token = backend
821            .create_upload_session(&id, &metadata, upload_length)
822            .await
823            .unwrap()
824            .unwrap();
825        let token = Session {
826            object_id: id.clone(),
827            upload_length,
828            backend_token: token,
829        };
830        assert_eq!(
831            backend.upload_offset(&token).await.unwrap(),
832            UploadProgress::Incomplete { offset: 0 }
833        );
834
835        // Upload a prefix. The session advances, and the old object remains visible.
836        assert_eq!(
837            backend
838                .put_chunk(&token, 0, 1, stream::single("a"))
839                .await
840                .unwrap(),
841            UploadProgress::Incomplete { offset: 1 }
842        );
843        assert_eq!(backend.get(&id).expect_object().1, "old");
844        assert_eq!(
845            backend
846                .put_chunk(&token, 0, 0, stream::single(""))
847                .await
848                .unwrap(),
849            UploadProgress::Incomplete { offset: 1 }
850        );
851        let error = backend
852            .put_chunk(&token, 0, 1, stream::single("a"))
853            .await
854            .unwrap_err();
855        assert_eq!(error.kind(), ErrorKind::UploadOffsetMismatch { offset: 1 });
856
857        // Upload the suffix. Publication replaces the object and deletes the session.
858        assert_eq!(
859            backend
860                .put_chunk(&token, 1, 2, stream::single("bc"))
861                .await
862                .unwrap(),
863            UploadProgress::Complete
864        );
865        let (actual, bytes) = backend.get(&id).expect_object();
866        assert_eq!(bytes, "abc");
867        assert_eq!(actual.custom, metadata.custom);
868        assert_eq!(actual.size, metadata.size);
869        assert_eq!(
870            backend.upload_offset(&token).await.unwrap_err().kind(),
871            ErrorKind::UnknownUploadSession
872        );
873    }
874
875    #[tokio::test]
876    async fn resumable_cancel_and_invalid_sessions() {
877        let backend = InMemoryBackend::new("test");
878        let id = make_id();
879        let token = create_session(&backend, &id, 3).await;
880
881        // The backend rejects a chunk whose declared range exceeds the session length.
882        let error = backend
883            .put_chunk(&token, 3, 1, stream::single("x"))
884            .await
885            .unwrap_err();
886        assert_eq!(
887            error.kind(),
888            ErrorKind::ChunkExceedsUploadLength {
889                offset: 3,
890                content_length: 1,
891                upload_length: 3
892            }
893        );
894
895        // Cancellation discards partial progress and makes the token unknown to the backend.
896        backend
897            .put_chunk(&token, 0, 1, stream::single("a"))
898            .await
899            .unwrap();
900        backend.cancel_upload(&token).await.unwrap();
901        assert!(!backend.contains(&id));
902        assert_eq!(
903            backend.upload_offset(&token).await.unwrap_err().kind(),
904            ErrorKind::UnknownUploadSession
905        );
906    }
907
908    #[tokio::test]
909    async fn failed_resumable_chunk_preserves_partial_progress() {
910        let backend = InMemoryBackend::new("test");
911        let id = make_id();
912        let token = create_session(&backend, &id, 4).await;
913
914        // Persist a prefix, then disconnect after writing one byte of the next chunk.
915        backend
916            .put_chunk(&token, 0, 2, stream::single("ab"))
917            .await
918            .unwrap();
919        let body = futures_util::stream::iter([
920            Ok(Bytes::from_static(b"c")),
921            Err(stream::ClientError::new(std::io::Error::other(
922                "interrupted",
923            ))),
924        ])
925        .boxed();
926        assert_eq!(
927            backend
928                .put_chunk(&token, 2, 2, body)
929                .await
930                .unwrap_err()
931                .kind(),
932            ErrorKind::ClientStream
933        );
934
935        // Resume from the partial byte and verify the complete object.
936        assert_eq!(
937            backend.upload_offset(&token).await.unwrap(),
938            UploadProgress::Incomplete { offset: 3 }
939        );
940        assert_eq!(
941            backend
942                .put_chunk(&token, 3, 1, stream::single("d"))
943                .await
944                .unwrap(),
945            UploadProgress::Complete
946        );
947        assert_eq!(backend.get(&id).expect_object().1, "abcd");
948    }
949
950    #[tokio::test]
951    async fn resumable_serializes_session_operations() {
952        let backend = InMemoryBackend::new("test");
953        let id = make_id();
954        let token = create_session(&backend, &id, 1).await;
955        let (sender, receiver) = tokio::sync::oneshot::channel();
956        let body = futures_util::stream::once(async { Ok(receiver.await.unwrap()) }).boxed();
957
958        // An offset query waits while a chunk holds the session lock.
959        let request = backend.put_chunk(&token, 0, 1, body);
960        tokio::pin!(request);
961        assert!(futures_util::poll!(&mut request).is_pending());
962        let query = backend.upload_offset(&token);
963        tokio::pin!(query);
964        assert!(futures_util::poll!(&mut query).is_pending());
965
966        // Completion wakes the query, which observes that the session is now missing.
967        sender.send(Bytes::from_static(b"x")).unwrap();
968        assert_eq!(request.await.unwrap(), UploadProgress::Complete);
969        assert_eq!(
970            query.await.unwrap_err().kind(),
971            ErrorKind::UnknownUploadSession
972        );
973        assert_eq!(backend.get(&id).expect_object().1, "x");
974    }
975
976    #[cfg(feature = "storage-cogs")]
977    #[tokio::test]
978    async fn resumable_emits_only_publication() {
979        let (backend, producer) = backend_with_change_stream();
980        let id = make_id();
981        let token = create_session(&backend, &id, 2).await;
982
983        // Partial session state is not reported as a stored object.
984        backend
985            .put_chunk(&token, 0, 1, stream::single("a"))
986            .await
987            .unwrap();
988        assert!(producer.records().is_empty());
989
990        // Completion emits exactly one write for the published object.
991        backend
992            .put_chunk(&token, 1, 1, stream::single("b"))
993            .await
994            .unwrap();
995        let records = producer.records();
996        assert_eq!(records.len(), 1);
997        assert_eq!(records[0].op_type, OpType::Write);
998    }
999
1000    #[tokio::test]
1001    async fn set_expiry() {
1002        let access_time = Timestamp::from_unix_secs(1_700_000_000).unwrap();
1003        for policy in [
1004            ExpirationPolicy::TimeToLive(Duration::from_hours(1)),
1005            ExpirationPolicy::TimeToIdle(Duration::from_hours(1)),
1006        ] {
1007            let original_expiry = access_time + Duration::from_hours(1);
1008            let requested = original_expiry + Duration::from_hours(1) + Duration::from_nanos(999);
1009            for target in [
1010                ExpiryTarget::At(requested),
1011                ExpiryTarget::FromCreation(Duration::from_hours(2) + Duration::from_nanos(999)),
1012            ] {
1013                let backend = InMemoryBackend::new("test");
1014                let id = make_id();
1015                let metadata = Metadata {
1016                    expiration_policy: policy,
1017                    time_created: Some(access_time),
1018                    time_expires: Some(original_expiry),
1019                    custom: [("preserved".into(), "yes".into())].into(),
1020                    ..Default::default()
1021                };
1022                backend
1023                    .put_object(&id, &metadata, stream::single("payload"), access_time)
1024                    .await
1025                    .unwrap();
1026
1027                let restricted = ExpiryUpdate {
1028                    max: Some(Duration::ZERO),
1029                    target,
1030                };
1031                assert_eq!(
1032                    backend
1033                        .set_expiry(&id, restricted, access_time)
1034                        .await
1035                        .unwrap_err()
1036                        .kind(),
1037                    ErrorKind::InvalidMetadata,
1038                );
1039                let target = ExpiryUpdate {
1040                    max: Some(Duration::from_hours(100)),
1041                    target,
1042                };
1043                assert_eq!(
1044                    backend.set_expiry(&id, target, access_time).await.unwrap(),
1045                    SetExpiryResponse::Satisfied(requested)
1046                );
1047                assert_eq!(
1048                    backend
1049                        .set_expiry(&id, restricted, access_time)
1050                        .await
1051                        .unwrap_err()
1052                        .kind(),
1053                    ErrorKind::InvalidMetadata,
1054                );
1055                let (updated, payload) = backend.get(&id).expect_object();
1056                let expected_policy = match policy {
1057                    ExpirationPolicy::TimeToLive(_) => {
1058                        ExpirationPolicy::TimeToLive(Duration::from_hours(2))
1059                    }
1060                    other => other,
1061                };
1062                assert_eq!(updated.expiration_policy, expected_policy);
1063                assert_eq!(updated.time_created, metadata.time_created);
1064                assert_eq!(updated.custom, metadata.custom);
1065                assert_eq!(payload, Bytes::from_static(b"payload"));
1066                assert_eq!(updated.time_expires, Some(requested));
1067
1068                assert_eq!(
1069                    backend
1070                        .set_expiry(&id, ExpiryTarget::At(original_expiry).into(), access_time)
1071                        .await
1072                        .unwrap(),
1073                    SetExpiryResponse::Satisfied(original_expiry)
1074                );
1075                assert_eq!(backend.get(&id).expect_object().0, updated);
1076            }
1077        }
1078    }
1079
1080    #[test]
1081    fn extend_object_expiry_preserves_metadata_on_error() {
1082        let access_time = Timestamp::from_unix_secs(1_700_000_000).unwrap();
1083        let original_expiry = access_time + Duration::from_hours(1);
1084        let requested = original_expiry + Duration::from_hours(1);
1085        for (time_created, ttl) in [
1086            (
1087                Some(requested + Duration::from_secs(1)),
1088                Duration::from_hours(1),
1089            ),
1090            (None, Duration::MAX),
1091        ] {
1092            let original = Metadata {
1093                expiration_policy: ExpirationPolicy::TimeToLive(ttl),
1094                time_created,
1095                time_expires: Some(original_expiry),
1096                ..Default::default()
1097            };
1098            let mut metadata = original.clone();
1099            for _ in 0..2 {
1100                assert_eq!(
1101                    extend_object_expiry(
1102                        &mut metadata,
1103                        ExpiryTarget::At(requested).into(),
1104                        access_time,
1105                    )
1106                    .unwrap_err()
1107                    .kind(),
1108                    ErrorKind::CorruptData
1109                );
1110                assert_eq!(metadata, original);
1111            }
1112        }
1113    }
1114
1115    #[tokio::test]
1116    async fn set_expiry_rejected() {
1117        let access_time = Timestamp::from_unix_secs(1_700_000_000).unwrap();
1118        let backend = InMemoryBackend::new("test");
1119        let absent = make_id();
1120        assert_eq!(
1121            backend
1122                .set_expiry(
1123                    &absent,
1124                    ExpiryTarget::At(access_time + Duration::from_hours(1)).into(),
1125                    access_time,
1126                )
1127                .await
1128                .unwrap(),
1129            SetExpiryResponse::NotFound
1130        );
1131
1132        let missing_creation = make_id();
1133        let original_expiry = access_time + Duration::from_hours(1);
1134        backend
1135            .put_object(
1136                &missing_creation,
1137                &Metadata {
1138                    expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_hours(1)),
1139                    time_expires: Some(original_expiry),
1140                    ..Default::default()
1141                },
1142                stream::single("legacy"),
1143                access_time,
1144            )
1145            .await
1146            .unwrap();
1147        let absolute = original_expiry + Duration::from_hours(1);
1148        assert_eq!(
1149            backend
1150                .set_expiry(
1151                    &missing_creation,
1152                    ExpiryTarget::At(absolute).into(),
1153                    access_time
1154                )
1155                .await
1156                .unwrap(),
1157            SetExpiryResponse::Satisfied(absolute)
1158        );
1159        let updated = backend.get(&missing_creation).expect_object().0;
1160        assert_eq!(
1161            updated.expiration_policy,
1162            ExpirationPolicy::TimeToLive(Duration::from_hours(2))
1163        );
1164        assert_eq!(updated.time_created, None);
1165        assert_eq!(
1166            backend
1167                .set_expiry(
1168                    &missing_creation,
1169                    ExpiryTarget::FromCreation(Duration::ZERO).into(),
1170                    access_time,
1171                )
1172                .await
1173                .unwrap(),
1174            SetExpiryResponse::Rejected
1175        );
1176
1177        let manual = make_id();
1178        backend
1179            .put_object(
1180                &manual,
1181                &Metadata::default(),
1182                stream::single("manual"),
1183                access_time,
1184            )
1185            .await
1186            .unwrap();
1187        assert_eq!(
1188            backend
1189                .set_expiry(
1190                    &manual,
1191                    ExpiryTarget::At(access_time + Duration::from_hours(1)).into(),
1192                    access_time,
1193                )
1194                .await
1195                .unwrap(),
1196            SetExpiryResponse::Rejected
1197        );
1198
1199        let expired = make_id();
1200        let metadata = Metadata {
1201            expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_hours(1)),
1202            time_expires: Some(access_time - Duration::from_secs(1)),
1203            ..Default::default()
1204        };
1205        backend
1206            .put_object(&expired, &metadata, stream::single("expired"), access_time)
1207            .await
1208            .unwrap();
1209        let deadline = metadata.time_expires.unwrap();
1210        for time in [deadline - Duration::from_secs(1), deadline] {
1211            assert!(
1212                backend
1213                    .get_object(&expired, time, None)
1214                    .await
1215                    .unwrap()
1216                    .is_some()
1217            );
1218        }
1219        assert!(
1220            backend
1221                .get_object(&expired, access_time, None)
1222                .await
1223                .unwrap()
1224                .is_none()
1225        );
1226        assert_eq!(
1227            backend
1228                .set_expiry(
1229                    &expired,
1230                    ExpiryTarget::At(access_time + Duration::from_hours(1)).into(),
1231                    access_time,
1232                )
1233                .await
1234                .unwrap(),
1235            SetExpiryResponse::NotFound
1236        );
1237    }
1238
1239    #[tokio::test]
1240    async fn redirect_expiry() {
1241        let access_time = Timestamp::from_unix_secs(1_700_000_000).unwrap();
1242        let backend = InMemoryBackend::new("test");
1243        let id = make_id();
1244        let target = make_id();
1245        let other = make_id();
1246        let old_expiry = access_time + Duration::from_hours(1);
1247        backend
1248            .compare_and_write(
1249                &id,
1250                None,
1251                TieredWrite::Tombstone(Tombstone {
1252                    target: target.clone(),
1253                    time_expires: Some(old_expiry),
1254                }),
1255                access_time,
1256            )
1257            .await
1258            .unwrap();
1259
1260        assert!(
1261            backend
1262                .get_object(&id, access_time, None)
1263                .await
1264                .is_err_and(|error| error.kind() == ErrorKind::UnexpectedTombstone)
1265        );
1266        assert!(
1267            backend
1268                .get_metadata(&id, access_time)
1269                .await
1270                .is_err_and(|error| error.kind() == ErrorKind::UnexpectedTombstone)
1271        );
1272
1273        let new_expiry = old_expiry + Duration::from_hours(1);
1274        assert_eq!(
1275            backend
1276                .compare_and_update(
1277                    &id,
1278                    Some(&other),
1279                    TieredUpdate::SetExpiry(ExpiryTarget::At(new_expiry).into()),
1280                    access_time,
1281                )
1282                .await
1283                .unwrap(),
1284            SetExpiryResponse::Rejected
1285        );
1286        assert_eq!(
1287            backend.get(&id).expect_tombstone().time_expires,
1288            Some(old_expiry)
1289        );
1290        assert_eq!(
1291            backend
1292                .compare_and_update(
1293                    &id,
1294                    Some(&target),
1295                    TieredUpdate::SetExpiry(
1296                        ExpiryTarget::FromCreation(Duration::from_hours(2)).into()
1297                    ),
1298                    access_time,
1299                )
1300                .await
1301                .unwrap(),
1302            SetExpiryResponse::Rejected
1303        );
1304        assert_eq!(
1305            backend
1306                .compare_and_update(
1307                    &id,
1308                    Some(&target),
1309                    TieredUpdate::SetExpiry(ExpiryTarget::At(new_expiry).into()),
1310                    access_time
1311                )
1312                .await
1313                .unwrap(),
1314            SetExpiryResponse::Satisfied(new_expiry)
1315        );
1316        assert_eq!(
1317            backend.get(&id).expect_tombstone().time_expires,
1318            Some(new_expiry)
1319        );
1320    }
1321
1322    #[cfg(feature = "storage-cogs")]
1323    fn backend_with_change_stream() -> (InMemoryBackend, DummyProducer) {
1324        use crate::change_stream::CostTrackerStreamConfig;
1325
1326        let (streams, producer) = crate::change_stream::dummy_factory();
1327        let change_stream = streams.build(Some(&CostTrackerStreamConfig {
1328            shared_resource_id: "in_memory_objectstore".into(),
1329            sample_rate: 1.0,
1330        }));
1331        (
1332            InMemoryBackend::new("test").with_change_stream(change_stream),
1333            producer,
1334        )
1335    }
1336
1337    #[cfg(feature = "storage-cogs")]
1338    #[tokio::test]
1339    async fn change_stream_reports_writes_and_deletes() {
1340        let (backend, producer) = backend_with_change_stream();
1341        let id = make_id();
1342        let metadata = Metadata::default();
1343        let payload = b"hello";
1344
1345        backend
1346            .put_object(
1347                &id,
1348                &metadata,
1349                stream::single(payload.to_vec()),
1350                Timestamp::now(),
1351            )
1352            .await
1353            .unwrap();
1354        backend.delete_object(&id, Timestamp::now()).await.unwrap();
1355        // The object is already gone, so this reports nothing.
1356        backend.delete_object(&id, Timestamp::now()).await.unwrap();
1357
1358        let records = producer.records();
1359        assert_eq!(records.len(), 2);
1360        assert_eq!(records[0].op_type, OpType::Write);
1361        assert_eq!(records[0].app_feature, "testing");
1362        assert_eq!(
1363            records[0].size,
1364            Some((json_len(&metadata) + payload.len()) as u64),
1365            "the reported size covers metadata as well as the payload"
1366        );
1367        assert!(json_len(&metadata) > 0, "metadata must contribute bytes");
1368        assert_eq!(records[1].op_type, OpType::Delete);
1369    }
1370
1371    #[cfg(feature = "storage-cogs")]
1372    #[tokio::test]
1373    async fn change_stream_reports_expiry_extension_as_update() {
1374        let (backend, producer) = backend_with_change_stream();
1375        let id = make_id();
1376        let access_time = Timestamp::now();
1377        let expires = access_time + Duration::from_secs(3600);
1378        let metadata = Metadata {
1379            time_expires: Some(expires),
1380            ..Default::default()
1381        };
1382
1383        backend
1384            .put_object(
1385                &id,
1386                &metadata,
1387                stream::single(b"hello".to_vec()),
1388                access_time,
1389            )
1390            .await
1391            .unwrap();
1392        producer.clear();
1393
1394        let extended = expires + Duration::from_secs(3600);
1395        assert_eq!(
1396            backend
1397                .set_expiry(&id, ExpiryTarget::At(extended).into(), access_time)
1398                .await
1399                .unwrap(),
1400            SetExpiryResponse::Satisfied(extended)
1401        );
1402
1403        let records = producer.records();
1404        assert_eq!(records.len(), 1);
1405        assert_eq!(records[0].op_type, OpType::Update);
1406
1407        // A deadline that already covers the request writes nothing, so it reports nothing.
1408        producer.clear();
1409        let requested = expires + Duration::from_secs(60);
1410        assert_eq!(
1411            backend
1412                .set_expiry(&id, ExpiryTarget::At(requested).into(), access_time)
1413                .await
1414                .unwrap(),
1415            SetExpiryResponse::Satisfied(requested)
1416        );
1417        assert!(producer.records().is_empty());
1418    }
1419
1420    #[cfg(feature = "storage-cogs")]
1421    #[tokio::test]
1422    async fn change_stream_reports_tombstone_expiry_extension() {
1423        let (backend, producer) = backend_with_change_stream();
1424        let id = make_id();
1425        let target = make_id();
1426        let access_time = Timestamp::now();
1427        let expires = access_time + Duration::from_secs(3600);
1428
1429        backend
1430            .compare_and_write(
1431                &id,
1432                None,
1433                TieredWrite::Tombstone(Tombstone {
1434                    target: target.clone(),
1435                    time_expires: Some(expires),
1436                }),
1437                access_time,
1438            )
1439            .await
1440            .unwrap();
1441        producer.clear();
1442
1443        let extended = expires + Duration::from_secs(3600);
1444        assert_eq!(
1445            backend
1446                .compare_and_update(
1447                    &id,
1448                    Some(&target),
1449                    TieredUpdate::SetExpiry(ExpiryTarget::At(extended).into()),
1450                    access_time,
1451                )
1452                .await
1453                .unwrap(),
1454            SetExpiryResponse::Satisfied(extended)
1455        );
1456
1457        let records = producer.records();
1458        assert_eq!(records.len(), 1);
1459        assert_eq!(records[0].op_type, OpType::Update);
1460    }
1461}
objectstore_service/backend/in_memory.rs

objectstore_service/backend/
in_memory.rs
