feat(spooler): Implement backpressure mechanism via notify

relay-server/src/service.rs

@@ -1,6 +1,5 @@
 use std::convert::Infallible;
 use std::fmt;
-use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -30,7 +29,7 @@ use crate::services::relays::{RelayCache, RelayCacheService};
 use crate::services::store::StoreService;
 use crate::services::test_store::{TestStore, TestStoreService};
 use crate::services::upstream::{UpstreamRelay, UpstreamRelayService};
-use crate::utils::{MemoryChecker, MemoryStat};
+use crate::utils::{MemoryChecker, MemoryStat, Waiter};
 
 /// Indicates the type of failure of the server.
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, thiserror::Error)]
@@ -244,7 +243,7 @@ impl ServiceState {
 
         // We initialize a shared boolean that is used to manage backpressure between the
         // EnvelopeBufferService and the ProjectCacheService.
-        let project_cache_ready = Arc::new(AtomicBool::new(true));
+        let project_cache_ready = Waiter::new(false);
         let envelope_buffer = EnvelopeBufferService::new(
             config.clone(),
             MemoryChecker::new(memory_stat.clone(), config.clone()),

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

@@ -26,8 +26,8 @@ use crate::services::project_cache::ProjectCache;
 use crate::services::project_cache::UpdateProject;
 use crate::services::test_store::TestStore;
 use crate::statsd::RelayCounters;
-use crate::utils::ManagedEnvelope;
 use crate::utils::MemoryChecker;
+use crate::utils::{ManagedEnvelope, Waiter};
 
 pub use envelope_buffer::EnvelopeBufferError;
 // pub for benchmarks
@@ -110,7 +110,7 @@ pub struct EnvelopeBufferService {
     services: Services,
     has_capacity: Arc<AtomicBool>,
     sleep: Duration,
-    project_cache_ready: Arc<AtomicBool>,
+    project_cache_ready: Waiter,
 }
 
 /// The maximum amount of time between evaluations of dequeue conditions.
@@ -129,7 +129,7 @@ impl EnvelopeBufferService {
         memory_checker: MemoryChecker,
         global_config_rx: watch::Receiver<global_config::Status>,
         services: Services,
-        project_cache_ready: Arc<AtomicBool>,
+        project_cache_ready: Waiter,
     ) -> Option<Self> {
         config.spool_v2().then(|| Self {
             config,
@@ -172,17 +172,14 @@ impl EnvelopeBufferService {
             tokio::time::sleep(self.sleep).await;
         }
 
-        // In case the project cache is not ready, we defer popping to first try and handle incoming
-        // messages and only come back to this in case within the timeout no data was received.
-        while !self.project_cache_ready.load(Ordering::Relaxed) {
-            tokio::time::sleep(Duration::from_millis(10)).await;
-        }
-
         relay_statsd::metric!(
             counter(RelayCounters::BufferReadyToPop) += 1,
             status = "slept"
         );
 
+        // We try to wait for the project cache.
+        self.project_cache_ready.try_wait().await;
+
         relay_statsd::metric!(
             counter(RelayCounters::BufferReadyToPop) += 1,
             status = "checked"
@@ -245,7 +242,7 @@ impl EnvelopeBufferService {
                     .expect("Element disappeared despite exclusive excess");
                 // We assume that the project cache is now busy to process this envelope, so we flip
                 // the boolean flag, which will prioritize writes.
-                self.project_cache_ready.store(false, Ordering::SeqCst);
+                self.project_cache_ready.set_wait(true);
                 self.services.project_cache.send(DequeuedEnvelope(envelope));
 
                 self.sleep = Duration::ZERO; // try next pop immediately
@@ -452,7 +449,7 @@ mod tests {
         watch::Sender<global_config::Status>,
         mpsc::UnboundedReceiver<ProjectCache>,
         mpsc::UnboundedReceiver<TrackOutcome>,
-        Arc<AtomicBool>,
+        Waiter,
     ) {
         let config = Arc::new(
             Config::from_json_value(serde_json::json!({
@@ -468,7 +465,7 @@ mod tests {
         let (global_tx, global_rx) = watch::channel(global_config::Status::Pending);
         let (project_cache, project_cache_rx) = Addr::custom();
         let (outcome_aggregator, outcome_aggregator_rx) = Addr::custom();
-        let project_cache_ready = Arc::new(AtomicBool::new(true));
+        let project_cache_ready = Waiter::new(false);
         (
             EnvelopeBufferService::new(
                 config,
@@ -563,7 +560,7 @@ mod tests {
         )));
 
         let (project_cache, project_cache_rx) = Addr::custom();
-        let project_cache_ready = Arc::new(AtomicBool::new(true));
+        let project_cache_ready = Waiter::new(false);
         let service = EnvelopeBufferService::new(
             config,
             memory_checker,
@@ -609,7 +606,7 @@ mod tests {
         let (global_tx, global_rx) = watch::channel(global_config::Status::Pending);
         let (project_cache, project_cache_rx) = Addr::custom();
         let (outcome_aggregator, mut outcome_aggregator_rx) = Addr::custom();
-        let project_cache_ready = Arc::new(AtomicBool::new(true));
+        let project_cache_ready = Waiter::new(false);
         let service = EnvelopeBufferService::new(
             config,
             memory_checker,

relay-server/src/services/project_cache.rs

@@ -1,6 +1,5 @@
 use std::collections::{BTreeMap, BTreeSet};
 use std::error::Error;
-use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -41,7 +40,9 @@ use crate::services::test_store::TestStore;
 use crate::services::upstream::UpstreamRelay;
 
 use crate::statsd::{RelayCounters, RelayGauges, RelayHistograms, RelayTimers};
-use crate::utils::{GarbageDisposal, ManagedEnvelope, MemoryChecker, RetryBackoff, SleepHandle};
+use crate::utils::{
+    GarbageDisposal, ManagedEnvelope, MemoryChecker, RetryBackoff, SleepHandle, Waiter,
+};
 
 /// Requests a refresh of a project state from one of the available sources.
 ///
@@ -611,8 +612,8 @@ struct ProjectCacheBroker {
     spool_v1: Option<SpoolV1>,
     /// Status of the global configuration, used to determine readiness for processing.
     global_config: GlobalConfigStatus,
-    /// Atomic boolean signaling whether the project cache is ready to accept a new envelope.
-    project_cache_ready: Arc<AtomicBool>,
+    /// Waiter signaling whether the project cache is ready to accept a new envelope.
+    project_cache_ready: Waiter,
 }
 
 #[derive(Debug)]
@@ -1323,7 +1324,7 @@ impl ProjectCacheBroker {
                         }
 
                         // We mark the project cache as ready to accept new traffic.
-                        self.project_cache_ready.store(true, Ordering::SeqCst);
+                        self.project_cache_ready.set_wait(false);
                     }
                     ProjectCache::UpdateProject(project) => self.handle_update_project(project),
                 }
@@ -1340,7 +1341,7 @@ pub struct ProjectCacheService {
     services: Services,
     global_config_rx: watch::Receiver<global_config::Status>,
     redis: Option<RedisPool>,
-    project_cache_ready: Arc<AtomicBool>,
+    project_cache_ready: Waiter,
 }
 
 impl ProjectCacheService {
@@ -1351,7 +1352,7 @@ impl ProjectCacheService {
         services: Services,
         global_config_rx: watch::Receiver<global_config::Status>,
         redis: Option<RedisPool>,
-        project_cache_ready: Arc<AtomicBool>,
+        project_cache_ready: Waiter,
     ) -> Self {
         Self {
             config,
@@ -1651,7 +1652,7 @@ mod tests {
                     buffer_unspool_backoff: RetryBackoff::new(Duration::from_millis(100)),
                 }),
                 global_config: GlobalConfigStatus::Pending,
-                project_cache_ready: Arc::new(AtomicBool::new(true)),
+                project_cache_ready: Waiter::new(false),
             },
             buffer,
         )

relay-server/src/utils/mod.rs

@@ -19,6 +19,7 @@ mod native;
 mod serde;
 #[cfg(feature = "processing")]
 mod unreal;
+mod waiter;
 
 pub use self::api::*;
 pub use self::dynamic_sampling::*;
@@ -40,3 +41,4 @@ pub use self::statsd::*;
 pub use self::thread_pool::*;
 #[cfg(feature = "processing")]
 pub use self::unreal::*;
+pub use self::waiter::*;

relay-server/src/utils/waiter.rs

@@ -0,0 +1,45 @@
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use tokio::sync::Notify;
+
+/// Struct holding both the boolean state and the notifier.
+#[derive(Debug)]
+struct Inner {
+    wait: AtomicBool,
+    notify: Notify,
+}
+
+/// Construct used to wait on a boolean variable only if it's set to `true`.
+#[derive(Debug, Clone)]
+pub struct Waiter {
+    inner: Arc<Inner>,
+}
+
+impl Waiter {
+    pub fn new(wait: bool) -> Self {
+        Self {
+            inner: Arc::new(Inner {
+                wait: AtomicBool::new(wait),
+                notify: Notify::new(),
+            }),
+        }
+    }
+
+    pub fn set_wait(&self, wait: bool) {
+        self.inner.wait.store(wait, Ordering::SeqCst);
+
+        // If we are not waiting anymore, notify all the subscribed futures
+        if !wait {
+            self.inner.notify.notify_waiters();
+        }
+    }
+
+    pub async fn try_wait(&self) {
+        if !self.inner.wait.load(Ordering::Relaxed) {
+            return;
+        }
+
+        // Suspend on waiting for a notification if waiting is required
+        self.inner.notify.notified().await;
+    }
+}