wip

getsentry · Sep 11, 2024 · fd0ae7a · fd0ae7a
1 parent 25a3b7a
commit fd0ae7a
Show file tree

Hide file tree

Showing 23 changed files with 143 additions and 75 deletions.
diff --git a/relay-server/src/lib.rs b/relay-server/src/lib.rs
@@ -278,6 +278,7 @@ mod testutils;
 
 use std::sync::Arc;
 
+use futures::StreamExt;
 use relay_config::Config;
 use relay_system::{Controller, Service};
 
@@ -303,6 +304,13 @@ pub fn run(config: Config) -> anyhow::Result<()> {
         Controller::start(config.shutdown_timeout());
         let service = ServiceState::start(config.clone())?;
         HttpServer::new(config, service.clone())?.start();
+
+        for x in service.join_handles() {}
+        // while let Some(res) = service.join_handles().next() {
+
+        // }
+
+        // TODO: await simultaneously
         Controller::shutdown_handle().finished().await;
         anyhow::Ok(())
     })?;

diff --git a/relay-server/src/service.rs b/relay-server/src/service.rs
@@ -9,12 +9,14 @@ use crate::services::stats::RelayStats;
 use anyhow::{Context, Result};
 use axum::extract::FromRequestParts;
 use axum::http::request::Parts;
+use futures::stream::FuturesUnordered;
 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::task::JoinHandle;
 
 use crate::services::cogs::{CogsService, CogsServiceRecorder};
 use crate::services::global_config::{GlobalConfigManager, GlobalConfigService};
@@ -141,6 +143,7 @@ struct StateInner {
     config: Arc<Config>,
     memory_checker: MemoryChecker,
     registry: Registry,
+    join_handles: Vec<JoinHandle<()>>,
 }
 
 /// Server state.
