feat(spool): Check preconditions before unspooling

relay-server/src/endpoints/common.rs

@@ -350,9 +350,7 @@ pub async fn handle_envelope(
  }
 
  // TODO(jjbayer): Remove this check once spool v1 is removed.
- 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.
+ if state.envelope_buffer().is_none() && state.memory_checker().check_memory().is_exceeded() {
  return Err(BadStoreRequest::QueueFailed);
  };
 

relay-server/src/service.rs

@@ -244,6 +244,7 @@ impl ServiceState {
  let envelope_buffer = EnvelopeBufferService::new(
  config.clone(),
  MemoryChecker::new(memory_stat.clone(), config.clone()),
+ global_config_rx.clone(),
  project_cache.clone(),
  )
  .map(|b| b.start_observable());

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

@@ -36,11 +36,18 @@ pub enum PolymorphicEnvelopeBuffer {
  /// An enveloper buffer that uses in-memory envelopes stacks.
  InMemory(EnvelopeBuffer<MemoryStackProvider>),
  /// An enveloper buffer that uses sqlite envelopes stacks.
- #[allow(dead_code)]
  Sqlite(EnvelopeBuffer<SqliteStackProvider>),
 }
 
 impl PolymorphicEnvelopeBuffer {
+ /// Returns true if the implementation stores envelopes on external storage (e.g. disk).
+ pub fn is_external(&self) -> bool {
+ match self {
+ PolymorphicEnvelopeBuffer::InMemory(_) => false,
+ PolymorphicEnvelopeBuffer::Sqlite(_) => true,
+ }
+ }
+
  /// Creates either a memory-based or a disk-based envelope buffer,
  /// depending on the given configuration.
  pub async fn from_config(

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

@@ -10,9 +10,11 @@ use relay_config::Config;
 use relay_system::Request;
 use relay_system::SendError;
 use relay_system::{Addr, FromMessage, Interface, NoResponse, Receiver, Service};
+use tokio::sync::watch;
 
 use crate::envelope::Envelope;
 use crate::services::buffer::envelope_buffer::Peek;
+use crate::services::global_config;
 use crate::services::project_cache::DequeuedEnvelope;
 use crate::services::project_cache::ProjectCache;
 use crate::services::project_cache::UpdateProject;
@@ -89,13 +91,18 @@ impl ObservableEnvelopeBuffer {
 pub struct EnvelopeBufferService {
  config: Arc<Config>,
  memory_checker: MemoryChecker,
+ global_config_rx: watch::Receiver<global_config::Status>,
  project_cache: Addr<ProjectCache>,
  has_capacity: Arc<AtomicBool>,
  sleep: Duration,
  project_cache_ready: Option<Request<()>>,
 }
 
-const DEFAULT_SLEEP: Duration = Duration::from_millis(100);
+/// The maximum amount of time between evaluations of dequeue conditions.
+///
+/// Some condition checks are sync (`has_capacity`), so cannot be awaited. The sleep in cancelled
+/// whenever a new message or a global config update comes in.
+const DEFAULT_SLEEP: Duration = Duration::from_secs(1);
 
 impl EnvelopeBufferService {
  /// Creates a memory or disk based [`EnvelopeBufferService`], depending on the given config.
@@ -105,11 +112,14 @@ impl EnvelopeBufferService {
  pub fn new(
  config: Arc<Config>,
  memory_checker: MemoryChecker,
+ global_config_rx: watch::Receiver<global_config::Status>,
  project_cache: Addr<ProjectCache>,
  ) -> Option<Self> {
  config.spool_v2().then(|| Self {
  config,
  memory_checker,
+
+ global_config_rx,
  project_cache,
  has_capacity: Arc::new(AtomicBool::new(true)),
  sleep: Duration::ZERO,
@@ -127,7 +137,11 @@ 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> {
+ async fn ready_to_pop(
+ &mut self,
+ buffer: &mut PolymorphicEnvelopeBuffer,
+ ) -> Result<(), SendError> {
+ self.system_ready(buffer).await;
  tokio::time::sleep(self.sleep).await;
  if let Some(project_cache_ready) = self.project_cache_ready.as_mut() {
  project_cache_ready.await?;
@@ -137,6 +151,26 @@ impl EnvelopeBufferService {
  Ok(())
  }
 
+ /// Waits until preconditions for unspooling are met.
+ ///
+ /// - We should not pop from disk into memory when relay's overall memory capacity
+ /// has been reached.
+ /// - We need a valid global config to unspool.
+ async fn system_ready(&self, buffer: &PolymorphicEnvelopeBuffer) {
+ loop {
+ // We should not unspool from external storage if memory capacity has been reached.
+ // But if buffer storage is in memory, unspooling can reduce memory usage.
+ let memory_ready =
+ !buffer.is_external() || self.memory_checker.check_memory().has_capacity();
+ let global_config_ready = self.global_config_rx.borrow().is_ready();
+
+ if memory_ready && global_config_ready {
+ return;
+ }
+ tokio::time::sleep(DEFAULT_SLEEP).await;
+ }
+ }
+
  /// Tries to pop an envelope for a ready project.
  async fn try_pop(
  &mut self,
@@ -228,6 +262,7 @@ impl Service for EnvelopeBufferService {
  fn spawn_handler(mut self, mut rx: Receiver<Self::Interface>) {
  let config = self.config.clone();
  let memory_checker = self.memory_checker.clone();
+ let mut global_config_rx = self.global_config_rx.clone();
  tokio::spawn(async move {
  let buffer = PolymorphicEnvelopeBuffer::from_config(&config, memory_checker).await;
 
@@ -244,15 +279,17 @@ impl Service for EnvelopeBufferService {
  buffer.initialize().await;
 
  relay_log::info!("EnvelopeBufferService start");
+ let mut iteration = 0;
  loop {
- relay_log::trace!("EnvelopeBufferService loop");
+ iteration += 1;
+ relay_log::trace!("EnvelopeBufferService loop iteration {iteration}");
 
  tokio::select! {
  // NOTE: we do not select a bias here.
  // On the one hand, we might want to prioritize dequeing over enqueing
  // 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() => {
+ Ok(()) = self.ready_to_pop(&mut buffer) => {
  if let Err(e) = self.try_pop(&mut buffer).await {
  relay_log::error!(
  error = &e as &dyn std::error::Error,
@@ -263,6 +300,10 @@ impl Service for EnvelopeBufferService {
  Some(message) = rx.recv() => {
  self.handle_message(&mut buffer, message).await;
  }
+ _ = global_config_rx.changed() => {
+ relay_log::trace!("EnvelopeBufferService received global config");
+ self.sleep = Duration::ZERO; // Try to pop
+ }
 
  else => break,
  }
@@ -279,14 +320,20 @@ impl Service for EnvelopeBufferService {
 mod tests {
  use std::time::Duration;
 
+ use relay_dynamic_config::GlobalConfig;
+ use tokio::sync::mpsc;
+ use uuid::Uuid;
+
  use crate::testutils::new_envelope;
  use crate::MemoryStat;
 
  use super::*;
 
- #[tokio::test]
- async fn capacity_is_updated() {
- tokio::time::pause();
+ fn buffer_service() -> (
+ EnvelopeBufferService,
+ watch::Sender<global_config::Status>,
+ mpsc::UnboundedReceiver<ProjectCache>,
+ ) {
  let config = Arc::new(
  Config::from_json_value(serde_json::json!({
  "spool": {
@@ -298,7 +345,20 @@ mod tests {
  .unwrap(),
  );
  let memory_checker = MemoryChecker::new(MemoryStat::default(), config.clone());
- let service = EnvelopeBufferService::new(config, memory_checker, Addr::dummy()).unwrap();
+ let (global_tx, global_rx) = watch::channel(global_config::Status::Pending);
+ let (project_cache_addr, project_cache_rx) = Addr::custom();
+ (
+ EnvelopeBufferService::new(config, memory_checker, global_rx, project_cache_addr)
+ .unwrap(),
+ global_tx,
+ project_cache_rx,
+ )
+ }
+
+ #[tokio::test]
+ async fn capacity_is_updated() {
+ tokio::time::pause();
+ let (service, _global_rx, _project_cache_tx) = buffer_service();
 
  // Set capacity to false:
  service.has_capacity.store(false, Ordering::Relaxed);
@@ -318,22 +378,81 @@ mod tests {
  }
 
  #[tokio::test]
- async fn output_is_throttled() {
+ async fn pop_requires_global_config() {
+ tokio::time::pause();
+ let (service, global_tx, project_cache_rx) = buffer_service();
+
+ let addr = service.start();
+
+ // Send five messages:
+ let envelope = new_envelope(false, "foo");
+ let project_key = envelope.meta().public_key();
+ addr.send(EnvelopeBuffer::Push(envelope.clone()));
+ addr.send(EnvelopeBuffer::Ready(project_key));
+
+ tokio::time::sleep(Duration::from_millis(1000)).await;
+
+ // Nothing was dequeued, global config not ready:
+ assert_eq!(project_cache_rx.len(), 0);
+
+ global_tx.send_replace(global_config::Status::Ready(Arc::new(
+ GlobalConfig::default(),
+ )));
+
+ tokio::time::sleep(Duration::from_millis(1000)).await;
+
+ // Dequeued, global config ready:
+ assert_eq!(project_cache_rx.len(), 1);
+ }
+
+ #[tokio::test]
+ async fn pop_requires_memory_capacity() {
  tokio::time::pause();
+
  let config = Arc::new(
  Config::from_json_value(serde_json::json!({
  "spool": {
  "envelopes": {
- "version": "experimental"
+ "version": "experimental",
+ "path": std::env::temp_dir().join(Uuid::new_v4().to_string()),
  }
+ },
+ "health": {
+ "max_memory_bytes": 0,
  }
  }))
  .unwrap(),
  );
  let memory_checker = MemoryChecker::new(MemoryStat::default(), config.clone());
- let (project_cache_addr, mut project_cache_rx) = Addr::custom();
+ let (_, global_rx) = watch::channel(global_config::Status::Ready(Arc::new(
+ GlobalConfig::default(),
+ )));
+
+ let (project_cache_addr, project_cache_rx) = Addr::custom();
  let service =
- EnvelopeBufferService::new(config, memory_checker, project_cache_addr).unwrap();
+ EnvelopeBufferService::new(config, memory_checker, global_rx, project_cache_addr)
+ .unwrap();
+ let addr = service.start();
+
+ // Send five messages:
+ let envelope = new_envelope(false, "foo");
+ let project_key = envelope.meta().public_key();
+ addr.send(EnvelopeBuffer::Push(envelope.clone()));
+ addr.send(EnvelopeBuffer::Ready(project_key));
+
+ tokio::time::sleep(Duration::from_millis(1000)).await;
+
+ // Nothing was dequeued, memory not ready:
+ assert_eq!(project_cache_rx.len(), 0);
+ }
+
+ #[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();
 

relay-server/src/services/global_config.rs

@@ -141,7 +141,8 @@ pub enum Status {
 }
 
 impl Status {
- fn is_ready(&self) -> bool {
+ /// Returns `true` if the global config is ready to be read.
+ pub fn is_ready(&self) -> bool {
  matches!(self, Self::Ready(_))
  }
 }

tests/integration/test_query.py

@@ -120,6 +120,9 @@ def test_query_retry(failure_type, mini_sentry, relay):
 
  @mini_sentry.app.endpoint("get_project_config")
  def get_project_config():
+ if flask_request.json.get("global") is True:
+ return original_endpoint()
+
  nonlocal retry_count
  retry_count += 1
  print("RETRY", retry_count)