Fix

relay-server/src/service.rs

@@ -5,17 +5,6 @@ use std::time::Duration;
 
 use crate::metrics::{MetricOutcomes, MetricStats};
 use crate::services::buffer::{EnvelopeBufferService, ObservableEnvelopeBuffer};
-use crate::services::stats::RelayStats;
-use anyhow::{Context, Result};
-use axum::extract::FromRequestParts;
-use axum::http::request::Parts;
-use rayon::ThreadPool;
-use relay_cogs::Cogs;
-use relay_config::{Config, RedisConnection, RedisPoolConfigs};
-use relay_redis::{RedisConfigOptions, RedisError, RedisPool, RedisPools};
-use relay_system::{channel, Addr, Service};
-use tokio::runtime::Runtime;
-
 use crate::services::cogs::{CogsService, CogsServiceRecorder};
 use crate::services::global_config::{GlobalConfigManager, GlobalConfigService};
 use crate::services::health_check::{HealthCheck, HealthCheckService};
@@ -25,11 +14,22 @@ use crate::services::outcome_aggregator::OutcomeAggregator;
 use crate::services::processor::{self, EnvelopeProcessor, EnvelopeProcessorService};
 use crate::services::project_cache::{ProjectCache, ProjectCacheService, Services};
 use crate::services::relays::{RelayCache, RelayCacheService};
+use crate::services::stats::RelayStats;
 #[cfg(feature = "processing")]
 use crate::services::store::StoreService;
 use crate::services::test_store::{TestStore, TestStoreService};
 use crate::services::upstream::{UpstreamRelay, UpstreamRelayService};
 use crate::utils::{MemoryChecker, MemoryStat};
+use anyhow::{Context, Result};
+use axum::extract::FromRequestParts;
+use axum::http::request::Parts;
+use rayon::ThreadPool;
+use relay_cogs::Cogs;
+use relay_config::{Config, RedisConnection, RedisPoolConfigs};
+use relay_redis::{RedisConfigOptions, RedisError, RedisPool, RedisPools};
+use relay_system::{channel, Addr, Service};
+use tokio::runtime::Runtime;
+use tokio::sync::watch;
 
 /// Indicates the type of failure of the server.
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, thiserror::Error)]
@@ -240,10 +240,16 @@ impl ServiceState {
  )
  .spawn_handler(processor_rx);
 
+ // We create a watch channel that is used to monitor availability of the project cache
+ // service for processing new envelopes. This is used as a simple backpressure mechanism
+ // between services.
+ let (envelope_processing_availability, _) = watch::channel(true);
+
  let envelope_buffer = EnvelopeBufferService::new(
  config.clone(),
  MemoryChecker::new(memory_stat.clone(), config.clone()),
  project_cache.clone(),
+ envelope_processing_availability.subscribe(),
  )
  .map(|b| b.start_observable());
 
@@ -266,6 +272,7 @@ impl ServiceState {
  redis_pools
  .as_ref()
  .map(|pools| pools.project_configs.clone()),
+ envelope_processing_availability,
  )
  .spawn_handler(project_cache_rx);
 

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

@@ -1,15 +1,15 @@
 //! Types for buffering envelopes.
 
-use std::sync::atomic::AtomicBool;
-use std::sync::atomic::Ordering;
-use std::sync::Arc;
-use std::time::Duration;
-
 use relay_base_schema::project::ProjectKey;
 use relay_config::Config;
-use relay_system::Request;
 use relay_system::SendError;
 use relay_system::{Addr, FromMessage, Interface, NoResponse, Receiver, Service};
+use std::future;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::watch;
 
 use crate::envelope::Envelope;
 use crate::services::buffer::envelope_buffer::Peek;
@@ -91,11 +91,11 @@ pub struct EnvelopeBufferService {
  memory_checker: MemoryChecker,
  project_cache: Addr<ProjectCache>,
  has_capacity: Arc<AtomicBool>,
- sleep: Duration,
- project_cache_ready: Option<Request<()>>,
+ pop_delay: Duration,
+ project_cache_envelope_processing_availability: watch::Receiver<bool>,
 }
 