@@ -221,7 +224,7 @@ impl ServiceState {
         let cogs = CogsService::new(&config);
         let cogs = Cogs::new(CogsServiceRecorder::new(&config, cogs.start()));
 
-        EnvelopeProcessorService::new(
+        let processor_handle = EnvelopeProcessorService::new(
             create_processor_pool(&config)?,
             config.clone(),
             global_config_handle,
@@ -251,7 +254,7 @@ impl ServiceState {
 
         // Keep all the services in one context.
         let project_cache_services = Services {
-            envelope_buffer: envelope_buffer.as_ref().map(ObservableEnvelopeBuffer::addr),
+            envelope_buffer: envelope_buffer.as_ref().map(|(b, _)| b.addr()),
             aggregator: aggregator.clone(),
             envelope_processor: processor.clone(),
             outcome_aggregator: outcome_aggregator.clone(),
@@ -301,13 +304,20 @@ impl ServiceState {
             global_config,
             project_cache,
             upstream_relay,
-            envelope_buffer,
+            envelope_buffer: envelope_buffer.as_ref().map(|(b, _)| b.clone()),
         };
 
         let state = StateInner {
             config: config.clone(),
             memory_checker: MemoryChecker::new(memory_stat, config.clone()),
             registry,
+            join_handles: {
+                let mut j = Vec::from_iter([processor_handle]);
+                if let Some((_, handle)) = envelope_buffer {
+                    j.push(handle);
+                }
+                j
+            },
         };
 
         Ok(ServiceState {
@@ -376,6 +386,10 @@ impl ServiceState {
     pub fn outcome_aggregator(&self) -> &Addr<TrackOutcome> {
         &self.inner.registry.outcome_aggregator
     }
+
+    pub fn join_handles(&self) -> &[JoinHandle<()>] {
+        self.inner.join_handles.as_slice()
+    }
 }
 
 fn create_redis_pool(

diff --git a/relay-server/src/services/buffer/mod.rs b/relay-server/src/services/buffer/mod.rs
@@ -11,6 +11,7 @@ use relay_system::Request;
 use relay_system::SendError;
 use relay_system::{Addr, FromMessage, Interface, NoResponse, Receiver, Service};
 use tokio::sync::watch;
+use tokio::task::JoinHandle;
 
 use crate::envelope::Envelope;
 use crate::services::buffer::envelope_buffer::Peek;
@@ -128,12 +129,10 @@ impl EnvelopeBufferService {
     }
 
     /// Returns both the [`Addr`] to this service, and a reference to the capacity flag.
-    pub fn start_observable(self) -> ObservableEnvelopeBuffer {
+    pub fn start_observable(self) -> (ObservableEnvelopeBuffer, JoinHandle<()>) {
         let has_capacity = self.has_capacity.clone();
-        ObservableEnvelopeBuffer {
-            addr: self.start(),
-            has_capacity,
-        }
+        let (addr, join_handle) = self.start_joinable();
+        (ObservableEnvelopeBuffer { addr, has_capacity }, join_handle)
     }
 
     /// Wait for the configured amount of time and make sure the project cache is ready to receive.
@@ -259,7 +258,7 @@ impl EnvelopeBufferService {
 impl Service for EnvelopeBufferService {
     type Interface = EnvelopeBuffer;
 
-    fn spawn_handler(mut self, mut rx: Receiver<Self::Interface>) {
+    fn spawn_handler(mut self, mut rx: Receiver<Self::Interface>) -> JoinHandle<()> {
         let config = self.config.clone();
         let memory_checker = self.memory_checker.clone();
         let mut global_config_rx = self.global_config_rx.clone();
@@ -312,7 +311,7 @@ impl Service for EnvelopeBufferService {
             }
 
             relay_log::info!("EnvelopeBufferService stop");
-        });
+        })
     }
 }
 
@@ -364,7 +363,7 @@ mod tests {
         service.has_capacity.store(false, Ordering::Relaxed);
 
         // Observable has correct value:
-        let ObservableEnvelopeBuffer { addr, has_capacity } = service.start_observable();
+        let (ObservableEnvelopeBuffer { addr, has_capacity }, _) = service.start_observable();
         assert!(!has_capacity.load(Ordering::Relaxed));
 
         // Send a message to trigger update of `has_capacity` flag:

diff --git a/relay-server/src/services/cogs.rs b/relay-server/src/services/cogs.rs
@@ -1,8 +1,8 @@
-use std::sync::atomic::{AtomicBool, Ordering};
-
 use relay_cogs::{CogsMeasurement, CogsRecorder, ResourceId};
 use relay_config::Config;
 use relay_system::{Addr, FromMessage, Interface, Service};
+use std::sync::atomic::{AtomicBool, Ordering};
+use tokio::task::JoinHandle;
 
 use crate::statsd::RelayCounters;
 
@@ -54,12 +54,12 @@ impl CogsService {
 impl Service for CogsService {
     type Interface = CogsReport;
 
-    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) {
+    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) -> JoinHandle<()> {
         tokio::spawn(async move {
             while let Some(message) = rx.recv().await {
                 self.handle_report(message);
             }
-        });
+        })
     }
 }
 

diff --git a/relay-server/src/services/global_config.rs b/relay-server/src/services/global_config.rs
@@ -22,6 +22,7 @@ use relay_system::{Addr, AsyncResponse, Controller, FromMessage, Interface, Serv
 use reqwest::Method;
 use serde::{Deserialize, Serialize};
 use tokio::sync::{mpsc, watch};
+use tokio::task::JoinHandle;
 use tokio::time::Instant;
 
 use crate::services::upstream::{
@@ -338,7 +339,7 @@ impl GlobalConfigService {
 impl Service for GlobalConfigService {
     type Interface = GlobalConfigManager;
 
-    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) {
+    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) -> JoinHandle<()> {
         tokio::spawn(async move {
             let mut shutdown_handle = Controller::shutdown_handle();
 
@@ -384,7 +385,7 @@ impl Service for GlobalConfigService {
                 }
             }
             relay_log::info!("global config service stopped");
-        });
+        })
     }
 }
 

diff --git a/relay-server/src/services/health_check.rs b/relay-server/src/services/health_check.rs
@@ -4,6 +4,7 @@ use relay_config::Config;
 use relay_system::{Addr, AsyncResponse, Controller, FromMessage, Interface, Sender, Service};
 use std::future::Future;
 use tokio::sync::watch;
+use tokio::task::JoinHandle;
 use tokio::time::{timeout, Instant};
 
 use crate::services::metrics::RouterHandle;
@@ -189,13 +190,13 @@ impl HealthCheckService {
 impl Service for HealthCheckService {
     type Interface = HealthCheck;
 
-    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) {
+    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) -> JoinHandle<()> {
         let (update_tx, update_rx) = watch::channel(StatusUpdate::new(Status::Unhealthy));
         let check_interval = self.config.health_refresh_interval();
         // Add 10% buffer to the internal timeouts to avoid race conditions.
         let status_timeout = (check_interval + self.config.health_probe_timeout()).mul_f64(1.1);
 
-        tokio::spawn(async move {
+        let j1 = tokio::spawn(async move {
             let shutdown = Controller::shutdown_handle();
 
             while shutdown.get().is_none() {
@@ -212,7 +213,7 @@ impl Service for HealthCheckService {
             update_tx.send(StatusUpdate::new(Status::Unhealthy)).ok();
         });
 
-        tokio::spawn(async move {
+        let _j2 = tokio::spawn(async move {
             while let Some(HealthCheck(message, sender)) = rx.recv().await {
                 let update = update_rx.borrow();
 
@@ -225,6 +226,8 @@ impl Service for HealthCheckService {
                 });
             }
         });
+
+        j1 // TODO: should return j1 + j2
     }
 }
 

diff --git a/relay-server/src/services/metrics/aggregator.rs b/relay-server/src/services/metrics/aggregator.rs
@@ -246,7 +246,10 @@ impl AggregatorService {
 impl Service for AggregatorService {
     type Interface = Aggregator;
 
-    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) {
+    fn spawn_handler(
+        mut self,
+        mut rx: relay_system::Receiver<Self::Interface>,
+    ) -> tokio::task::JoinHandle<()> {
         tokio::spawn(async move {
             let mut ticker = tokio::time::interval(Duration::from_millis(self.flush_interval_ms));
             let mut shutdown = Controller::shutdown_handle();
@@ -264,7 +267,7 @@ impl Service for AggregatorService {
                     else => break,
                 }
             }
-        });
+        })
     }
 }
 
@@ -361,7 +364,10 @@ mod tests {
     impl Service for TestReceiver {
         type Interface = TestInterface;
 
-        fn spawn_handler(self, mut rx: relay_system::Receiver<Self::Interface>) {
+        fn spawn_handler(
+            self,
+            mut rx: relay_system::Receiver<Self::Interface>,
+        ) -> tokio::task::JoinHandle<()> {
             tokio::spawn(async move {
                 while let Some(message) = rx.recv().await {
                     let buckets = message.0.buckets;
@@ -370,7 +376,7 @@ mod tests {
                         self.add_buckets(buckets);
                     }
                 }
-            });
+            })
         }
     }
 

diff --git a/relay-server/src/services/metrics/router.rs b/relay-server/src/services/metrics/router.rs
@@ -5,6 +5,7 @@ use relay_config::aggregator::Condition;
 use relay_config::{AggregatorServiceConfig, ScopedAggregatorConfig};
 use relay_metrics::MetricNamespace;
 use relay_system::{Addr, NoResponse, Recipient, Service};
+use tokio::task::JoinHandle;
 
 use crate::services::metrics::{
     Aggregator, AggregatorHandle, AggregatorService, FlushBuckets, MergeBuckets,
@@ -53,7 +54,7 @@ impl RouterService {
 impl Service for RouterService {
     type Interface = Aggregator;
 
-    fn spawn_handler(self, mut rx: relay_system::Receiver<Self::Interface>) {
+    fn spawn_handler(self, mut rx: relay_system::Receiver<Self::Interface>) -> JoinHandle<()> {
         tokio::spawn(async move {
             let mut router = StartedRouter::start(self);
             relay_log::info!("metrics router started");
@@ -72,7 +73,7 @@ impl Service for RouterService {
                 }
             }
             relay_log::info!("metrics router stopped");
-        });
+        })
     }
 }
 

diff --git a/relay-server/src/services/outcome.rs b/relay-server/src/services/outcome.rs
@@ -29,6 +29,7 @@ use relay_sampling::evaluation::MatchedRuleIds;
 use relay_statsd::metric;
 use relay_system::{Addr, FromMessage, Interface, NoResponse, Service};
 use serde::{Deserialize, Serialize};
+use tokio::task::JoinHandle;
 
 #[cfg(feature = "processing")]
 use crate::service::ServiceError;
@@ -682,7 +683,7 @@ impl HttpOutcomeProducer {
 impl Service for HttpOutcomeProducer {
     type Interface = TrackRawOutcome;
 
-    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) {
+    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) -> JoinHandle<()> {
         tokio::spawn(async move {
             loop {
                 tokio::select! {
@@ -694,7 +695,7 @@ impl Service for HttpOutcomeProducer {
                     else => break,
                 }
             }
-        });
+        })
     }
 }
 
@@ -775,7 +776,7 @@ impl ClientReportOutcomeProducer {
 impl Service for ClientReportOutcomeProducer {
     type Interface = TrackOutcome;
 
-    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) {
+    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) -> JoinHandle<()> {
         tokio::spawn(async move {
             loop {
                 tokio::select! {
@@ -787,7 +788,7 @@ impl Service for ClientReportOutcomeProducer {
                     else => break,
                 }
             }
-        });
+        })
     }
 }
 
@@ -1034,7 +1035,7 @@ impl OutcomeProducerService {
 impl Service for OutcomeProducerService {
     type Interface = OutcomeProducer;
 
-    fn spawn_handler(self, mut rx: relay_system::Receiver<Self::Interface>) {
+    fn spawn_handler(self, mut rx: relay_system::Receiver<Self::Interface>) -> JoinHandle<()> {
         let Self { config, inner } = self;
 
         tokio::spawn(async move {
@@ -1045,7 +1046,7 @@ impl Service for OutcomeProducerService {
                 broker.handle_message(message, &config);
             }
             relay_log::info!("OutcomeProducer stopped.");
-        });
+        })
     }
 }
 

diff --git a/relay-server/src/services/outcome_aggregator.rs b/relay-server/src/services/outcome_aggregator.rs
@@ -10,6 +10,7 @@ use relay_config::{Config, EmitOutcomes};
 use relay_quotas::{DataCategory, Scoping};
 use relay_statsd::metric;
 use relay_system::{Addr, Controller, Service, Shutdown};
+use tokio::task::JoinHandle;
 
 use crate::services::outcome::{Outcome, OutcomeProducer, TrackOutcome};
 use crate::statsd::RelayTimers;
@@ -138,7 +139,7 @@ impl OutcomeAggregator {
 impl Service for OutcomeAggregator {
     type Interface = TrackOutcome;
 
-    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) {
+    fn spawn_handler(mut self, mut rx: relay_system::Receiver<Self::Interface>) -> JoinHandle<()> {
         tokio::spawn(async move {
             let mut shutdown = Controller::shutdown_handle();
             relay_log::info!("outcome aggregator started");
@@ -157,6 +158,6 @@ impl Service for OutcomeAggregator {
             }
 
             relay_log::info!("outcome aggregator stopped");
-        });
+        })
     }
 }