feat(spooler): Add capacity check for the envelope buffer

CHANGELOG.md

@@ -24,6 +24,7 @@
 - Extract client sdk from transaction into profiles. ([#3915](https://github.com/getsentry/relay/pull/3915))
 - Extract `user.geo.subregion` into span metrics/indexed. ([#3914](https://github.com/getsentry/relay/pull/3914))
 - Add `last_peek` field to the `Priority` struct. ([#3922](https://github.com/getsentry/relay/pull/3922))
+- Add `EnvelopeStore` trait and implement `DiskUsage` for tracking disk usage. ([#3925](https://github.com/getsentry/relay/pull/3925))
 
 ## 24.7.1
 

relay-config/src/config.rs

@@ -879,7 +879,7 @@ pub struct EnvelopeSpool {
     min_connections: u32,
     /// The maximum size of the buffer to keep, in bytes.
     ///
-    /// If not set the befault is 524288000 bytes (500MB).
+    /// If not set the default is 524288000 bytes (500MB).
     #[serde(default = "spool_envelopes_max_disk_size")]
     max_disk_size: ByteSize,
     /// The maximum bytes to keep in the memory buffer before spooling envelopes to disk, in bytes.

relay-server/benches/benches.rs

@@ -4,13 +4,15 @@ use relay_config::Config;
 use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions};
 use sqlx::{Pool, Sqlite};
 use std::path::PathBuf;
+use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tempfile::TempDir;
 use tokio::runtime::Runtime;
 
 use relay_base_schema::project::ProjectKey;
 use relay_server::{
-    Envelope, EnvelopeStack, PolymorphicEnvelopeBuffer, SqliteEnvelopeStack, SqliteEnvelopeStore,
+    Envelope, EnvelopeStack, MemoryChecker, MemoryStat, PolymorphicEnvelopeBuffer,
+    SqliteEnvelopeStack, SqliteEnvelopeStore,
 };
 
 fn setup_db(path: &PathBuf) -> Pool<Sqlite> {
@@ -72,7 +74,7 @@ fn benchmark_sqlite_envelope_stack(c: &mut Criterion) {
     let temp_dir = TempDir::new().unwrap();
     let db_path = temp_dir.path().join("test.db");
     let db = setup_db(&db_path);
-    let envelope_store = SqliteEnvelopeStore::new(db.clone());
+    let envelope_store = SqliteEnvelopeStore::new(db.clone(), 100);
 
     let runtime = Runtime::new().unwrap();
 
@@ -221,6 +223,17 @@ fn benchmark_envelope_buffer(c: &mut Criterion) {
     let num_projects = 100000;
     let envelopes_per_project = 10;
 
+    let config: Arc<Config> = Config::from_json_value(serde_json::json!({
+        "spool": {
+            "health": {
+                "max_memory_percent": 1.0
+            }
+        }
+    }))
+    .unwrap()
+    .into();
+    let memory_checker = MemoryChecker::new(MemoryStat::default(), config.clone());
+
     group.throughput(Throughput::Elements(
         num_projects * envelopes_per_project as u64,
     ));
@@ -245,7 +258,8 @@ fn benchmark_envelope_buffer(c: &mut Criterion) {
             },
             |envelopes| {
                 runtime.block_on(async {
-                    let mut buffer = PolymorphicEnvelopeBuffer::from_config(&Config::default());
+                    let mut buffer =
+                        PolymorphicEnvelopeBuffer::from_config(&config, memory_checker.clone());
                     for envelope in envelopes.into_iter() {
                         buffer.push(envelope).await.unwrap();
                     }
@@ -274,7 +288,8 @@ fn benchmark_envelope_buffer(c: &mut Criterion) {
             },
             |envelopes| {
                 runtime.block_on(async {
-                    let mut buffer = PolymorphicEnvelopeBuffer::from_config(&Config::default());
+                    let mut buffer =
+                        PolymorphicEnvelopeBuffer::from_config(&config, memory_checker.clone());
                     let n = envelopes.len();
                     for envelope in envelopes.into_iter() {
                         let public_key = envelope.meta().public_key();

relay-server/src/endpoints/common.rs

@@ -259,7 +259,7 @@ pub fn event_id_from_items(items: &Items) -> Result<Option<EventId>, BadStoreReq
 ///
 /// Queueing can fail if the queue exceeds `envelope_buffer_size`. In this case, `Err` is
 /// returned and the envelope is not queued.
-fn queue_envelope(
+async fn queue_envelope(
     state: &ServiceState,
     mut managed_envelope: ManagedEnvelope,
 ) -> Result<(), BadStoreRequest> {
@@ -307,12 +307,14 @@ fn queue_envelope(
 
         match state.envelope_buffer() {
             Some(buffer) => {
+                if !buffer.has_capacity().await {
+                    return Err(BadStoreRequest::QueueFailed);
+                }
+
                 // NOTE: This assumes that a `prefetch` has already been scheduled for both the
                 // envelope's projects. See `handle_check_envelope`.
                 relay_log::trace!("Pushing envelope to V2 buffer");
 
-                // TODO: Sync-check whether the buffer has capacity.
-                // Otherwise return `QueueFailed`.
                 buffer.defer_push(envelope);
             }
             None => {
@@ -347,13 +349,6 @@ pub async fn handle_envelope(
         )
     }
 
-    // TODO(jjbayer): Move this check to spool impl
-    if state.memory_checker().check_memory().is_exceeded() {
-        // NOTE: Long-term, we should not reject the envelope here, but spool it to disk instead.
-        // This will be fixed with the new spool implementation.
-        return Err(BadStoreRequest::QueueFailed);
-    };
-
     let mut managed_envelope = ManagedEnvelope::new(
         envelope,
         state.outcome_aggregator().clone(),
@@ -393,7 +388,7 @@ pub async fn handle_envelope(
         return Err(BadStoreRequest::Overflow(offender));
     }
 
-    queue_envelope(state, managed_envelope)?;
+    queue_envelope(state, managed_envelope).await?;
 
     if checked.rate_limits.is_limited() {
         // Even if some envelope items have been queued, there might be active rate limits on

relay-server/src/lib.rs

@@ -271,6 +271,7 @@ pub use self::services::buffer::{
     EnvelopeStack, PolymorphicEnvelopeBuffer, SqliteEnvelopeStack, SqliteEnvelopeStore,
 }; // pub for benchmarks
 pub use self::services::spooler::spool_utils;
+pub use self::utils::{MemoryChecker, MemoryStat}; // pub for benchmarks
 
 #[cfg(test)]
 mod testutils;

relay-server/src/service.rs

@@ -254,7 +254,11 @@ impl ServiceState {
             upstream_relay.clone(),
             global_config.clone(),
         );
-        let envelope_buffer = GuardedEnvelopeBuffer::from_config(&config).map(Arc::new);
+        let envelope_buffer = GuardedEnvelopeBuffer::from_config(
+            &config,
+            MemoryChecker::new(memory_stat.clone(), config.clone()),
+        )
+        .map(Arc::new);
         ProjectCacheService::new(
             config.clone(),
             MemoryChecker::new(memory_stat.clone(), config.clone()),

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

@@ -8,11 +8,13 @@ use relay_config::Config;
 
 use crate::envelope::Envelope;
 use crate::services::buffer::envelope_stack::sqlite::SqliteEnvelopeStackError;
-use crate::services::buffer::envelope_stack::{EnvelopeStack, StackProvider};
-use crate::services::buffer::sqlite_envelope_store::SqliteEnvelopeStoreError;
-use crate::services::buffer::stack_provider::memory::MemoryStackProvider;
-use crate::services::buffer::stack_provider::sqlite::SqliteStackProvider;
+use crate::services::buffer::envelope_stack::EnvelopeStack;
+use crate::services::buffer::envelope_store::sqlite::SqliteEnvelopeStoreError;
+use crate::services::buffer::stacks_manager::memory::MemoryStacksManager;
+use crate::services::buffer::stacks_manager::sqlite::SqliteStacksManager;
+use crate::services::buffer::stacks_manager::{Capacity, StacksManager};
 use crate::statsd::{RelayCounters, RelayGauges};
+use crate::utils::MemoryChecker;
 
 /// Polymorphic envelope buffering interface.
 ///
@@ -25,20 +27,21 @@ use crate::statsd::{RelayCounters, RelayGauges};
 #[allow(private_interfaces)]
 pub enum PolymorphicEnvelopeBuffer {
     /// An enveloper buffer that uses in-memory envelopes stacks.
-    InMemory(EnvelopeBuffer<MemoryStackProvider>),
+    InMemory(EnvelopeBuffer<MemoryStacksManager>),
     /// An enveloper buffer that uses sqlite envelopes stacks.
     #[allow(dead_code)]
-    Sqlite(EnvelopeBuffer<SqliteStackProvider>),
+    Sqlite(EnvelopeBuffer<SqliteStacksManager>),
 }
 
 impl PolymorphicEnvelopeBuffer {
     /// Creates either a memory-based or a disk-based envelope buffer,
     /// depending on the given configuration.
-    pub fn from_config(config: &Config) -> Self {
+    pub fn from_config(config: &Config, memory_checker: MemoryChecker) -> Self {
         if config.spool_envelopes_path().is_some() {
             panic!("Disk backend not yet supported for spool V2");
         }
-        Self::InMemory(EnvelopeBuffer::<MemoryStackProvider>::new())
+
+        Self::InMemory(EnvelopeBuffer::<MemoryStacksManager>::new(memory_checker))
     }
 
     /// Adds an envelope to the buffer.
@@ -71,13 +74,21 @@ impl PolymorphicEnvelopeBuffer {
 
     /// Marks a project as ready or not ready.
     ///
-    /// The buffer reprioritizes its envelopes based on this information.
+    /// The buffer re-prioritizes its envelopes based on this information.
     pub fn mark_ready(&mut self, project: &ProjectKey, is_ready: bool) -> bool {
         match self {
             Self::Sqlite(buffer) => buffer.mark_ready(project, is_ready),
             Self::InMemory(buffer) => buffer.mark_ready(project, is_ready),
         }
     }
+
+    /// Returns `true` whether the buffer has capacity to accept new [`Envelope`]s.
+    pub fn has_capacity(&self) -> bool {
+        match self {
+            Self::Sqlite(buffer) => buffer.has_capacity(),
+            Self::InMemory(buffer) => buffer.has_capacity(),
+        }
+    }
 }
 
 /// Error that occurs while interacting with the envelope buffer.
@@ -86,6 +97,9 @@ pub enum EnvelopeBufferError {
     #[error("sqlite")]
     Sqlite(#[from] SqliteEnvelopeStackError),
 
+    #[error("failed to push envelope to the buffer")]
+    PushFailed,
+
     #[error("impossible")]
     Impossible(#[from] Infallible),
 }
@@ -95,7 +109,7 @@ pub enum EnvelopeBufferError {
 /// Envelope stacks are organized in a priority queue, and are reprioritized every time an envelope
 /// is pushed, popped, or when a project becomes ready.
 #[derive(Debug)]
-struct EnvelopeBuffer<P: StackProvider> {
+struct EnvelopeBuffer<P: StacksManager> {
     /// The central priority queue.
     priority_queue: priority_queue::PriorityQueue<QueueItem<StackKey, P::Stack>, Priority>,
     /// A lookup table to find all stacks involving a project.
@@ -107,32 +121,32 @@ struct EnvelopeBuffer<P: StackProvider> {
     stack_provider: P,
 }
 
-impl EnvelopeBuffer<MemoryStackProvider> {
-    /// Creates an empty buffer.
-    pub fn new() -> Self {
+impl EnvelopeBuffer<MemoryStacksManager> {
+    /// Creates an empty memory-based buffer.
+    pub fn new(memory_checker: MemoryChecker) -> Self {
         Self {
             stacks_by_project: Default::default(),
             priority_queue: Default::default(),
-            stack_provider: MemoryStackProvider,
+            stack_provider: MemoryStacksManager::new(memory_checker),
         }
     }
 }
 
 #[allow(dead_code)]
-impl EnvelopeBuffer<SqliteStackProvider> {
-    /// Creates an empty buffer.
+impl EnvelopeBuffer<SqliteStacksManager> {
+    /// Creates an empty sqlite-based buffer.
     pub async fn new(config: &Config) -> Result<Self, SqliteEnvelopeStoreError> {
         Ok(Self {
             stacks_by_project: Default::default(),
             priority_queue: Default::default(),
-            stack_provider: SqliteStackProvider::new(config).await?,
+            stack_provider: SqliteStacksManager::new(config).await?,
         })
     }
 }
 
-impl<P: StackProvider> EnvelopeBuffer<P>
+impl<P: StacksManager> EnvelopeBuffer<P>
 where
-    EnvelopeBufferError: std::convert::From<<P::Stack as EnvelopeStack>::Error>,
+    EnvelopeBufferError: From<<P::Stack as EnvelopeStack>::Error>,
 {
     /// Pushes an envelope to the appropriate envelope stack and reprioritizes the stack.
     ///
@@ -221,7 +235,7 @@ where
         Ok(Some(envelope))
     }
 
-    /// Reprioritizes all stacks that involve the given project key by setting it to "ready".
+    /// Re-prioritizes all stacks that involve the given project key by setting it to "ready".
     pub fn mark_ready(&mut self, project: &ProjectKey, is_ready: bool) -> bool {
         let mut changed = false;
         if let Some(stack_keys) = self.stacks_by_project.get(project) {
@@ -272,6 +286,10 @@ where
         );
     }
 
+    pub fn has_capacity(&self) -> bool {
+        matches!(self.stack_provider.capacity(), Capacity::Free)
+    }
+
     fn pop_stack(&mut self, stack_key: StackKey) {
         for project_key in stack_key.iter() {
             self.stacks_by_project
@@ -420,14 +438,16 @@ impl Readiness {
 #[cfg(test)]
 mod tests {
     use std::str::FromStr;
-    use uuid::Uuid;
+    use std::sync::Arc;
 
     use relay_common::Dsn;
     use relay_event_schema::protocol::EventId;
     use relay_sampling::DynamicSamplingContext;
+    use uuid::Uuid;
 
     use crate::envelope::{Item, ItemType};
     use crate::extractors::RequestMeta;
+    use crate::utils::MemoryStat;
 
     use super::*;
 
@@ -461,9 +481,23 @@ mod tests {
         envelope
     }
 
+    fn mock_memory_checker() -> MemoryChecker {
+        let config: Arc<_> = Config::from_json_value(serde_json::json!({
+            "spool": {
+                "health": {
+                    "max_memory_percent": 1.0
+                }
+            }
+        }))
+        .unwrap()
+        .into();
+
+        MemoryChecker::new(MemoryStat::default(), config.clone())
+    }
+
     #[tokio::test]
     async fn test_insert_pop() {
-        let mut buffer = EnvelopeBuffer::<MemoryStackProvider>::new();
+        let mut buffer = EnvelopeBuffer::<MemoryStacksManager>::new(mock_memory_checker());
 
         let project_key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fed").unwrap();
         let project_key2 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fee").unwrap();
@@ -547,7 +581,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_project_internal_order() {
-        let mut buffer = EnvelopeBuffer::<MemoryStackProvider>::new();
+        let mut buffer = EnvelopeBuffer::<MemoryStacksManager>::new(mock_memory_checker());
 
         let project_key = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fed").unwrap();
 
@@ -574,7 +608,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_sampling_projects() {
-        let mut buffer = EnvelopeBuffer::<MemoryStackProvider>::new();
+        let mut buffer = EnvelopeBuffer::<MemoryStacksManager>::new(mock_memory_checker());
 
         let project_key1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fed").unwrap();
         let project_key2 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fef").unwrap();
@@ -652,7 +686,7 @@ mod tests {
 
         assert_ne!(stack_key1, stack_key2);
 
-        let mut buffer = EnvelopeBuffer::<MemoryStackProvider>::new();
+        let mut buffer = EnvelopeBuffer::<MemoryStacksManager>::new(mock_memory_checker());
         buffer
             .push(new_envelope(project_key1, Some(project_key2), None))
             .await
@@ -666,7 +700,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_last_peek_internal_order() {
-        let mut buffer = EnvelopeBuffer::<MemoryStackProvider>::new();
+        let mut buffer = EnvelopeBuffer::<MemoryStacksManager>::new(mock_memory_checker());
 
         let project_key_1 = ProjectKey::parse("a94ae32be2584e0bbd7a4cbb95971fed").unwrap();
         let event_id_1 = EventId::new();

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

@@ -20,9 +20,3 @@ pub trait EnvelopeStack: Send + std::fmt::Debug {
     /// Pops the [`Envelope`] on top of the stack.
     fn pop(&mut self) -> impl Future<Output = Result<Option<Box<Envelope>>, Self::Error>>;
 }
-
-pub trait StackProvider: std::fmt::Debug {
-    type Stack: EnvelopeStack;
-
-    fn create_stack(&self, envelope: Box<Envelope>) -> Self::Stack;
-}