feat(spooler): Implement backpressure mechanism based on atomics

relay-server/src/service.rs

@@ -1,5 +1,6 @@
 use std::convert::Infallible;
 use std::fmt;
+use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -241,6 +242,9 @@ impl ServiceState {
  )
  .spawn_handler(processor_rx);
 
+ // 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 envelope_buffer = EnvelopeBufferService::new(
  config.clone(),
  MemoryChecker::new(memory_stat.clone(), config.clone()),
@@ -250,6 +254,7 @@ impl ServiceState {
  outcome_aggregator: outcome_aggregator.clone(),
  test_store: test_store.clone(),
  },
+ project_cache_ready.clone(),
  )
  .map(|b| b.start_observable());
 
@@ -272,6 +277,7 @@ impl ServiceState {
  redis_pools
  .as_ref()
  .map(|pools| pools.project_configs.clone()),
+ project_cache_ready,
  )
  .spawn_handler(project_cache_rx);
 

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

@@ -1,6 +1,7 @@
 //! Types for buffering envelopes.
 
 use std::error::Error;
+use std::future;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
@@ -110,6 +111,7 @@ pub struct EnvelopeBufferService {
  services: Services,
  has_capacity: Arc<AtomicBool>,
  sleep: Duration,
+ project_cache_ready: Arc<AtomicBool>,
 }
 
 /// The maximum amount of time between evaluations of dequeue conditions.
