feat(spooler): Add CachingEnvelopeStack

relay-server/src/services/buffer/envelope_stack/caching.rs

@@ -0,0 +1,101 @@
+use chrono::{DateTime, Utc};
+
+use super::EnvelopeStack;
+use crate::envelope::Envelope;
+
+/// An envelope stack implementation that caches one element in memory and delegates
+/// to another envelope stack for additional storage.
+#[derive(Debug)]
+pub struct CachingEnvelopeStack<S> {
+    /// The underlying envelope stack
+    inner: S,
+    /// The cached envelope (if any)
+    cached: Option<Box<Envelope>>,
+}
+
+impl<S> CachingEnvelopeStack<S>
+where
+    S: EnvelopeStack,
+{
+    /// Creates a new [`CachingEnvelopeStack`] wrapping the provided envelope stack
+    pub fn new(inner: S) -> Self {
+        Self {
+            inner,
+            cached: None,
+        }
+    }
+}
+
+impl<S> EnvelopeStack for CachingEnvelopeStack<S>
+where
+    S: EnvelopeStack,
+{
+    type Error = S::Error;
+
+    async fn push(&mut self, envelope: Box<Envelope>) -> Result<(), Self::Error> {
+        if self.cached.is_none() {
+            // If we don't have a cached envelope, store this one in cache
+            self.cached = Some(envelope);
+            Ok(())
+        } else {
+            // If we already have a cached envelope, push the current cached one
+            // to the inner stack and keep the new one in cache
+            let cached = self.cached.take().unwrap();
+            self.inner.push(cached).await?;
+            self.cached = Some(envelope);
+            Ok(())
+        }
+    }
+
+    async fn peek(&mut self) -> Result<Option<DateTime<Utc>>, Self::Error> {
+        if let Some(ref envelope) = self.cached {
+            Ok(Some(envelope.received_at()))
+        } else {
+            self.inner.peek().await
+        }
+    }
+
+    async fn pop(&mut self) -> Result<Option<Box<Envelope>>, Self::Error> {
+        if let Some(envelope) = self.cached.take() {
+            Ok(Some(envelope))
+        } else {
+            self.inner.pop().await
+        }
+    }
+
+    async fn flush(mut self) {
+        if let Some(envelope) = self.cached {
+            let _ = self.inner.push(envelope).await;
+        }
+        self.inner.flush().await;
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::services::buffer::envelope_stack::memory::MemoryEnvelopeStack;
+    use crate::services::buffer::testutils::utils::mock_envelope;
+
+    #[tokio::test]
+    async fn test_caching_stack() {
+        let inner = MemoryEnvelopeStack::new();
+        let mut stack = CachingEnvelopeStack::new(inner);
+
+        // Create test envelopes with different timestamps
+        let env1 = mock_envelope(Utc::now());
+        let env2 = mock_envelope(Utc::now());
+        let env3 = mock_envelope(Utc::now());
+
+        // Push envelopes
+        stack.push(env1).await.unwrap();
+        stack.push(env2).await.unwrap();
+        stack.push(env3).await.unwrap();
+
+        // Pop them back
+        assert!(stack.pop().await.unwrap().is_some()); // Should come from cache
+        assert!(stack.pop().await.unwrap().is_some()); // Should come from inner
+        assert!(stack.pop().await.unwrap().is_some()); // Should come from inner
+        assert!(stack.pop().await.unwrap().is_none()); // Should be empty
+    }
+}

relay-server/src/services/buffer/envelope_stack/memory.rs

@@ -24,7 +24,7 @@ impl EnvelopeStack for MemoryEnvelopeStack {
     }
 
     async fn peek(&mut self) -> Result<Option<DateTime<Utc>>, Self::Error> {
-        Ok(self.0.last().map(|e| e.meta().received_at()))
+        Ok(self.0.last().map(|e| e.received_at()))
     }
 
     async fn pop(&mut self) -> Result<Option<Box<Envelope>>, Self::Error> {

relay-server/src/services/buffer/envelope_stack/mod.rs

@@ -4,6 +4,7 @@ use chrono::{DateTime, Utc};
 
 use crate::envelope::Envelope;
 
+pub mod caching;
 pub mod memory;
 pub mod sqlite;
 

relay-server/src/services/buffer/stack_provider/sqlite.rs

@@ -3,6 +3,7 @@ use std::error::Error;
 use relay_config::Config;
 
 use crate::services::buffer::common::ProjectKeyPair;
+use crate::services::buffer::envelope_stack::caching::CachingEnvelopeStack;
 use crate::services::buffer::envelope_store::sqlite::{
     SqliteEnvelopeStore, SqliteEnvelopeStoreError,
 };
@@ -38,7 +39,7 @@ impl SqliteStackProvider {
 }
 
 impl StackProvider for SqliteStackProvider {
-    type Stack = SqliteEnvelopeStack;
+    type Stack = CachingEnvelopeStack<SqliteEnvelopeStack>;
 
     async fn initialize(&self) -> InitializationState {
         match self.envelope_store.project_key_pairs().await {
@@ -58,7 +59,7 @@ impl StackProvider for SqliteStackProvider {
         stack_creation_type: StackCreationType,
         project_key_pair: ProjectKeyPair,
     ) -> Self::Stack {
-        SqliteEnvelopeStack::new(
+        let inner = SqliteEnvelopeStack::new(
             self.envelope_store.clone(),
             self.batch_size_bytes,
             project_key_pair.own_key,
@@ -69,7 +70,9 @@ impl StackProvider for SqliteStackProvider {
             // it was empty, or we never had data on disk for that stack, so we assume by default
             // that there is no need to check disk until some data is spooled.
             Self::assume_data_on_disk(stack_creation_type),
-        )
+        );
+
+        CachingEnvelopeStack::new(inner)
     }
 
     fn has_store_capacity(&self) -> bool {