-const DEFAULT_SLEEP: Duration = Duration::from_millis(100);
+const DEFAULT_POP_DELAY: Duration = Duration::from_millis(100);
 
 impl EnvelopeBufferService {
  /// Creates a memory or disk based [`EnvelopeBufferService`], depending on the given config.
@@ -106,14 +106,15 @@ impl EnvelopeBufferService {
  config: Arc<Config>,
  memory_checker: MemoryChecker,
  project_cache: Addr<ProjectCache>,
+ project_cache_envelope_processing_availability: watch::Receiver<bool>,
  ) -> Option<Self> {
  config.spool_v2().then(|| Self {
  config,
  memory_checker,
  project_cache,
  has_capacity: Arc::new(AtomicBool::new(true)),
- sleep: Duration::ZERO,
- project_cache_ready: None,
+ pop_delay: Duration::ZERO,
+ project_cache_envelope_processing_availability,
  })
  }
 
@@ -128,11 +129,19 @@ impl EnvelopeBufferService {
 
  /// Wait for the configured amount of time and make sure the project cache is ready to receive.
  async fn ready_to_pop(&mut self) -> Result<(), SendError> {
- tokio::time::sleep(self.sleep).await;
- if let Some(project_cache_ready) = self.project_cache_ready.take() {
- project_cache_ready.await?;
+ tokio::time::sleep(self.pop_delay).await;
+
+ if let Ok(availability) = self
+ .project_cache_envelope_processing_availability
+ .wait_for(|availability| *availability)
+ .await
+ {
+ if *availability {
+ return Ok(());
+ }
  }
- Ok(())
+
+ future::pending().await
  }
 
  /// Tries to pop an envelope for a ready project.
@@ -144,7 +153,8 @@ impl EnvelopeBufferService {
  match buffer.peek().await? {
  Peek::Empty => {
  relay_log::trace!("EnvelopeBufferService empty");
- self.sleep = Duration::MAX; // wait for reset by `handle_message`.
+ // Wait for reset by `handle_message`.
+ self.pop_delay = Duration::MAX;
  }
  Peek::Ready(_) => {
  relay_log::trace!("EnvelopeBufferService pop");
@@ -153,9 +163,9 @@ impl EnvelopeBufferService {
  .await?
  .expect("Element disappeared despite exclusive excess");
 
- self.project_cache_ready
-  .replace(self.project_cache.send(DequeuedEnvelope(envelope)));
- self.sleep = Duration::ZERO; // try next pop immediately
+ self.project_cache.send(DequeuedEnvelope(envelope));
+ // Try next pop immediately.
+ self.pop_delay = Duration::ZERO;
  }
  Peek::NotReady(stack_key, envelope) => {
  relay_log::trace!("EnvelopeBufferService request update");
@@ -168,9 +178,9 @@ impl EnvelopeBufferService {
  self.project_cache.send(UpdateProject(sampling_key));
  }
  }
- // deprioritize the stack to prevent head-of-line blocking
+ // Deprioritize the stack to prevent head-of-line blocking.
  buffer.mark_seen(&stack_key);
- self.sleep = DEFAULT_SLEEP;
+ self.pop_delay = DEFAULT_POP_DELAY;
  }
  }
 
@@ -202,7 +212,7 @@ impl EnvelopeBufferService {
  buffer.mark_ready(&project_key, true);
  }
  };
- self.sleep = Duration::ZERO;
+ self.pop_delay = Duration::ZERO;
  }
 
  /// Pushes an [`Envelope`] to the [`PolymorphicEnvelopeBuffer`].
@@ -235,6 +245,8 @@ impl Service for EnvelopeBufferService {
  let config = self.config.clone();
  let memory_checker = self.memory_checker.clone();
  tokio::spawn(async move {
+ let mut ticker = tokio::time::interval(Duration::from_millis(100));
+
  let buffer = PolymorphicEnvelopeBuffer::from_config(&config, memory_checker).await;
 
  let mut buffer = match buffer {
@@ -249,15 +261,14 @@ impl Service for EnvelopeBufferService {
  };
  buffer.initialize().await;
 
- let mut ticker = tokio::time::interval(Duration::from_millis(100));
-
  relay_log::info!("EnvelopeBufferService start");
  loop {
  relay_log::trace!("EnvelopeBufferService loop");
 
  tokio::select! {
  // NOTE: we do not select a bias here.
- // On the one hand, we might want to prioritize dequeing over enqueing
+ //
+ // On the one hand, we might want to prioritize dequeuing over enqueuing
  // so we do not exceed the buffer capacity by starving the dequeue.
  // on the other hand, prioritizing old messages violates the LIFO design.
  Ok(()) = self.ready_to_pop() => {
@@ -308,7 +319,13 @@ mod tests {
  .unwrap(),
  );
  let memory_checker = MemoryChecker::new(MemoryStat::default(), config.clone());
- let service = EnvelopeBufferService::new(config, memory_checker, Addr::dummy()).unwrap();
+ let service = EnvelopeBufferService::new(
+ config,
+ memory_checker,
+ Addr::dummy(),
+ watch::channel(true).0.subscribe(),
+ )
+ .unwrap();
 
  // Set capacity to false:
  service.has_capacity.store(false, Ordering::Relaxed);
@@ -342,24 +359,47 @@ mod tests {
  );
  let memory_checker = MemoryChecker::new(MemoryStat::default(), config.clone());
  let (project_cache_addr, mut project_cache_rx) = Addr::custom();
- let service =
- EnvelopeBufferService::new(config, memory_checker, project_cache_addr).unwrap();
+ let (envelope_processing_availability, _) = watch::channel(true);
+ let service = EnvelopeBufferService::new(
+ config,
+ memory_checker,
+ project_cache_addr,
+ envelope_processing_availability.subscribe(),
+ )
+ .unwrap();
 
  let addr = service.start();
 
- // Send five messages:
+ // Send one message with availability set to true.
  let envelope = new_envelope(false, "foo");
  let project_key = envelope.meta().public_key();
+ addr.send(EnvelopeBuffer::Push(envelope.clone()));
+
+ // We wait for the message to be processed.
+ tokio::time::sleep(Duration::from_millis(100)).await;
+
+ // We simulate the availability set to false.
+ envelope_processing_availability.send(false).unwrap();
+
+ // We send five more messages.
  for _ in 0..5 {
  addr.send(EnvelopeBuffer::Push(envelope.clone()));
  }
  addr.send(EnvelopeBuffer::Ready(project_key));
 
- tokio::time::sleep(Duration::from_millis(1000)).await;
-
- // Project cache received only one envelope:
- assert_eq!(project_cache_rx.len(), 1); // without throttling, this would be 5.
- assert!(project_cache_rx.try_recv().is_ok());
+ // We expect the project cache to have received only one envelope since the other five
+ // have been blocked on the check.
+ assert_eq!(project_cache_rx.len(), 1);
+ assert!(project_cache_rx.recv().await.is_some());
  assert_eq!(project_cache_rx.len(), 0);
+
+ // We set the availability back to true.
+ envelope_processing_availability.send(true).unwrap();
+
+ // We wait for all the envelopes to be processed.
+ tokio::time::sleep(Duration::from_millis(100)).await;
+
+ // We expect to have received all five of them.
+ assert_eq!(project_cache_rx.len(), 5);
  }
 }