@@ -128,6 +130,7 @@ impl EnvelopeBufferService {
  memory_checker: MemoryChecker,
  global_config_rx: watch::Receiver<global_config::Status>,
  services: Services,
+ project_cache_ready: Arc<AtomicBool>,
  ) -> Option<Self> {
  config.spool_v2().then(|| Self {
  config,
@@ -136,6 +139,7 @@ impl EnvelopeBufferService {
  services,
  has_capacity: Arc::new(AtomicBool::new(true)),
  sleep: Duration::ZERO,
+ project_cache_ready,
  })
  }
 
@@ -169,6 +173,11 @@ impl EnvelopeBufferService {
  tokio::time::sleep(self.sleep).await;
  }
 
+ // In case the project cache is not ready, we don't want to attempt popping.
+ if !self.project_cache_ready.load(Ordering::Relaxed) {
+ return future::pending().await;
+ }
+
  relay_statsd::metric!(
  counter(RelayCounters::BufferReadyToPop) += 1,
  status = "slept"
@@ -234,7 +243,11 @@ impl EnvelopeBufferService {
  .pop()
  .await?
  .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.services.project_cache.send(DequeuedEnvelope(envelope));
+
  self.sleep = Duration::ZERO; // try next pop immediately
  }
  Peek::NotReady(stack_key, envelope) => {
@@ -439,6 +452,7 @@ mod tests {
  watch::Sender<global_config::Status>,
  mpsc::UnboundedReceiver<ProjectCache>,
  mpsc::UnboundedReceiver<TrackOutcome>,
+ Arc<AtomicBool>,
  ) {
  let config = Arc::new(
  Config::from_json_value(serde_json::json!({
@@ -454,6 +468,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));
  (
  EnvelopeBufferService::new(
  config,
@@ -464,18 +479,20 @@ mod tests {
  outcome_aggregator,
  test_store: Addr::dummy(),
  },
+ project_cache_ready.clone(),
  )
  .unwrap(),
  global_tx,
  project_cache_rx,
  outcome_aggregator_rx,
+ project_cache_ready,
  )
  }
 
  #[tokio::test]
  async fn capacity_is_updated() {
  tokio::time::pause();
- let (service, _global_rx, _project_cache_tx, _) = buffer_service();
+ let (service, _global_rx, _project_cache_tx, _, _) = buffer_service();
 
  // Set capacity to false:
  service.has_capacity.store(false, Ordering::Relaxed);
@@ -497,7 +514,7 @@ mod tests {
  #[tokio::test]
  async fn pop_requires_global_config() {
  tokio::time::pause();
- let (service, global_tx, project_cache_rx, _) = buffer_service();
+ let (service, global_tx, project_cache_rx, _, _) = buffer_service();
 
  let addr = service.start();
 
@@ -546,6 +563,7 @@ mod tests {
  )));
 
  let (project_cache, project_cache_rx) = Addr::custom();
+ let project_cache_ready = Arc::new(AtomicBool::new(true));
  let service = EnvelopeBufferService::new(
  config,
  memory_checker,
@@ -555,6 +573,7 @@ mod tests {
  outcome_aggregator: Addr::dummy(),
  test_store: Addr::dummy(),
  },
+ project_cache_ready,
  )
  .unwrap();
  let addr = service.start();
@@ -590,6 +609,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 service = EnvelopeBufferService::new(
  config,
  memory_checker,
@@ -599,6 +619,7 @@ mod tests {
  outcome_aggregator,
  test_store: Addr::dummy(),
  },
+ project_cache_ready,
  )
  .unwrap();
 
@@ -623,4 +644,36 @@ mod tests {
  assert_eq!(outcome.category, DataCategory::TransactionIndexed);
  assert_eq!(outcome.quantity, 1);
  }
+
+ #[tokio::test]
+ async fn output_is_throttled() {
+ tokio::time::pause();
+ let (service, global_tx, mut project_cache_rx, _, _) = buffer_service();
+ global_tx.send_replace(global_config::Status::Ready(Arc::new(
+ GlobalConfig::default(),
+ )));
+
+ let addr = service.start();
+
+ // Send five messages:
+ let envelope = new_envelope(false, "foo");
+ let project_key = envelope.meta().public_key();
+ for _ in 0..5 {
+ addr.send(EnvelopeBuffer::Push(envelope.clone()));
+ }
+ addr.send(EnvelopeBuffer::Ready(project_key));
+
+ tokio::time::sleep(Duration::from_millis(100)).await;
+
+ let mut messages = vec![];
+ project_cache_rx.recv_many(&mut messages, 100).await;
+
+ assert_eq!(
+ messages
+ .iter()
+ .filter(|message| matches!(message, ProjectCache::HandleDequeuedEnvelope(..)))
+ .count(),
+ 1
+ );
+ }
 }

relay-server/src/services/project_cache.rs

@@ -1,5 +1,6 @@
 use std::collections::{BTreeMap, BTreeSet};
 use std::error::Error;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -292,7 +293,7 @@ pub enum ProjectCache {
  UpdateSpoolIndex(UpdateSpoolIndex),
  SpoolHealth(Sender<bool>),
  RefreshIndexCache(RefreshIndexCache),
- HandleDequeuedEnvelope(Box<Envelope>, Sender<()>),
+ HandleDequeuedEnvelope(Box<Envelope>),
  UpdateProject(ProjectKey),
 }
 
@@ -311,7 +312,7 @@ impl ProjectCache {
  Self::UpdateSpoolIndex(_) => "UpdateSpoolIndex",
  Self::SpoolHealth(_) => "SpoolHealth",
  Self::RefreshIndexCache(_) => "RefreshIndexCache",
- Self::HandleDequeuedEnvelope(_, _) => "HandleDequeuedEnvelope",
+ Self::HandleDequeuedEnvelope(_) => "HandleDequeuedEnvelope",
  Self::UpdateProject(_) => "UpdateProject",
  }
  }
@@ -419,11 +420,11 @@ impl FromMessage<SpoolHealth> for ProjectCache {
 }
 
 impl FromMessage<DequeuedEnvelope> for ProjectCache {
- type Response = relay_system::AsyncResponse<()>;
+ type Response = relay_system::NoResponse;
 
- fn from_message(message: DequeuedEnvelope, sender: Sender<()>) -> Self {
+ fn from_message(message: DequeuedEnvelope, _: ()) -> Self {
  let DequeuedEnvelope(envelope) = message;
- Self::HandleDequeuedEnvelope(envelope, sender)
+ Self::HandleDequeuedEnvelope(envelope)
  }
 }
 
@@ -610,6 +611,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>,
 }
 
 #[derive(Debug)]
@@ -1305,7 +1308,7 @@ impl ProjectCacheBroker {
  ProjectCache::RefreshIndexCache(message) => {
  self.handle_refresh_index_cache(message)
  }
- ProjectCache::HandleDequeuedEnvelope(message, sender) => {
+ ProjectCache::HandleDequeuedEnvelope(message) => {
  let envelope_buffer = self
  .services
  .envelope_buffer
@@ -1318,8 +1321,9 @@ impl ProjectCacheBroker {
  "Failed to handle popped envelope"
  );
  }
- // Return response to signal readiness for next envelope:
- sender.send(())
+
+ // We mark the project cache as ready to accept new traffic.
+ self.project_cache_ready.store(true, Ordering::SeqCst);
  }
  ProjectCache::UpdateProject(project) => self.handle_update_project(project),
  }
@@ -1336,6 +1340,7 @@ pub struct ProjectCacheService {
  services: Services,
  global_config_rx: watch::Receiver<global_config::Status>,
  redis: Option<RedisPool>,
+ project_cache_ready: Arc<AtomicBool>,
 }
 
 impl ProjectCacheService {
@@ -1346,13 +1351,15 @@ impl ProjectCacheService {
  services: Services,
  global_config_rx: watch::Receiver<global_config::Status>,
  redis: Option<RedisPool>,
+ project_cache_ready: Arc<AtomicBool>,
  ) -> Self {
  Self {
  config,
  memory_checker,
  services,
  global_config_rx,
  redis,
+ project_cache_ready,
  }
  }
 }
@@ -1367,6 +1374,7 @@ impl Service for ProjectCacheService {
  services,
  mut global_config_rx,
  redis,
+ project_cache_ready,
  } = self;
  let project_cache = services.project_cache.clone();
  let outcome_aggregator = services.outcome_aggregator.clone();
@@ -1448,6 +1456,7 @@ impl Service for ProjectCacheService {
  spool_v1_unspool_handle: SleepHandle::idle(),
  spool_v1,
  global_config,
+ project_cache_ready,
  };
 
  loop {
@@ -1642,6 +1651,7 @@ mod tests {
  buffer_unspool_backoff: RetryBackoff::new(Duration::from_millis(100)),
  }),
  global_config: GlobalConfigStatus::Pending,
+ project_cache_ready: Arc::new(AtomicBool::new(true)),
  },
  buffer,
  )