feat(memory): Add a way to track memory usage

getsentry · iambriccardo · Jul 22, 2024 · Jul 15, 2024 · Jul 16, 2024 · Jul 16, 2024
commit b0ca49db84d68ccae8f31b329a31edba8223a6be
@@ -75,7 +75,7 @@ pub enum BadStoreRequest {
     InvalidEventId,
 
     #[error("failed to queue envelope")]
-    QueueFailed(#[from] BufferError),
+    QueueFailed,
 
     #[error(
         "envelope exceeded size limits for type '{0}' (https://develop.sentry.dev/sdk/envelopes/#size-limits)"
@@ -299,14 +299,18 @@ fn queue_envelope(
     // Split off the envelopes by item type.
     let envelopes = ProcessingGroup::split_envelope(*managed_envelope.take_envelope());
     for (group, envelope) in envelopes {
-        let envelope = buffer_guard
-            .enter(
-                envelope,
-                state.outcome_aggregator().clone(),
-                state.test_store().clone(),
-                group,
-            )
-            .map_err(BadStoreRequest::QueueFailed)?;
+        state.memory_stat().increment();
+        if !state.memory_stat().has_enough_memory() {
+            return Err(BadStoreRequest::QueueFailed);
+        }
+
+        let envelope = ManagedEnvelope::standalone(
+            envelope,
+            state.outcome_aggregator().clone(),
+            state.test_store().clone(),
+            group,
+        );
+
         state.project_cache().send(ValidateEnvelope::new(envelope));
     }
     // The entire envelope is taken for a split above, and it's empty at this point, we can just

@@ -26,7 +26,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::BufferGuard;
+use crate::utils::{BufferGuard, MemoryStat};
 
 /// Indicates the type of failure of the server.
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, thiserror::Error)]
@@ -85,6 +85,7 @@ pub fn create_runtime(name: &str, threads: usize) -> Runtime {
 struct StateInner {
     config: Arc<Config>,
     buffer_guard: Arc<BufferGuard>,
+    memory_stat: MemoryStat,
     registry: Registry,
 }
 
@@ -109,6 +110,8 @@ impl ServiceState {
 
         let buffer_guard = Arc::new(BufferGuard::new(config.envelope_buffer_size()));
 
+        let memory_stat = MemoryStat::new(config.clone());
+
         // Create an address for the `EnvelopeProcessor`, which can be injected into the
         // other services.
         let (processor, processor_rx) = channel(EnvelopeProcessorService::name());
@@ -234,8 +237,9 @@ impl ServiceState {
         };
 
         let state = StateInner {
-            buffer_guard,
             config,
+            buffer_guard,
+            memory_stat,
             registry,
         };
 
@@ -257,6 +261,10 @@ impl ServiceState {
         &self.inner.buffer_guard
     }
 
+    pub fn memory_stat(&self) -> &MemoryStat {
+        &self.inner.memory_stat
+    }
+
     /// Returns the address of the [`ProjectCache`] service.
     pub fn project_cache(&self) -> &Addr<ProjectCache> {
         &self.inner.registry.project_cache

@@ -1,33 +1,110 @@
 use relay_config::Config;
-use std::sync::{Arc, Mutex};
+use std::fmt;
+use std::fmt::Formatter;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, RwLock};
+use sysinfo::System;
+
+/// Count after which the [`MemoryStat`] data will be refreshed.
+const UPDATE_COUNT_THRESHOLD: u64 = 1000;
 
 struct MemoryStatInner {
-    pub lock: Mutex<()>,
     pub used: u64,
     pub total: u64,
 }
 
 impl MemoryStatInner {
     fn new() -> Self {
-        Self {
-            lock: Mutex::new(()),
-            used: 0,
-            total: 0,
-        }
+        Self { used: 0, total: 0 }
+    }
+
+    fn used_percent(&self) -> f32 {
+        (self.used as f32 / self.total as f32).clamp(0.0, 1.0)
     }
 }
 
 #[derive(Clone)]
 pub struct MemoryStat {
-    inner: Arc<MemoryStatInner>,
+    data: Arc<RwLock<MemoryStatInner>>,
+    current_count: Arc<AtomicU64>,
     config: Arc<Config>,
+    system: Arc<System>,
+}
+
+impl fmt::Debug for MemoryStat {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "MemoryStat")
+    }
 }
 
 impl MemoryStat {
     pub fn new(config: Arc<Config>) -> Self {
+        // sysinfo docs suggest to use a single instance of `System` across the program.
+        let system = System::new();
         Self {
-            inner: Arc::new(MemoryStatInner::new()),
+            data: Arc::new(RwLock::new(Self::build_data(&system))),
+            current_count: Arc::new(AtomicU64::new(0)),
             config: config.clone(),
+            system: Arc::new(system),
         }
     }
+
+    pub fn has_enough_memory(&self) -> bool {
+        self.try_update();
+
+        // TODO: we could make an optimization that if we already updated the memory, we can avoid
+        //  acquiring a read lock and just evaluate the bottom expression on the newly added data.
+        let Ok(data) = self.data.read() else {
+            return false;
+        };
+
+        data.used_percent() < self.config.health_max_memory_watermark_percent()
+    }
+
+    pub fn increment(&self) {
+        self.current_count.fetch_add(1, Ordering::Relaxed);
+    }
+
+    fn build_data(system: &System) -> MemoryStatInner {
+        match system.cgroup_limits() {
+            Some(cgroup) => MemoryStatInner {
+                used: cgroup.rss,
+                total: cgroup.total_memory,
+            },
+            None => MemoryStatInner {
+                used: system.used_memory(),
+                total: system.total_memory(),
+            },
+        }
+    }
+
+    fn try_update(&self) {
+        let current_count = self.current_count.load(Ordering::Relaxed);
+
+        if current_count < UPDATE_COUNT_THRESHOLD {
+            return;
+        }
+
+        let Ok(mut data) = self.data.write() else {
+            return;
+        };
+
+        let Ok(_) = self.current_count.compare_exchange_weak(
+            current_count,
+            0,
+            Ordering::Relaxed,
+            Ordering::Relaxed,
+        ) else {
+            return;
+        };
+
+        *data = Self::build_data(&self.system);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    #[test]
+    fn test_data() {}
 }
@@ -26,6 +26,7 @@ pub use self::buffer::*;
 pub use self::dynamic_sampling::*;
 pub use self::garbage::*;
 pub use self::managed_envelope::*;
+pub use self::memory::*;
 pub use self::multipart::*;
 #[cfg(feature = "processing")]
 pub use self::native::*;