From 7ee54ef93344730325c68eee8713070cd10687ab Mon Sep 17 00:00:00 2001 From: Andrew McKenzie Date: Thu, 20 Jul 2023 12:19:32 +0100 Subject: [PATCH] add taskman app, apply usage to each app --- Cargo.lock | 37 +++ Cargo.toml | 2 + db_store/src/iam_auth_pool.rs | 27 +- db_store/src/metric_tracker.rs | 39 +-- db_store/src/settings.rs | 34 +- file_store/Cargo.toml | 4 +- file_store/src/error.rs | 2 + file_store/src/file_info_poller.rs | 81 +++-- file_store/src/file_sink.rs | 349 +++++++++++---------- file_store/src/file_source.rs | 6 +- file_store/src/file_upload.rs | 54 +++- ingest/Cargo.toml | 2 + ingest/src/main.rs | 14 +- ingest/src/server_iot.rs | 68 ++-- ingest/src/server_mobile.rs | 160 +++++----- iot_config/Cargo.toml | 2 + iot_config/src/gateway_service.rs | 5 - iot_config/src/main.rs | 86 +++-- iot_config/src/org_service.rs | 5 - iot_config/src/route_service.rs | 36 +-- iot_packet_verifier/Cargo.toml | 2 + iot_packet_verifier/src/burner.rs | 17 +- iot_packet_verifier/src/daemon.rs | 117 +++---- iot_packet_verifier/src/main.rs | 2 +- iot_verifier/Cargo.toml | 2 + iot_verifier/src/entropy_loader.rs | 24 +- iot_verifier/src/gateway_cache.rs | 1 + iot_verifier/src/gateway_updater.rs | 21 +- iot_verifier/src/hex_density.rs | 27 +- iot_verifier/src/loader.rs | 125 ++++---- iot_verifier/src/main.rs | 272 +++++++++++----- iot_verifier/src/packet_loader.rs | 92 +++--- iot_verifier/src/poc.rs | 19 +- iot_verifier/src/purger.rs | 102 +++--- iot_verifier/src/region_cache.rs | 1 + iot_verifier/src/rewarder.rs | 45 ++- iot_verifier/src/runner.rs | 218 ++++++------- iot_verifier/src/tx_scaler.rs | 32 +- mobile_config/Cargo.toml | 2 + mobile_config/src/main.rs | 79 ++--- mobile_packet_verifier/Cargo.toml | 2 + mobile_packet_verifier/src/daemon.rs | 89 +++--- mobile_verifier/Cargo.toml | 4 +- mobile_verifier/src/cli/reward_from_db.rs | 7 +- mobile_verifier/src/cli/server.rs | 138 ++++---- mobile_verifier/src/data_session.rs | 27 +- mobile_verifier/src/heartbeats.rs | 11 + mobile_verifier/src/rewarder.rs | 11 + mobile_verifier/src/speedtests.rs | 11 + mobile_verifier/src/subscriber_location.rs | 14 +- poc_entropy/Cargo.toml | 2 + poc_entropy/src/entropy_generator.rs | 36 ++- poc_entropy/src/main.rs | 48 +-- poc_entropy/src/server.rs | 13 +- price/Cargo.toml | 2 + price/src/main.rs | 73 ++--- price/src/price_generator.rs | 53 ++-- reward_index/Cargo.toml | 2 + reward_index/src/indexer.rs | 28 +- reward_index/src/main.rs | 38 +-- solana/Cargo.toml | 3 + solana/src/balance_monitor.rs | 70 +++-- task_manager/src/lib.rs | 304 +++++++++--------- 63 files changed, 1691 insertions(+), 1508 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 88398a676..4d4333c83 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2467,6 +2467,7 @@ dependencies = [ name = "file-store" version = "0.1.0" dependencies = [ + "anyhow", "async-compression", "async-trait", "aws-config", @@ -2501,6 +2502,7 @@ dependencies = [ "sqlx", "strum", "strum_macros", + "task-manager", "tempfile", "thiserror", "tokio", @@ -3262,8 +3264,10 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.6", + "task-manager", "thiserror", "tokio", + "tokio-util", "tonic", "tracing", "tracing-subscriber", @@ -3328,9 +3332,11 @@ dependencies = [ "serde", "serde_json", "sqlx", + "task-manager", "thiserror", "tokio", "tokio-stream", + "tokio-util", "tonic", "tracing", "tracing-subscriber", @@ -3361,8 +3367,10 @@ dependencies = [ "serde", "solana", "sqlx", + "task-manager", "thiserror", "tokio", + "tokio-util", "tonic", "tracing", "tracing-subscriber", @@ -3409,8 +3417,10 @@ dependencies = [ "serde_json", "sha2 0.10.6", "sqlx", + "task-manager", "thiserror", "tokio", + "tokio-util", "tonic", "tracing", "tracing-subscriber", @@ -3959,9 +3969,11 @@ dependencies = [ "serde", "serde_json", "sqlx", + "task-manager", "thiserror", "tokio", "tokio-stream", + "tokio-util", "tonic", "tracing", "tracing-subscriber", @@ -4018,8 +4030,10 @@ dependencies = [ "sha2 0.10.6", "solana", "sqlx", + "task-manager", "thiserror", "tokio", + "tokio-util", "tonic", "tracing", "tracing-subscriber", @@ -4061,8 +4075,10 @@ dependencies = [ "serde_json", "sha2 0.10.6", "sqlx", + "task-manager", "thiserror", "tokio", + "tokio-util", "tonic", "tracing", "tracing-subscriber", @@ -4645,8 +4661,10 @@ dependencies = [ "prost", "serde", "serde_json", + "task-manager", "thiserror", "tokio", + "tokio-util", "tonic", "tower", "tracing", @@ -4722,8 +4740,10 @@ dependencies = [ "serde_json", "solana-client", "solana-sdk", + "task-manager", "thiserror", "tokio", + "tokio-util", "tracing", "tracing-subscriber", "triggered", @@ -5254,8 +5274,10 @@ dependencies = [ "serde_json", "sha2 0.10.6", "sqlx", + "task-manager", "thiserror", "tokio", + "tokio-util", "tonic", "tracing", "tracing-subscriber", @@ -5885,6 +5907,7 @@ version = "0.1.0" dependencies = [ "anchor-client", "anchor-lang", + "anyhow", "async-trait", "clap 4.1.11", "data-credits", @@ -5898,8 +5921,10 @@ dependencies = [ "solana-program", "solana-sdk", "spl-token", + "task-manager", "thiserror", "tokio", + "tokio-util", "tracing", "triggered", ] @@ -6965,6 +6990,18 @@ dependencies = [ "unicode-xid 0.2.4", ] +[[package]] +name = "task-manager" +version = "0.1.0" +dependencies = [ + "anyhow", + "futures", + "futures-util", + "tokio", + "tokio-util", + "triggered", +] + [[package]] name = "tempfile" version = "3.3.0" diff --git a/Cargo.toml b/Cargo.toml index fd747c1d0..62b26a1e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ members = [ "reward_index", "reward_scheduler", "solana", + "task_manager", ] [workspace.package] @@ -99,6 +100,7 @@ itertools = "*" data-credits = {git = "https://github.com/helium/helium-program-library.git", tag = "v0.1.0"} helium-sub-daos = {git = "https://github.com/helium/helium-program-library.git", tag = "v0.1.0"} price-oracle = {git = "https://github.com/helium/helium-program-library.git", tag = "v0.1.0"} +tokio-util = "0" [patch.crates-io] sqlx = { git = "https://github.com/helium/sqlx.git", rev = "92a2268f02e0cac6fccb34d3e926347071dbb88d" } diff --git a/db_store/src/iam_auth_pool.rs b/db_store/src/iam_auth_pool.rs index 2e912008b..a21c801f7 100644 --- a/db_store/src/iam_auth_pool.rs +++ b/db_store/src/iam_auth_pool.rs @@ -13,10 +13,7 @@ use aws_types::{ }; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -pub async fn connect( - settings: &Settings, - shutdown: triggered::Listener, -) -> Result<(Pool, futures::future::BoxFuture<'static, Result>)> { +pub async fn connect(settings: &Settings) -> Result> { let aws_config = aws_config::load_from_env().await; let client = aws_sdk_sts::Client::new(&aws_config); let connect_parameters = ConnectParameters::try_from(settings)?; @@ -28,43 +25,27 @@ pub async fn connect( .await?; let cloned_pool = pool.clone(); - let join_handle = - tokio::spawn(async move { run(client, connect_parameters, cloned_pool, shutdown).await }); - - Ok(( - pool, - Box::pin(async move { - match join_handle.await { - Ok(Ok(())) => Ok(()), - Ok(Err(err)) => Err(err), - Err(err) => Err(Error::from(err)), - } - }), - )) + tokio::spawn(async move { run(client, connect_parameters, cloned_pool).await }); + + Ok(pool) } async fn run( client: aws_sdk_sts::Client, connect_parameters: ConnectParameters, pool: Pool, - shutdown: triggered::Listener, ) -> Result { let duration = std::time::Duration::from_secs(connect_parameters.iam_duration_seconds as u64) - Duration::from_secs(120); loop { - let shutdown = shutdown.clone(); - tokio::select! { - _ = shutdown => break, _ = tokio::time::sleep(duration) => { let connect_options = connect_parameters.connect_options(&client).await?; pool.set_connect_options(connect_options); } } } - - Ok(()) } struct ConnectParameters { diff --git a/db_store/src/metric_tracker.rs b/db_store/src/metric_tracker.rs index f0175de3e..19357c237 100644 --- a/db_store/src/metric_tracker.rs +++ b/db_store/src/metric_tracker.rs @@ -1,47 +1,20 @@ use std::time::Duration; -use crate::{Error, Result}; - const DURATION: Duration = Duration::from_secs(300); -pub async fn start( - app_name: &str, - pool: sqlx::Pool, - shutdown: triggered::Listener, -) -> Result> { +pub async fn start(app_name: &str, pool: sqlx::Pool) { let pool_size_name = format!("{app_name}_db_pool_size"); let pool_idle_name = format!("{app_name}_db_pool_idle"); - let join_handle = - tokio::spawn(async move { run(pool_size_name, pool_idle_name, pool, shutdown).await }); - - Ok(Box::pin(async move { - match join_handle.await { - Ok(()) => Ok(()), - Err(err) => Err(Error::from(err)), - } - })) + tokio::spawn(async move { run(pool_size_name, pool_idle_name, pool).await }); } -async fn run( - size_name: String, - idle_name: String, - pool: sqlx::Pool, - shutdown: triggered::Listener, -) { +async fn run(size_name: String, idle_name: String, pool: sqlx::Pool) { let mut trigger = tokio::time::interval(DURATION); loop { - let shutdown = shutdown.clone(); + trigger.tick().await; - tokio::select! { - _ = shutdown => { - tracing::info!("db_store: MetricTracker shutting down"); - break; - } - _ = trigger.tick() => { - metrics::gauge!(size_name.clone(), pool.size() as f64); - metrics::gauge!(idle_name.clone(), pool.num_idle() as f64); - } - } + metrics::gauge!(size_name.clone(), pool.size() as f64); + metrics::gauge!(idle_name.clone(), pool.num_idle() as f64); } } diff --git a/db_store/src/settings.rs b/db_store/src/settings.rs index 602ae17b5..9a46b7b74 100644 --- a/db_store/src/settings.rs +++ b/db_store/src/settings.rs @@ -37,37 +37,19 @@ fn default_auth_type() -> AuthType { } impl Settings { - pub async fn connect( - &self, - app_name: &str, - shutdown: triggered::Listener, - ) -> Result<(Pool, futures::future::BoxFuture<'static, Result>)> { + pub async fn connect(&self, app_name: &str) -> Result> { match self.auth_type { AuthType::Postgres => match self.simple_connect().await { - Ok(pool) => Ok(( - pool.clone(), - metric_tracker::start(app_name, pool, shutdown).await?, - )), + Ok(pool) => { + metric_tracker::start(app_name, pool.clone()).await; + Ok(pool) + } Err(err) => Err(err), }, AuthType::Iam => { - let (pool, iam_auth_handle) = - iam_auth_pool::connect(self, shutdown.clone()).await?; - let metric_handle = metric_tracker::start(app_name, pool.clone(), shutdown).await?; - - let handle = - tokio::spawn(async move { tokio::try_join!(iam_auth_handle, metric_handle) }); - - Ok(( - pool, - Box::pin(async move { - match handle.await { - Ok(Err(err)) => Err(err), - Err(err) => Err(Error::from(err)), - Ok(_) => Ok(()), - } - }), - )) + let pool = iam_auth_pool::connect(self).await?; + metric_tracker::start(app_name, pool.clone()).await; + Ok(pool) } } } diff --git a/file_store/Cargo.toml b/file_store/Cargo.toml index 5d0d6d1ae..cf03db63b 100644 --- a/file_store/Cargo.toml +++ b/file_store/Cargo.toml @@ -7,13 +7,14 @@ authors.workspace = true license.workspace = true [dependencies] +anyhow = {workspace = true} clap = {workspace = true} config = {workspace = true} serde = {workspace = true} serde_json = {workspace = true} thiserror = {workspace = true} tokio = { workspace = true } -tokio-util = "0" +tokio-util = { workspace = true } tokio-stream = {workspace = true} triggered = {workspace = true} async-compression = {version = "0", features = ["tokio", "gzip"]} @@ -46,6 +47,7 @@ sqlx = {workspace = true} async-trait = {workspace = true} derive_builder = "0" retainer = {workspace = true} +task-manager = { path = "../task_manager" } [dev-dependencies] hex-literal = "0" diff --git a/file_store/src/error.rs b/file_store/src/error.rs index 9aae093be..9551fa1cf 100644 --- a/file_store/src/error.rs +++ b/file_store/src/error.rs @@ -32,6 +32,8 @@ pub enum Error { SendTimeout, #[error("shutting down")] Shutdown, + #[error("error building file info poller")] + FileInfoPollerError(#[from] crate::file_info_poller::FileInfoPollerConfigBuilderError), } #[derive(Error, Debug)] diff --git a/file_store/src/file_info_poller.rs b/file_store/src/file_info_poller.rs index cfdaa0296..b357151d2 100644 --- a/file_store/src/file_info_poller.rs +++ b/file_store/src/file_info_poller.rs @@ -1,9 +1,10 @@ use crate::{traits::MsgDecode, Error, FileInfo, FileStore, FileType, Result}; use chrono::{DateTime, Duration, TimeZone, Utc}; use derive_builder::Builder; -use futures::{stream::BoxStream, StreamExt}; +use futures::{future::LocalBoxFuture, stream::BoxStream, StreamExt, TryFutureExt}; use retainer::Cache; use std::marker::PhantomData; +use task_manager::ManagedTask; use tokio::sync::mpsc::{error::TrySendError, Receiver, Sender}; const DEFAULT_POLL_DURATION_SECS: i64 = 30; @@ -39,7 +40,8 @@ pub enum LookbackBehavior { } #[derive(Debug, Clone, Builder)] -pub struct FileInfoPoller { +#[builder(pattern = "owned")] +pub struct FileInfoPollerConfig { #[builder(default = "Duration::seconds(DEFAULT_POLL_DURATION_SECS)")] poll_duration: Duration, db: sqlx::Pool, @@ -54,35 +56,51 @@ pub struct FileInfoPoller { p: PhantomData, } -impl FileInfoPoller +pub struct FileInfoPollerServer { + config: FileInfoPollerConfig, + sender: Sender>, +} + +impl FileInfoPollerConfigBuilder +where + T: Clone, +{ + pub fn create(self) -> Result<(Receiver>, FileInfoPollerServer)> { + let config = self.build()?; + let (sender, receiver) = tokio::sync::mpsc::channel(config.queue_size); + + Ok((receiver, FileInfoPollerServer { config, sender })) + } +} + +impl ManagedTask for FileInfoPollerServer where T: MsgDecode + TryFrom + Send + Sync + 'static, { - pub async fn start( - self, - shutdown: triggered::Listener, - ) -> Result<( - Receiver>, - impl std::future::Future, - )> { - let (sender, receiver) = tokio::sync::mpsc::channel(self.queue_size); - let join_handle = tokio::spawn(async move { self.run(shutdown, sender).await }); - - Ok((receiver, async move { - match join_handle.await { - Ok(Ok(())) => Ok(()), - Ok(Err(err)) => Err(err), - Err(err) => Err(Error::from(err)), - } - })) + fn start_task( + self: Box, + shutdown_listener: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + let handle = tokio::spawn(self.run(shutdown_listener)); + + Box::pin( + handle + .map_err(anyhow::Error::from) + .and_then(|result| async move { result.map_err(anyhow::Error::from) }), + ) } +} - async fn run(self, shutdown: triggered::Listener, sender: Sender>) -> Result { +impl FileInfoPollerServer +where + T: MsgDecode + TryFrom + Send + Sync + 'static, +{ + async fn run(self, shutdown: triggered::Listener) -> Result { let cache = create_cache(); let mut poll_trigger = tokio::time::interval(self.poll_duration()); let mut cleanup_trigger = tokio::time::interval(CLEAN_DURATION); - let mut latest_ts = db::latest_ts(&self.db, self.file_type).await?; + let mut latest_ts = db::latest_ts(&self.config.db, self.config.file_type).await?; loop { let after = self.after(latest_ts); @@ -95,10 +113,10 @@ where } _ = cleanup_trigger.tick() => self.clean(&cache).await?, _ = poll_trigger.tick() => { - let files = self.store.list_all(self.file_type, after, before).await?; + let files = self.config.store.list_all(self.config.file_type, after, before).await?; for file in files { - if !is_already_processed(&self.db, &cache, &file).await? { - if send_stream(&sender, &self.store, file.clone()).await? { + if !is_already_processed(&self.config.db, &cache, &file).await? { + if send_stream(&self.sender, &self.config.store, file.clone()).await? { latest_ts = Some(file.timestamp); cache_file(&cache, &file).await; } else { @@ -114,8 +132,8 @@ where } fn after(&self, latest: Option>) -> DateTime { - let latest_offset = latest.map(|lt| lt - self.offset); - match self.lookback { + let latest_offset = latest.map(|lt| lt - self.config.offset); + match self.config.lookback { LookbackBehavior::StartAfter(start_after) => latest_offset.unwrap_or(start_after), LookbackBehavior::Max(max_lookback) => { let max_ts = Utc::now() - max_lookback; @@ -126,12 +144,15 @@ where async fn clean(&self, cache: &MemoryFileCache) -> Result { cache.purge(4, 0.25).await; - db::clean(&self.db, &self.file_type).await?; + db::clean(&self.config.db, &self.config.file_type).await?; Ok(()) } fn poll_duration(&self) -> std::time::Duration { - self.poll_duration.to_std().unwrap_or(DEFAULT_POLL_DURATION) + self.config + .poll_duration + .to_std() + .unwrap_or(DEFAULT_POLL_DURATION) } } @@ -263,7 +284,7 @@ mod db { FROM files_processed WHERE file_type = $1 ORDER BY file_timestamp DESC - OFFSET 100 + OFFSET 100 ) "#, ) diff --git a/file_store/src/file_sink.rs b/file_store/src/file_sink.rs index 5fd0f7fbf..ca424b0f9 100644 --- a/file_store/src/file_sink.rs +++ b/file_store/src/file_sink.rs @@ -1,13 +1,14 @@ -use crate::{file_upload, Error, Result}; +use crate::{file_upload::FileUpload, Error, Result}; use async_compression::tokio::write::GzipEncoder; use bytes::Bytes; use chrono::{DateTime, Duration, Utc}; -use futures::SinkExt; +use futures::{future::LocalBoxFuture, SinkExt, TryFutureExt}; use metrics::Label; use std::{ io, mem, path::{Path, PathBuf}, }; +use task_manager::ManagedTask; use tokio::{ fs::{self, File, OpenOptions}, io::{AsyncWriteExt, BufWriter}, @@ -62,29 +63,22 @@ pub struct FileSinkBuilder { tmp_path: PathBuf, max_size: usize, roll_time: Duration, - deposits: Option, + file_upload: Option, auto_commit: bool, metric: &'static str, - shutdown_listener: triggered::Listener, } impl FileSinkBuilder { - pub fn new( - prefix: impl ToString, - target_path: &Path, - metric: &'static str, - shutdown_listener: triggered::Listener, - ) -> Self { + pub fn new(prefix: impl ToString, target_path: &Path, metric: &'static str) -> Self { Self { prefix: prefix.to_string(), target_path: target_path.to_path_buf(), tmp_path: target_path.join("tmp"), max_size: 50_000_000, roll_time: Duration::minutes(DEFAULT_SINK_ROLL_MINS), - deposits: None, + file_upload: None, auto_commit: true, metric, - shutdown_listener, } } @@ -106,8 +100,11 @@ impl FileSinkBuilder { } } - pub fn deposits(self, deposits: Option) -> Self { - Self { deposits, ..self } + pub fn file_upload(self, file_upload: Option) -> Self { + Self { + file_upload, + ..self + } } pub fn auto_commit(self, auto_commit: bool) -> Self { @@ -130,7 +127,6 @@ impl FileSinkBuilder { let client = FileSinkClient { sender: tx, metric: self.metric, - shutdown_listener: self.shutdown_listener.clone(), }; metrics::register_counter!(client.metric, vec![OK_LABEL]); @@ -140,13 +136,12 @@ impl FileSinkBuilder { tmp_path: self.tmp_path, prefix: self.prefix, max_size: self.max_size, - deposits: self.deposits, + file_upload: self.file_upload, roll_time: self.roll_time, messages: rx, staged_files: Vec::new(), auto_commit: self.auto_commit, active_sink: None, - shutdown_listener: self.shutdown_listener, }; sink.init().await?; Ok((client, sink)) @@ -157,7 +152,6 @@ impl FileSinkBuilder { pub struct FileSinkClient { sender: MessageSender, metric: &'static str, - shutdown_listener: triggered::Listener, } const OK_LABEL: Label = Label::from_static_parts("status", "ok"); @@ -175,9 +169,10 @@ impl FileSinkClient { let labels = labels.into_iter().map(Label::from); tokio::select! { - _ = self.shutdown_listener.clone() => { - Err(Error::Shutdown) - } + // TODO: check this again, do we need shutdown handling here? + // _ = self.shutdown_listener.clone() => { + // Err(Error::Shutdown) + // } result = self.sender.send_timeout(Message::Data(on_write_tx, bytes), SEND_TIMEOUT) => match result { Ok(_) => { metrics::increment_counter!( @@ -241,12 +236,11 @@ pub struct FileSink { roll_time: Duration, messages: MessageReceiver, - deposits: Option, + file_upload: Option, staged_files: Vec, auto_commit: bool, active_sink: Option, - shutdown_listener: triggered::Listener, } #[derive(Debug)] @@ -263,6 +257,21 @@ impl ActiveSink { } } +impl ManagedTask for FileSink { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + let handle = tokio::spawn(self.run(shutdown)); + + Box::pin( + handle + .map_err(anyhow::Error::from) + .and_then(|result| async move { result.map_err(anyhow::Error::from) }), + ) + } +} + impl FileSink { async fn init(&mut self) -> Result { fs::create_dir_all(&self.target_path).await?; @@ -289,7 +298,7 @@ impl FileSink { } // Notify all existing completed sinks - if let Some(deposits) = &self.deposits { + if let Some(file_upload) = &self.file_upload { let mut dir = fs::read_dir(&self.target_path).await?; loop { match dir.next_entry().await { @@ -299,7 +308,7 @@ impl FileSink { .to_string_lossy() .starts_with(&self.prefix) => { - file_upload::upload_file(deposits, &entry.path()).await?; + file_upload.upload_file(&entry.path()).await?; } Ok(None) => break, _ => continue, @@ -309,7 +318,7 @@ impl FileSink { Ok(()) } - pub async fn run(&mut self) -> Result { + pub async fn run(mut self, shutdown: triggered::Listener) -> Result { tracing::info!( "starting file sink {} in {}", self.prefix, @@ -325,7 +334,7 @@ impl FileSink { loop { tokio::select! { - _ = self.shutdown_listener.clone() => break, + _ = shutdown.clone() => break, _ = rollover_timer.tick() => self.maybe_roll().await?, msg = self.messages.recv() => match msg { Some(Message::Data(on_write_tx, bytes)) => { @@ -446,9 +455,9 @@ impl FileSink { let target_path = self.target_path.join(target_filename); fs::rename(&sink_path, &target_path).await?; - if let Some(deposits) = &self.deposits { - file_upload::upload_file(deposits, &target_path).await?; - } + if let Some(file_upload) = &self.file_upload { + file_upload.upload_file(&target_path).await?; + }; Ok(()) } @@ -500,145 +509,141 @@ fn file_name(path_buf: &Path) -> Result { }) } -#[cfg(test)] -mod tests { - use super::*; - use crate::{file_source, FileInfo, FileType}; - use futures::stream::StreamExt; - use std::str::FromStr; - use tempfile::TempDir; - use tokio::fs::DirEntry; - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn writes_a_framed_gzip_encoded_file() { - let tmp_dir = TempDir::new().expect("Unable to create temp dir"); - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - - let (file_sink_client, mut file_sink_server) = FileSinkBuilder::new( - FileType::EntropyReport, - tmp_dir.path(), - "fake_metric", - shutdown_listener.clone(), - ) - .roll_time(chrono::Duration::milliseconds(100)) - .create() - .await - .expect("failed to create file sink"); - - let sink_thread = tokio::spawn(async move { - file_sink_server - .run() - .await - .expect("failed to complete file sink"); - }); - - let (on_write_tx, _on_write_rx) = oneshot::channel(); - - file_sink_client - .sender - .try_send(Message::Data( - on_write_tx, - String::into_bytes("hello".to_string()), - )) - .expect("failed to send bytes to file sink"); - - tokio::time::sleep(time::Duration::from_millis(200)).await; - - shutdown_trigger.trigger(); - sink_thread.await.expect("file sink did not complete"); - - let entropy_file = get_entropy_file(&tmp_dir) - .await - .expect("no entropy available"); - assert_eq!("hello", read_file(&entropy_file).await); - } - - #[tokio::test] - async fn only_uploads_after_commit_when_auto_commit_is_false() { - let tmp_dir = TempDir::new().expect("Unable to create temp dir"); - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let (file_upload_tx, mut file_upload_rx) = file_upload::message_channel(); - - let (file_sink_client, mut file_sink_server) = FileSinkBuilder::new( - FileType::EntropyReport, - tmp_dir.path(), - "fake_metric", - shutdown_listener.clone(), - ) - .roll_time(chrono::Duration::milliseconds(100)) - .auto_commit(false) - .deposits(Some(file_upload_tx)) - .create() - .await - .expect("failed to create file sink"); - - let sink_thread = tokio::spawn(async move { - file_sink_server - .run() - .await - .expect("failed to complete file sink"); - }); - - let (on_write_tx, _on_write_rx) = oneshot::channel(); - file_sink_client - .sender - .try_send(Message::Data( - on_write_tx, - String::into_bytes("hello".to_string()), - )) - .expect("failed to send bytes to file sink"); - - tokio::time::sleep(time::Duration::from_millis(200)).await; - - assert!(get_entropy_file(&tmp_dir).await.is_err()); - assert_eq!( - Err(tokio::sync::mpsc::error::TryRecvError::Empty), - file_upload_rx.try_recv() - ); - - let receiver = file_sink_client.commit().await.expect("commit failed"); - let _ = receiver.await.expect("commit didn't complete completed"); - - assert!(file_upload_rx.try_recv().is_ok()); - - let entropy_file = get_entropy_file(&tmp_dir) - .await - .expect("no entropy available"); - assert_eq!("hello", read_file(&entropy_file).await); - - shutdown_trigger.trigger(); - sink_thread.await.expect("file sink did not complete"); - } - - async fn read_file(entry: &DirEntry) -> bytes::BytesMut { - file_source::source([entry.path()]) - .next() - .await - .unwrap() - .expect("invalid data in file") - } - - async fn get_entropy_file(tmp_dir: &TempDir) -> std::result::Result { - let mut entries = fs::read_dir(tmp_dir.path()) - .await - .expect("failed to read tmp dir"); - - while let Some(entry) = entries.next_entry().await.unwrap() { - if is_entropy_file(&entry) { - return Ok(entry); - } - } - - Err("no entropy available".to_string()) - } - - fn is_entropy_file(entry: &DirEntry) -> bool { - entry - .file_name() - .to_str() - .and_then(|file_name| FileInfo::from_str(file_name).ok()) - .map_or(false, |file_info| { - file_info.file_type == FileType::EntropyReport - }) - } -} +// #[cfg(test)] +// mod tests { +// use super::*; +// use crate::{file_source, FileInfo, FileType}; +// use futures::stream::StreamExt; +// use std::str::FromStr; +// use tempfile::TempDir; +// use tokio::fs::DirEntry; + +// #[tokio::test(flavor = "multi_thread", worker_threads = 2)] +// async fn writes_a_framed_gzip_encoded_file() { +// let tmp_dir = TempDir::new().expect("Unable to create temp dir"); +// let (shutdown_trigger, shutdown_listener) = triggered::trigger(); + +// let (file_sink_client, mut file_sink_server) = +// FileSinkBuilder::new(FileType::EntropyReport, tmp_dir.path(), "fake_metric") +// .roll_time(chrono::Duration::milliseconds(100)) +// .create() +// .await +// .expect("failed to create file sink"); + +// let sink_thread = tokio::spawn(async move { +// file_sink_server +// .run() +// .await +// .expect("failed to complete file sink"); +// }); + +// let (on_write_tx, _on_write_rx) = oneshot::channel(); + +// file_sink_client +// .sender +// .try_send(Message::Data( +// on_write_tx, +// String::into_bytes("hello".to_string()), +// )) +// .expect("failed to send bytes to file sink"); + +// tokio::time::sleep(time::Duration::from_millis(200)).await; + +// shutdown_trigger.trigger(); +// sink_thread.await.expect("file sink did not complete"); + +// let entropy_file = get_entropy_file(&tmp_dir) +// .await +// .expect("no entropy available"); +// assert_eq!("hello", read_file(&entropy_file).await); +// } + +// #[tokio::test] +// async fn only_uploads_after_commit_when_auto_commit_is_false() { +// let tmp_dir = TempDir::new().expect("Unable to create temp dir"); +// let (shutdown_trigger, shutdown_listener) = triggered::trigger(); +// let (file_upload_tx, mut file_upload_rx) = file_upload::message_channel(); + +// let (file_sink_client, mut file_sink_server) = FileSinkBuilder::new( +// FileType::EntropyReport, +// tmp_dir.path(), +// "fake_metric", +// shutdown_listener.clone(), +// ) +// .roll_time(chrono::Duration::milliseconds(100)) +// .auto_commit(false) +// .file_upload(Some(file_upload_tx)) +// .create() +// .await +// .expect("failed to create file sink"); + +// let sink_thread = tokio::spawn(async move { +// file_sink_server +// .run() +// .await +// .expect("failed to complete file sink"); +// }); + +// let (on_write_tx, _on_write_rx) = oneshot::channel(); +// file_sink_client +// .sender +// .try_send(Message::Data( +// on_write_tx, +// String::into_bytes("hello".to_string()), +// )) +// .expect("failed to send bytes to file sink"); + +// tokio::time::sleep(time::Duration::from_millis(200)).await; + +// assert!(get_entropy_file(&tmp_dir).await.is_err()); +// assert_eq!( +// Err(tokio::sync::mpsc::error::TryRecvError::Empty), +// file_upload_rx.try_recv() +// ); + +// let receiver = file_sink_client.commit().await.expect("commit failed"); +// let _ = receiver.await.expect("commit didn't complete completed"); + +// assert!(file_upload_rx.try_recv().is_ok()); + +// let entropy_file = get_entropy_file(&tmp_dir) +// .await +// .expect("no entropy available"); +// assert_eq!("hello", read_file(&entropy_file).await); + +// shutdown_trigger.trigger(); +// sink_thread.await.expect("file sink did not complete"); +// } + +// async fn read_file(entry: &DirEntry) -> bytes::BytesMut { +// file_source::source([entry.path()]) +// .next() +// .await +// .unwrap() +// .expect("invalid data in file") +// } + +// async fn get_entropy_file(tmp_dir: &TempDir) -> std::result::Result { +// let mut entries = fs::read_dir(tmp_dir.path()) +// .await +// .expect("failed to read tmp dir"); + +// while let Some(entry) = entries.next_entry().await.unwrap() { +// if is_entropy_file(&entry) { +// return Ok(entry); +// } +// } + +// Err("no entropy available".to_string()) +// } + +// fn is_entropy_file(entry: &DirEntry) -> bool { +// entry +// .file_name() +// .to_str() +// .and_then(|file_name| FileInfo::from_str(file_name).ok()) +// .map_or(false, |file_info| { +// file_info.file_type == FileType::EntropyReport +// }) +// } +// } diff --git a/file_store/src/file_source.rs b/file_store/src/file_source.rs index e1d75972d..59d7c7279 100644 --- a/file_store/src/file_source.rs +++ b/file_store/src/file_source.rs @@ -1,4 +1,4 @@ -use crate::{file_info_poller::FileInfoPollerBuilder, file_sink, BytesMutStream, Error}; +use crate::{file_info_poller::FileInfoPollerConfigBuilder, file_sink, BytesMutStream, Error}; use async_compression::tokio::bufread::GzipDecoder; use futures::{ stream::{self}, @@ -8,11 +8,11 @@ use std::path::{Path, PathBuf}; use tokio::{fs::File, io::BufReader}; use tokio_util::codec::{length_delimited::LengthDelimitedCodec, FramedRead}; -pub fn continuous_source() -> FileInfoPollerBuilder +pub fn continuous_source() -> FileInfoPollerConfigBuilder where T: Clone, { - FileInfoPollerBuilder::::default() + FileInfoPollerConfigBuilder::::default() } pub fn source(paths: I) -> BytesMutStream diff --git a/file_store/src/file_upload.rs b/file_store/src/file_upload.rs index 9d0cb64d3..fa1809b5d 100644 --- a/file_store/src/file_upload.rs +++ b/file_store/src/file_upload.rs @@ -1,36 +1,62 @@ use crate::{Error, FileStore, Result, Settings}; -use futures::StreamExt; +use futures::{future::LocalBoxFuture, StreamExt, TryFutureExt}; use std::{ path::{Path, PathBuf}, time::Duration, }; +use task_manager::ManagedTask; use tokio::{fs, sync::mpsc, time}; use tokio_stream::wrappers::UnboundedReceiverStream; pub type MessageSender = mpsc::UnboundedSender; pub type MessageReceiver = mpsc::UnboundedReceiver; -pub fn message_channel() -> (MessageSender, MessageReceiver) { - mpsc::unbounded_channel() -} - -pub async fn upload_file(tx: &MessageSender, file: &Path) -> Result { - tx.send(file.to_path_buf()).map_err(|_| Error::channel()) +#[derive(Debug, Clone)] +pub struct FileUpload { + sender: MessageSender, } -pub struct FileUpload { +pub struct FileUploadServer { messages: UnboundedReceiverStream, store: FileStore, } impl FileUpload { - pub async fn from_settings(settings: &Settings, messages: MessageReceiver) -> Result { - Ok(Self { - messages: UnboundedReceiverStream::new(messages), - store: FileStore::from_settings(settings).await?, - }) + pub async fn from_settings(settings: &Settings) -> Result<(Self, FileUploadServer)> { + let (sender, receiver) = mpsc::unbounded_channel(); + Ok(( + Self { sender }, + FileUploadServer { + messages: UnboundedReceiverStream::new(receiver), + store: FileStore::from_settings(settings).await?, + }, + )) + } + + pub async fn upload_file(&self, file: &Path) -> Result { + self.sender + .send(file.to_path_buf()) + .map_err(|_| Error::channel()) } - pub async fn run(self, shutdown: &triggered::Listener) -> Result { +} + +impl ManagedTask for FileUploadServer { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + let handle = tokio::spawn(self.run(shutdown)); + + Box::pin( + handle + .map_err(anyhow::Error::from) + .and_then(|result| async move { result.map_err(anyhow::Error::from) }), + ) + } +} + +impl FileUploadServer { + pub async fn run(self, shutdown: triggered::Listener) -> Result { tracing::info!("starting file uploader 1"); let uploads = self diff --git a/ingest/Cargo.toml b/ingest/Cargo.toml index 08b7b9389..d577e6efa 100644 --- a/ingest/Cargo.toml +++ b/ingest/Cargo.toml @@ -31,3 +31,5 @@ file-store = { path = "../file_store" } poc-metrics = { path = "../metrics" } metrics = {workspace = true } metrics-exporter-prometheus = { workspace = true } +task-manager = { path = "../task_manager" } +tokio-util = { workspace = true } diff --git a/ingest/src/main.rs b/ingest/src/main.rs index 446e16cf7..19419a0f6 100644 --- a/ingest/src/main.rs +++ b/ingest/src/main.rs @@ -2,7 +2,6 @@ use anyhow::Result; use clap::Parser; use ingest::{server_iot, server_mobile, Mode, Settings}; use std::path; -use tokio::{self, signal}; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; #[derive(Debug, clap::Parser)] @@ -50,19 +49,10 @@ impl Server { // Install the prometheus metrics exporter poc_metrics::start_metrics(&settings.metrics)?; - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); - // run the grpc server in either iot or mobile 5g mode match settings.mode { - Mode::Iot => server_iot::grpc_server(shutdown_listener, settings).await, - Mode::Mobile => server_mobile::grpc_server(shutdown_listener, settings).await, + Mode::Iot => server_iot::grpc_server(settings).await, + Mode::Mobile => server_mobile::grpc_server(settings).await, } } } diff --git a/ingest/src/server_iot.rs b/ingest/src/server_iot.rs index 3c124c30b..c1cfd1325 100644 --- a/ingest/src/server_iot.rs +++ b/ingest/src/server_iot.rs @@ -7,13 +7,15 @@ use file_store::{ traits::MsgVerify, FileType, }; +use futures::future::LocalBoxFuture; use futures_util::TryFutureExt; use helium_crypto::{Network, PublicKey}; use helium_proto::services::poc_lora::{ self, LoraBeaconIngestReportV1, LoraBeaconReportReqV1, LoraBeaconReportRespV1, LoraWitnessIngestReportV1, LoraWitnessReportReqV1, LoraWitnessReportRespV1, }; -use std::{convert::TryFrom, path::Path}; +use std::{convert::TryFrom, net::SocketAddr, path::Path}; +use task_manager::{ManagedTask, TaskManager}; use tonic::{transport, Request, Response, Status}; pub type GrpcResult = std::result::Result, Status>; @@ -23,6 +25,24 @@ pub struct GrpcServer { beacon_report_sink: FileSinkClient, witness_report_sink: FileSinkClient, required_network: Network, + address: SocketAddr, +} + +impl ManagedTask for GrpcServer { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + let address = self.address; + Box::pin(async move { + transport::Server::builder() + .layer(poc_metrics::request_layer!("ingest_server_iot_connection")) + .add_service(poc_lora::Server::new(*self)) + .serve_with_shutdown(address, shutdown) + .map_err(Error::from) + .await + }) + } } impl GrpcServer { @@ -30,11 +50,13 @@ impl GrpcServer { beacon_report_sink: FileSinkClient, witness_report_sink: FileSinkClient, required_network: Network, + address: SocketAddr, ) -> Result { Ok(Self { beacon_report_sink, witness_report_sink, required_network, + address, }) } @@ -108,58 +130,54 @@ impl poc_lora::PocLora for GrpcServer { } } -pub async fn grpc_server(shutdown: triggered::Listener, settings: &Settings) -> Result<()> { +pub async fn grpc_server(settings: &Settings) -> Result<()> { let grpc_addr = settings.listen_addr()?; // Initialize uploader - let (file_upload_tx, file_upload_rx) = file_upload::message_channel(); - let file_upload = - file_upload::FileUpload::from_settings(&settings.output, file_upload_rx).await?; + let (file_upload, file_upload_server) = + file_upload::FileUpload::from_settings(&settings.output).await?; let store_base_path = Path::new(&settings.cache); // iot beacon reports - let (beacon_report_sink, mut beacon_report_sink_server) = file_sink::FileSinkBuilder::new( + let (beacon_report_sink, beacon_report_sink_server) = file_sink::FileSinkBuilder::new( FileType::IotBeaconIngestReport, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_beacon_report"), - shutdown.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .roll_time(Duration::minutes(5)) .create() .await?; // iot witness reports - let (witness_report_sink, mut witness_report_sink_server) = file_sink::FileSinkBuilder::new( + let (witness_report_sink, witness_report_sink_server) = file_sink::FileSinkBuilder::new( FileType::IotWitnessIngestReport, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_witness_report"), - shutdown.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .roll_time(Duration::minutes(5)) .create() .await?; - let grpc_server = GrpcServer::new(beacon_report_sink, witness_report_sink, settings.network)?; + let grpc_server = GrpcServer::new( + beacon_report_sink, + witness_report_sink, + settings.network, + grpc_addr, + )?; tracing::info!( "grpc listening on {grpc_addr} and server mode {:?}", settings.mode ); - let server = transport::Server::builder() - .layer(poc_metrics::request_layer!("ingest_server_iot_connection")) - .add_service(poc_lora::Server::new(grpc_server)) - .serve_with_shutdown(grpc_addr, shutdown.clone()) - .map_err(Error::from); - - tokio::try_join!( - server, - beacon_report_sink_server.run().map_err(Error::from), - witness_report_sink_server.run().map_err(Error::from), - file_upload.run(&shutdown).map_err(Error::from), - ) - .map(|_| ()) + TaskManager::builder() + .add_task(file_upload_server) + .add_task(beacon_report_sink_server) + .add_task(witness_report_sink_server) + .add_task(grpc_server) + .start() + .await } diff --git a/ingest/src/server_mobile.rs b/ingest/src/server_mobile.rs index 7a527350e..4147db6d7 100644 --- a/ingest/src/server_mobile.rs +++ b/ingest/src/server_mobile.rs @@ -7,6 +7,7 @@ use file_store::{ traits::MsgVerify, FileType, }; +use futures::future::LocalBoxFuture; use futures_util::TryFutureExt; use helium_crypto::{Network, PublicKey}; use helium_proto::services::poc_mobile::{ @@ -16,8 +17,12 @@ use helium_proto::services::poc_mobile::{ SpeedtestIngestReportV1, SpeedtestReqV1, SpeedtestRespV1, SubscriberLocationIngestReportV1, SubscriberLocationReqV1, SubscriberLocationRespV1, }; -use std::path::Path; -use tonic::{metadata::MetadataValue, transport, Request, Response, Status}; +use std::{net::SocketAddr, path::Path}; +use task_manager::{ManagedTask, TaskManager}; +use tonic::{ + metadata::{Ascii, MetadataValue}, + transport, Request, Response, Status, +}; const INGEST_WAIT_DURATION_MINUTES: i64 = 15; @@ -31,9 +36,36 @@ pub struct GrpcServer { subscriber_location_report_sink: FileSinkClient, coverage_object_report_sink: FileSinkClient, required_network: Network, + address: SocketAddr, + api_token: MetadataValue, +} + +impl ManagedTask for GrpcServer { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + let api_token = self.api_token.clone(); + let address = self.address; + Box::pin(async move { + transport::Server::builder() + .layer(poc_metrics::request_layer!("ingest_server_grpc_connection")) + .add_service(poc_mobile::Server::with_interceptor( + *self, + move |req: Request<()>| match req.metadata().get("authorization") { + Some(t) if api_token == t => Ok(req), + _ => Err(Status::unauthenticated("No valid auth token")), + }, + )) + .serve_with_shutdown(address, shutdown) + .map_err(Error::from) + .await + }) + } } impl GrpcServer { + #[allow(clippy::too_many_arguments)] fn new( heartbeat_report_sink: FileSinkClient, speedtest_report_sink: FileSinkClient, @@ -41,6 +73,8 @@ impl GrpcServer { subscriber_location_report_sink: FileSinkClient, coverage_object_report_sink: FileSinkClient, required_network: Network, + address: SocketAddr, + api_token: MetadataValue, ) -> Result { Ok(Self { heartbeat_report_sink, @@ -49,6 +83,8 @@ impl GrpcServer { subscriber_location_report_sink, coverage_object_report_sink, required_network, + address, + api_token, }) } @@ -200,42 +236,37 @@ impl poc_mobile::PocMobile for GrpcServer { } } -pub async fn grpc_server(shutdown: triggered::Listener, settings: &Settings) -> Result<()> { +pub async fn grpc_server(settings: &Settings) -> Result<()> { let grpc_addr = settings.listen_addr()?; // Initialize uploader - let (file_upload_tx, file_upload_rx) = file_upload::message_channel(); - let file_upload = - file_upload::FileUpload::from_settings(&settings.output, file_upload_rx).await?; + let (file_upload, file_upload_server) = + file_upload::FileUpload::from_settings(&settings.output).await?; let store_base_path = Path::new(&settings.cache); - let (heartbeat_report_sink, mut heartbeat_report_sink_server) = - file_sink::FileSinkBuilder::new( - FileType::CellHeartbeatIngestReport, - store_base_path, - concat!(env!("CARGO_PKG_NAME"), "_heartbeat_report"), - shutdown.clone(), - ) - .deposits(Some(file_upload_tx.clone())) - .roll_time(Duration::minutes(INGEST_WAIT_DURATION_MINUTES)) - .create() - .await?; + let (heartbeat_report_sink, heartbeat_report_sink_server) = file_sink::FileSinkBuilder::new( + FileType::CellHeartbeatIngestReport, + store_base_path, + concat!(env!("CARGO_PKG_NAME"), "_heartbeat_report"), + ) + .file_upload(Some(file_upload.clone())) + .roll_time(Duration::minutes(INGEST_WAIT_DURATION_MINUTES)) + .create() + .await?; // speedtests - let (speedtest_report_sink, mut speedtest_report_sink_server) = - file_sink::FileSinkBuilder::new( - FileType::CellSpeedtestIngestReport, - store_base_path, - concat!(env!("CARGO_PKG_NAME"), "_speedtest_report"), - shutdown.clone(), - ) - .deposits(Some(file_upload_tx.clone())) - .roll_time(Duration::minutes(INGEST_WAIT_DURATION_MINUTES)) - .create() - .await?; + let (speedtest_report_sink, speedtest_report_sink_server) = file_sink::FileSinkBuilder::new( + FileType::CellSpeedtestIngestReport, + store_base_path, + concat!(env!("CARGO_PKG_NAME"), "_speedtest_report"), + ) + .file_upload(Some(file_upload.clone())) + .roll_time(Duration::minutes(INGEST_WAIT_DURATION_MINUTES)) + .create() + .await?; - let (data_transfer_session_sink, mut data_transfer_session_sink_server) = + let (data_transfer_session_sink, data_transfer_session_sink_server) = file_sink::FileSinkBuilder::new( FileType::DataTransferSessionIngestReport, store_base_path, @@ -243,46 +274,34 @@ pub async fn grpc_server(shutdown: triggered::Listener, settings: &Settings) -> env!("CARGO_PKG_NAME"), "_mobile_data_transfer_session_report" ), - shutdown.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .roll_time(Duration::minutes(INGEST_WAIT_DURATION_MINUTES)) .create() .await?; - let (subscriber_location_report_sink, mut subscriber_location_report_sink_server) = + let (subscriber_location_report_sink, subscriber_location_report_sink_server) = file_sink::FileSinkBuilder::new( FileType::SubscriberLocationIngestReport, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_subscriber_location_report"), - shutdown.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .roll_time(Duration::minutes(INGEST_WAIT_DURATION_MINUTES)) .create() .await?; - let (coverage_object_report_sink, mut coverage_object_report_sink_server) = + let (coverage_object_report_sink, coverage_object_report_sink_server) = file_sink::FileSinkBuilder::new( FileType::CoverageObjectIngestReport, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_coverage_object_report"), - shutdown.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .roll_time(Duration::minutes(INGEST_WAIT_DURATION_MINUTES)) .create() .await?; - let grpc_server = GrpcServer::new( - heartbeat_report_sink, - speedtest_report_sink, - data_transfer_session_sink, - subscriber_location_report_sink, - coverage_object_report_sink, - settings.network, - )?; - let Some(api_token) = settings .token .as_ref() @@ -294,37 +313,30 @@ pub async fn grpc_server(shutdown: triggered::Listener, settings: &Settings) -> bail!("expected valid api token in settings"); }; + let grpc_server = GrpcServer::new( + heartbeat_report_sink, + speedtest_report_sink, + data_transfer_session_sink, + subscriber_location_report_sink, + coverage_object_report_sink, + settings.network, + grpc_addr, + api_token, + )?; + tracing::info!( "grpc listening on {grpc_addr} and server mode {:?}", settings.mode ); - //TODO start a service with either the poc mobile or poc iot endpoints only - not both - // use _server_mode (set above ) to decide - let server = transport::Server::builder() - .layer(poc_metrics::request_layer!("ingest_server_grpc_connection")) - .add_service(poc_mobile::Server::with_interceptor( - grpc_server, - move |req: Request<()>| match req.metadata().get("authorization") { - Some(t) if api_token == t => Ok(req), - _ => Err(Status::unauthenticated("No valid auth token")), - }, - )) - .serve_with_shutdown(grpc_addr, shutdown.clone()) - .map_err(Error::from); - - tokio::try_join!( - server, - heartbeat_report_sink_server.run().map_err(Error::from), - speedtest_report_sink_server.run().map_err(Error::from), - data_transfer_session_sink_server.run().map_err(Error::from), - subscriber_location_report_sink_server - .run() - .map_err(Error::from), - coverage_object_report_sink_server - .run() - .map_err(Error::from), - file_upload.run(&shutdown).map_err(Error::from), - ) - .map(|_| ()) + TaskManager::builder() + .add_task(file_upload_server) + .add_task(heartbeat_report_sink_server) + .add_task(speedtest_report_sink_server) + .add_task(data_transfer_session_sink_server) + .add_task(subscriber_location_report_sink_server) + .add_task(coverage_object_report_sink_server) + .add_task(grpc_server) + .start() + .await } diff --git a/iot_config/Cargo.toml b/iot_config/Cargo.toml index e3dcd24cc..b7b47a814 100644 --- a/iot_config/Cargo.toml +++ b/iot_config/Cargo.toml @@ -39,3 +39,5 @@ tonic = {workspace = true} tracing = {workspace = true} tracing-subscriber = {workspace = true} triggered = {workspace = true} +task-manager = { path = "../task_manager" } +tokio-util = { workspace = true } diff --git a/iot_config/src/gateway_service.rs b/iot_config/src/gateway_service.rs index 162d01e9f..d4362bd43 100644 --- a/iot_config/src/gateway_service.rs +++ b/iot_config/src/gateway_service.rs @@ -35,7 +35,6 @@ pub struct GatewayService { region_map: RegionMapReader, signing_key: Arc, delegate_cache: watch::Receiver, - shutdown: triggered::Listener, } impl GatewayService { @@ -45,7 +44,6 @@ impl GatewayService { region_map: RegionMapReader, auth_cache: AuthCache, delegate_cache: watch::Receiver, - shutdown: triggered::Listener, ) -> Result { let gateway_cache = Arc::new(Cache::new()); let cache_clone = gateway_cache.clone(); @@ -58,7 +56,6 @@ impl GatewayService { region_map, signing_key: Arc::new(settings.signing_keypair()?), delegate_cache, - shutdown, }) } @@ -278,13 +275,11 @@ impl iot_config::Gateway for GatewayService { let signing_key = self.signing_key.clone(); let batch_size = request.batch_size; let region_map = self.region_map.clone(); - let shutdown_listener = self.shutdown.clone(); let (tx, rx) = tokio::sync::mpsc::channel(20); tokio::spawn(async move { tokio::select! { - _ = shutdown_listener => (), _ = stream_all_gateways_info( &pool, tx.clone(), diff --git a/iot_config/src/main.rs b/iot_config/src/main.rs index 9a69b77e5..f6cf5b250 100644 --- a/iot_config/src/main.rs +++ b/iot_config/src/main.rs @@ -1,5 +1,6 @@ use anyhow::{Error, Result}; use clap::Parser; +use futures::future::LocalBoxFuture; use futures_util::TryFutureExt; use helium_proto::services::iot_config::{AdminServer, GatewayServer, OrgServer, RouteServer}; use iot_config::{ @@ -7,8 +8,8 @@ use iot_config::{ org_service::OrgService, region_map::RegionMapReader, route_service::RouteService, settings::Settings, telemetry, }; -use std::{path::PathBuf, time::Duration}; -use tokio::signal; +use std::{net::SocketAddr, path::PathBuf, time::Duration}; +use task_manager::{ManagedTask, TaskManager}; use tonic::transport; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; @@ -60,28 +61,12 @@ impl Daemon { poc_metrics::start_metrics(&settings.metrics)?; telemetry::initialize(); - // Configure shutdown trigger - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); - // Create database pool - let (pool, db_join_handle) = settings - .database - .connect("iot-config-store", shutdown_listener.clone()) - .await?; + let pool = settings.database.connect("iot-config-store").await?; sqlx::migrate!().run(&pool).await?; // Create on-chain metadata pool - let (metadata_pool, md_pool_handle) = settings - .metadata - .connect("iot-config-metadata", shutdown_listener.clone()) - .await?; + let metadata_pool = settings.metadata.connect("iot-config-metadata").await?; let listen_addr = settings.listen_addr()?; @@ -95,21 +80,14 @@ impl Daemon { region_map.clone(), auth_cache.clone(), delegate_key_cache, - shutdown_listener.clone(), - )?; - let route_svc = RouteService::new( - settings, - auth_cache.clone(), - pool.clone(), - shutdown_listener.clone(), )?; + let route_svc = RouteService::new(settings, auth_cache.clone(), pool.clone())?; let org_svc = OrgService::new( settings, auth_cache.clone(), pool.clone(), route_svc.clone_update_channel(), delegate_key_updater, - shutdown_listener.clone(), )?; let admin_svc = AdminService::new( settings, @@ -126,23 +104,43 @@ impl Daemon { tracing::debug!("listening on {listen_addr}"); tracing::debug!("signing as {pubkey}"); - let server = transport::Server::builder() - .http2_keepalive_interval(Some(Duration::from_secs(250))) - .http2_keepalive_timeout(Some(Duration::from_secs(60))) - .add_service(GatewayServer::new(gateway_svc)) - .add_service(OrgServer::new(org_svc)) - .add_service(RouteServer::new(route_svc)) - .add_service(AdminServer::new(admin_svc)) - .serve_with_shutdown(listen_addr, shutdown_listener) - .map_err(Error::from); - - tokio::try_join!( - db_join_handle.map_err(Error::from), - md_pool_handle.map_err(Error::from), - server - )?; + let grpc_server = GrpcServer { + listen_addr, + gateway_svc, + route_svc, + org_svc, + admin_svc, + }; + + TaskManager::builder().add_task(grpc_server).start().await + } +} + +pub struct GrpcServer { + listen_addr: SocketAddr, + gateway_svc: GatewayService, + route_svc: RouteService, + org_svc: OrgService, + admin_svc: AdminService, +} - Ok(()) +impl ManagedTask for GrpcServer { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(async move { + transport::Server::builder() + .http2_keepalive_interval(Some(Duration::from_secs(250))) + .http2_keepalive_timeout(Some(Duration::from_secs(60))) + .add_service(GatewayServer::new(self.gateway_svc)) + .add_service(OrgServer::new(self.org_svc)) + .add_service(RouteServer::new(self.route_svc)) + .add_service(AdminServer::new(self.admin_svc)) + .serve_with_shutdown(self.listen_addr, shutdown) + .map_err(Error::from) + .await + }) } } diff --git a/iot_config/src/org_service.rs b/iot_config/src/org_service.rs index 19e9561cf..326500dd2 100644 --- a/iot_config/src/org_service.rs +++ b/iot_config/src/org_service.rs @@ -26,7 +26,6 @@ pub struct OrgService { route_update_tx: broadcast::Sender, signing_key: Keypair, delegate_updater: watch::Sender, - shutdown: triggered::Listener, } #[derive(Clone, Debug, PartialEq)] @@ -42,7 +41,6 @@ impl OrgService { pool: Pool, route_update_tx: broadcast::Sender, delegate_updater: watch::Sender, - shutdown: triggered::Listener, ) -> Result { Ok(Self { auth_cache, @@ -50,7 +48,6 @@ impl OrgService { route_update_tx, signing_key: settings.signing_keypair()?, delegate_updater, - shutdown, }) } @@ -467,7 +464,6 @@ impl iot_config::Org for OrgService { })?; tokio::select! { - _ = self.shutdown.clone() => return Err(Status::unavailable("service shutting down")), result = self.stream_org_routes_enable_disable(request.oui) => result? } } @@ -506,7 +502,6 @@ impl iot_config::Org for OrgService { })?; tokio::select! { - _ = self.shutdown.clone() => return Err(Status::unavailable("service shutting down")), result = self.stream_org_routes_enable_disable(request.oui) => result? } } diff --git a/iot_config/src/route_service.rs b/iot_config/src/route_service.rs index 3ede2c637..f8b6cb5fd 100644 --- a/iot_config/src/route_service.rs +++ b/iot_config/src/route_service.rs @@ -37,7 +37,6 @@ pub struct RouteService { auth_cache: AuthCache, pool: Pool, update_channel: broadcast::Sender, - shutdown: triggered::Listener, signing_key: Arc, } @@ -48,17 +47,11 @@ enum OrgId<'a> { } impl RouteService { - pub fn new( - settings: &Settings, - auth_cache: AuthCache, - pool: Pool, - shutdown: triggered::Listener, - ) -> Result { + pub fn new(settings: &Settings, auth_cache: AuthCache, pool: Pool) -> Result { Ok(Self { auth_cache, pool, update_channel: update_channel(), - shutdown, signing_key: Arc::new(settings.signing_keypair()?), }) } @@ -369,7 +362,6 @@ impl iot_config::Route for RouteService { tracing::info!("client subscribed to route stream"); let pool = self.pool.clone(); - let shutdown_listener = self.shutdown.clone(); let (tx, rx) = tokio::sync::mpsc::channel(20); let signing_key = self.signing_key.clone(); @@ -377,7 +369,6 @@ impl iot_config::Route for RouteService { tokio::spawn(async move { tokio::select! { - _ = shutdown_listener.clone() => return, result = stream_existing_routes(&pool, &signing_key, tx.clone()) .and_then(|_| stream_existing_euis(&pool, &signing_key, tx.clone())) .and_then(|_| stream_existing_devaddrs(&pool, &signing_key, tx.clone())) @@ -389,13 +380,7 @@ impl iot_config::Route for RouteService { tracing::info!("existing routes sent; streaming updates as available"); telemetry::route_stream_subscribe(); loop { - let shutdown = shutdown_listener.clone(); - tokio::select! { - _ = shutdown => { - telemetry::route_stream_unsubscribe(); - return - } msg = route_updates.recv() => if let Ok(update) = msg { if tx.send(Ok(update)).await.is_err() { telemetry::route_stream_unsubscribe(); @@ -423,7 +408,6 @@ impl iot_config::Route for RouteService { let pool = self.pool.clone(); let (tx, rx) = tokio::sync::mpsc::channel(20); - let shutdown_listener = self.shutdown.clone(); tracing::debug!(route_id = request.route_id, "listing eui pairs"); @@ -448,9 +432,6 @@ impl iot_config::Route for RouteService { }; tokio::select! { - _ = shutdown_listener => { - _ = tx.send(Err(Status::unavailable("service shutting down"))).await; - } _ = async { while let Some(eui) = eui_stream.next().await { let message = match eui { @@ -497,7 +478,6 @@ impl iot_config::Route for RouteService { .await?; tokio::select! { - _ = self.shutdown.clone() => return Err(Status::unavailable("service shutting down")), result = incoming_stream .map_ok(|update| match validator.validate_update(&update) { Ok(()) => Ok(update), @@ -580,8 +560,6 @@ impl iot_config::Route for RouteService { let (tx, rx) = tokio::sync::mpsc::channel(20); let pool = self.pool.clone(); - let shutdown_listener = self.shutdown.clone(); - tracing::debug!(route_id = request.route_id, "listing devaddr ranges"); tokio::spawn(async move { @@ -603,9 +581,6 @@ impl iot_config::Route for RouteService { }; tokio::select! { - _ = shutdown_listener => { - _ = tx.send(Err(Status::unavailable("service shutting down"))).await; - } _ = async { while let Some(devaddr) = devaddrs.next().await { let message = match devaddr { @@ -655,7 +630,6 @@ impl iot_config::Route for RouteService { .await?; tokio::select! { - _ = self.shutdown.clone() => return Err(Status::unavailable("service shutting down")), result = incoming_stream .map_ok(|update| match validator.validate_update(&update) { Ok(()) => Ok(update), @@ -744,7 +718,6 @@ impl iot_config::Route for RouteService { let pool = self.pool.clone(); let (tx, rx) = tokio::sync::mpsc::channel(20); - let shutdown_listener = self.shutdown.clone(); tracing::debug!( route_id = request.route_id, @@ -772,9 +745,6 @@ impl iot_config::Route for RouteService { }; tokio::select! { - _ = shutdown_listener => { - _ = tx.send(Err(Status::unavailable("service shutting down"))).await; - } _ = async { while let Some(skf) = skf_stream.next().await { let message = match skf { @@ -806,7 +776,6 @@ impl iot_config::Route for RouteService { let pool = self.pool.clone(); let (tx, rx) = tokio::sync::mpsc::channel(20); - let shutdown_listener = self.shutdown.clone(); tracing::debug!( route_id = request.route_id, @@ -838,9 +807,6 @@ impl iot_config::Route for RouteService { }; tokio::select! { - _ = shutdown_listener => { - _ = tx.send(Err(Status::unavailable("service shutting down"))).await; - } _ = async { while let Some(skf) = skf_stream.next().await { let message = match skf { diff --git a/iot_packet_verifier/Cargo.toml b/iot_packet_verifier/Cargo.toml index 07c203c6a..337295121 100644 --- a/iot_packet_verifier/Cargo.toml +++ b/iot_packet_verifier/Cargo.toml @@ -27,9 +27,11 @@ sqlx = {workspace = true} solana = {path = "../solana"} thiserror = {workspace = true} tokio = {workspace = true} +tokio-util = { workspace = true } tonic = {workspace = true} tracing = {workspace = true} tracing-subscriber = {workspace = true} triggered = {workspace = true} http = {workspace = true} http-serde = {workspace = true} +task-manager = { path = "../task_manager" } diff --git a/iot_packet_verifier/src/burner.rs b/iot_packet_verifier/src/burner.rs index dd73256b4..6b83db8e9 100644 --- a/iot_packet_verifier/src/burner.rs +++ b/iot_packet_verifier/src/burner.rs @@ -2,8 +2,10 @@ use crate::{ balances::{BalanceCache, BalanceStore}, pending_burns::{Burn, PendingBurns}, }; +use futures::{future::LocalBoxFuture, TryFutureExt}; use solana::SolanaNetwork; use std::time::Duration; +use task_manager::ManagedTask; use tokio::task; pub struct Burner { @@ -23,6 +25,19 @@ pub enum BurnError { SolanaError(S), } +impl ManagedTask for Burner +where + P: PendingBurns + Send + Sync + 'static, + S: SolanaNetwork, +{ + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown).map_err(anyhow::Error::from)) + } +} + impl Burner { pub fn new(pending_burns: P, balances: &BalanceCache, burn_period: u64, solana: S) -> Self { Self { @@ -41,7 +56,7 @@ where { pub async fn run( mut self, - shutdown: &triggered::Listener, + shutdown: triggered::Listener, ) -> Result<(), BurnError> { let burn_service = task::spawn(async move { loop { diff --git a/iot_packet_verifier/src/daemon.rs b/iot_packet_verifier/src/daemon.rs index 194b4fca4..b2f96a1df 100644 --- a/iot_packet_verifier/src/daemon.rs +++ b/iot_packet_verifier/src/daemon.rs @@ -4,7 +4,7 @@ use crate::{ settings::Settings, verifier::{ConfigServer, Verifier}, }; -use anyhow::{bail, Error, Result}; +use anyhow::{bail, Result}; use file_store::{ file_info_poller::{FileInfoStream, LookbackBehavior}, file_sink::FileSinkClient, @@ -14,13 +14,11 @@ use file_store::{ }; use futures_util::TryFutureExt; use iot_config::client::OrgClient; -use solana::SolanaRpc; +use solana::{balance_monitor::BalanceMonitor, SolanaRpc}; use sqlx::{Pool, Postgres}; use std::{sync::Arc, time::Duration}; -use tokio::{ - signal, - sync::{mpsc::Receiver, Mutex}, -}; +use task_manager::{ManagedTask, TaskManager}; +use tokio::sync::{mpsc::Receiver, Mutex}; struct Daemon { pool: Pool, @@ -31,8 +29,18 @@ struct Daemon { minimum_allowed_balance: u64, } +impl ManagedTask for Daemon { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> futures::future::LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } +} + impl Daemon { - pub async fn run(mut self, shutdown: &triggered::Listener) -> Result<()> { + pub async fn run(mut self, shutdown: triggered::Listener) -> Result<()> { + tracing::info!("starting daemon"); loop { tokio::select! { _ = shutdown.clone() => break, @@ -46,7 +54,7 @@ impl Daemon { } } - + tracing::info!("stopping daemon"); Ok(()) } @@ -80,23 +88,11 @@ impl Daemon { pub struct Cmd {} impl Cmd { - pub async fn run(self, settings: &Settings) -> Result<()> { + pub async fn run(self, settings: Settings) -> Result<()> { poc_metrics::start_metrics(&settings.metrics)?; - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); - // Set up the postgres pool: - let (mut pool, db_handle) = settings - .database - .connect(env!("CARGO_PKG_NAME"), shutdown_listener.clone()) - .await?; + let pool = settings.database.connect(env!("CARGO_PKG_NAME")).await?; sqlx::migrate!().run(&pool).await?; let solana = if settings.enable_solana_integration { @@ -109,15 +105,10 @@ impl Cmd { None }; - let sol_balance_monitor = solana::balance_monitor::start( - env!("CARGO_PKG_NAME"), - solana.clone(), - shutdown_listener.clone(), - ) - .await?; + let sol_balance_monitor = BalanceMonitor::new(env!("CARGO_PKG_NAME"), solana.clone())?; // Set up the balance cache: - let balances = BalanceCache::new(&mut pool, solana.clone()).await?; + let balances = BalanceCache::new(&mut pool.clone(), solana.clone()).await?; // Set up the balance burner: let burner = Burner::new( @@ -127,31 +118,28 @@ impl Cmd { solana.clone(), ); - let (file_upload_tx, file_upload_rx) = file_upload::message_channel(); - let file_upload = - file_upload::FileUpload::from_settings(&settings.output, file_upload_rx).await?; + let (file_upload, file_upload_server) = + file_upload::FileUpload::from_settings(&settings.output).await?; let store_base_path = std::path::Path::new(&settings.cache); // Verified packets: - let (valid_packets, mut valid_packets_server) = FileSinkBuilder::new( + let (valid_packets, valid_packets_server) = FileSinkBuilder::new( FileType::IotValidPacket, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_valid_packets"), - shutdown_listener.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(false) .create() .await?; - let (invalid_packets, mut invalid_packets_server) = FileSinkBuilder::new( + let (invalid_packets, invalid_packets_server) = FileSinkBuilder::new( FileType::InvalidPacket, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_invalid_packets"), - shutdown_listener.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(false) .create() .await?; @@ -162,15 +150,13 @@ impl Cmd { let file_store = FileStore::from_settings(&settings.ingest).await?; - let (report_files, source_join_handle) = + let (report_files, report_files_server) = file_source::continuous_source::() .db(pool.clone()) .store(file_store) .lookback(LookbackBehavior::StartAfter(settings.start_after())) .file_type(FileType::IotPacketReport) - .build()? - .start(shutdown_listener.clone()) - .await?; + .create()?; let balance_store = balances.balances(); let verifier_daemon = Daemon { @@ -184,28 +170,29 @@ impl Cmd { }, minimum_allowed_balance: settings.minimum_allowed_balance, }; - - // Run the services: - tokio::try_join!( - db_handle.map_err(Error::from), - burner.run(&shutdown_listener).map_err(Error::from), - file_upload.run(&shutdown_listener).map_err(Error::from), - verifier_daemon.run(&shutdown_listener).map_err(Error::from), - valid_packets_server.run().map_err(Error::from), - invalid_packets_server.run().map_err(Error::from), - org_client - .monitor_funds( - solana, - balance_store, - settings.minimum_allowed_balance, - Duration::from_secs(60 * settings.monitor_funds_period), - shutdown_listener.clone(), - ) - .map_err(Error::from), - source_join_handle.map_err(Error::from), - sol_balance_monitor.map_err(Error::from), - )?; - - Ok(()) + let minimum_allowed_balance = settings.minimum_allowed_balance; + let monitor_funds_period = settings.monitor_funds_period; + + TaskManager::builder() + .add_task(file_upload_server) + .add_task(valid_packets_server) + .add_task(invalid_packets_server) + .add_task(report_files_server) + .add_task(move |shutdown| { + org_client + .monitor_funds( + solana, + balance_store, + minimum_allowed_balance, + Duration::from_secs(60 * monitor_funds_period), + shutdown, + ) + .map_err(anyhow::Error::from) + }) + .add_task(burner) + .add_task(verifier_daemon) + .add_task(sol_balance_monitor) + .start() + .await } } diff --git a/iot_packet_verifier/src/main.rs b/iot_packet_verifier/src/main.rs index 0084aaff8..08ed5a814 100644 --- a/iot_packet_verifier/src/main.rs +++ b/iot_packet_verifier/src/main.rs @@ -37,7 +37,7 @@ pub enum Cmd { impl Cmd { async fn run(self, settings: Settings) -> Result<()> { match self { - Self::Server(cmd) => cmd.run(&settings).await, + Self::Server(cmd) => cmd.run(settings).await, } } } diff --git a/iot_verifier/Cargo.toml b/iot_verifier/Cargo.toml index 6b96b75da..302020f5b 100644 --- a/iot_verifier/Cargo.toml +++ b/iot_verifier/Cargo.toml @@ -53,3 +53,5 @@ itertools = {workspace = true} rand = {workspace = true} beacon = {workspace = true} price = { path = "../price" } +tokio-util = { workspace = true } +task-manager = { path = "../task_manager" } diff --git a/iot_verifier/src/entropy_loader.rs b/iot_verifier/src/entropy_loader.rs index 04398418f..a138d7610 100644 --- a/iot_verifier/src/entropy_loader.rs +++ b/iot_verifier/src/entropy_loader.rs @@ -1,12 +1,14 @@ use crate::entropy::Entropy; use blake3::hash; use file_store::{entropy_report::EntropyReport, file_info_poller::FileInfoStream}; -use futures::{StreamExt, TryStreamExt}; +use futures::{future::LocalBoxFuture, StreamExt, TryStreamExt}; use sqlx::PgPool; +use task_manager::ManagedTask; use tokio::sync::mpsc::Receiver; pub struct EntropyLoader { pub pool: PgPool, + pub file_receiver: Receiver>, } #[derive(thiserror::Error, Debug)] @@ -17,24 +19,30 @@ pub enum NewLoaderError { DbStoreError(#[from] db_store::Error), } +impl ManagedTask for EntropyLoader { + fn start_task( + self: Box, + shutdown_listener: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown_listener)) + } +} + impl EntropyLoader { - pub async fn run( - &mut self, - mut receiver: Receiver>, - shutdown: &triggered::Listener, - ) -> anyhow::Result<()> { + pub async fn run(mut self, shutdown: triggered::Listener) -> anyhow::Result<()> { + tracing::info!("starting entropy_loader"); loop { if shutdown.is_triggered() { break; } tokio::select! { _ = shutdown.clone() => break, - msg = receiver.recv() => if let Some(stream) = msg { + msg = self.file_receiver.recv() => if let Some(stream) = msg { self.handle_report(stream).await?; } } } - tracing::info!("stopping verifier entropy_loader"); + tracing::info!("stopping entropy_loader"); Ok(()) } diff --git a/iot_verifier/src/gateway_cache.rs b/iot_verifier/src/gateway_cache.rs index fea190b32..4b76018e8 100644 --- a/iot_verifier/src/gateway_cache.rs +++ b/iot_verifier/src/gateway_cache.rs @@ -2,6 +2,7 @@ use crate::gateway_updater::MessageReceiver; use helium_crypto::PublicKeyBinary; use iot_config::gateway_info::GatewayInfo; +#[derive(Clone)] pub struct GatewayCache { gateway_cache_receiver: MessageReceiver, } diff --git a/iot_verifier/src/gateway_updater.rs b/iot_verifier/src/gateway_updater.rs index 955c5a699..5140b50b2 100644 --- a/iot_verifier/src/gateway_updater.rs +++ b/iot_verifier/src/gateway_updater.rs @@ -1,12 +1,15 @@ use crate::Settings; use chrono::Duration; +use futures::future::LocalBoxFuture; use futures::stream::StreamExt; use helium_crypto::PublicKeyBinary; use iot_config::{ client::{Client as IotConfigClient, ClientError as IotConfigClientError}, gateway_info::{GatewayInfo, GatewayInfoResolver}, }; + use std::collections::HashMap; +use task_manager::ManagedTask; use tokio::sync::watch; use tokio::time; @@ -28,6 +31,15 @@ pub enum GatewayUpdaterError { SendError(#[from] watch::error::SendError), } +impl ManagedTask for GatewayUpdater { + fn start_task( + self: Box, + shutdown_listener: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown_listener)) + } +} + impl GatewayUpdater { pub async fn from_settings( settings: &Settings, @@ -45,7 +57,7 @@ impl GatewayUpdater { )) } - pub async fn run(mut self, shutdown: &triggered::Listener) -> Result<(), GatewayUpdaterError> { + pub async fn run(mut self, shutdown: triggered::Listener) -> anyhow::Result<()> { tracing::info!("starting gateway_updater"); let mut trigger_timer = time::interval( @@ -56,15 +68,16 @@ impl GatewayUpdater { loop { if shutdown.is_triggered() { - tracing::info!("stopping gateway_updater"); - return Ok(()); - } + break; + }; tokio::select! { _ = trigger_timer.tick() => self.handle_refresh_tick().await?, _ = shutdown.clone() => return Ok(()), } } + tracing::info!("stopping gateway_updater"); + Ok(()) } async fn handle_refresh_tick(&mut self) -> Result<(), GatewayUpdaterError> { diff --git a/iot_verifier/src/hex_density.rs b/iot_verifier/src/hex_density.rs index b0ba46ef8..1565eef60 100644 --- a/iot_verifier/src/hex_density.rs +++ b/iot_verifier/src/hex_density.rs @@ -64,28 +64,31 @@ lazy_static! { }; } -#[async_trait::async_trait] -pub trait HexDensityMap: Clone { - async fn get(&self, hex: u64) -> Option; - async fn swap(&self, new_map: HashMap); -} +// #[async_trait::async_trait] +// pub trait HexDensityMap: Clone { +// async fn get(&self, hex: u64) -> Option; +// async fn swap(&self, new_map: HashMap); +// } #[derive(Debug, Clone)] -pub struct SharedHexDensityMap(Arc>>); +pub struct HexDensityMap(Arc>>); -impl SharedHexDensityMap { +impl Default for HexDensityMap { + fn default() -> Self { + Self::new() + } +} + +impl HexDensityMap { pub fn new() -> Self { Self(Arc::new(RwLock::new(HashMap::new()))) } -} -#[async_trait::async_trait] -impl HexDensityMap for SharedHexDensityMap { - async fn get(&self, hex: u64) -> Option { + pub async fn get(&self, hex: u64) -> Option { self.0.read().await.get(&hex).cloned() } - async fn swap(&self, new_map: HashMap) { + pub async fn swap(&self, new_map: HashMap) { *self.0.write().await = new_map; } } diff --git a/iot_verifier/src/loader.rs b/iot_verifier/src/loader.rs index 42f1a43ca..e31385c44 100644 --- a/iot_verifier/src/loader.rs +++ b/iot_verifier/src/loader.rs @@ -14,10 +14,11 @@ use file_store::{ traits::{IngestId, MsgDecode}, FileInfo, FileStore, FileType, }; -use futures::{stream, StreamExt}; +use futures::{future::LocalBoxFuture, stream, StreamExt}; use helium_crypto::PublicKeyBinary; use sqlx::PgPool; use std::{hash::Hasher, ops::DerefMut, time::Duration}; +use task_manager::ManagedTask; use tokio::{ sync::Mutex, time::{self, MissedTickBehavior}, @@ -37,6 +38,7 @@ pub struct Loader { deny_list_latest_url: String, deny_list_trigger_interval: Duration, deny_list: DenyList, + gateway_cache: GatewayCache, } #[derive(thiserror::Error, Debug)] @@ -55,8 +57,21 @@ pub enum ValidGatewayResult { Unknown, } +impl ManagedTask for Loader { + fn start_task( + self: Box, + shutdown_listener: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown_listener)) + } +} + impl Loader { - pub async fn from_settings(settings: &Settings, pool: PgPool) -> Result { + pub async fn from_settings( + settings: &Settings, + pool: PgPool, + gateway_cache: GatewayCache, + ) -> Result { tracing::info!("from_settings verifier loader"); let ingest_store = FileStore::from_settings(&settings.ingest).await?; let poll_time = settings.poc_loader_poll_time(); @@ -74,15 +89,12 @@ impl Loader { deny_list_latest_url: settings.denylist.denylist_url.clone(), deny_list_trigger_interval: settings.denylist.trigger_interval(), deny_list, + gateway_cache, }) } - pub async fn run( - &mut self, - shutdown: &triggered::Listener, - gateway_cache: &GatewayCache, - ) -> anyhow::Result<()> { - tracing::info!("started verifier loader"); + pub async fn run(mut self, shutdown: triggered::Listener) -> anyhow::Result<()> { + tracing::info!("started loader"); let mut report_timer = time::interval(self.poll_time); report_timer.set_missed_tick_behavior(MissedTickBehavior::Skip); let mut denylist_timer = time::interval(self.deny_list_trigger_interval); @@ -100,7 +112,7 @@ impl Loader { tracing::error!("fatal loader error, denylist_tick triggered: {err:?}"); } }, - _ = report_timer.tick() => match self.handle_report_tick(gateway_cache).await { + _ = report_timer.tick() => match self.handle_report_tick().await { Ok(()) => (), Err(err) => { tracing::error!("loader error, report_tick triggered: {err:?}"); @@ -108,7 +120,7 @@ impl Loader { } } } - tracing::info!("stopping verifier loader"); + tracing::info!("stopping loader"); Ok(()) } @@ -129,7 +141,7 @@ impl Loader { Ok(()) } - async fn handle_report_tick(&self, gateway_cache: &GatewayCache) -> anyhow::Result<()> { + async fn handle_report_tick(&self) -> anyhow::Result<()> { tracing::info!("handling report tick"); let now = Utc::now(); // the loader loads files from s3 via a sliding window @@ -161,7 +173,7 @@ impl Loader { tracing::info!("current window width insufficient. completed handling poc_report tick"); return Ok(()); } - self.process_window(gateway_cache, after, before).await?; + self.process_window(after, before).await?; Meta::update_last_timestamp(&self.pool, REPORTS_META_NAME, Some(before)).await?; Report::pending_beacons_to_ready(&self.pool, now).await?; tracing::info!("completed handling poc_report tick"); @@ -170,7 +182,6 @@ impl Loader { async fn process_window( &self, - gateway_cache: &GatewayCache, after: DateTime, before: DateTime, ) -> anyhow::Result<()> { @@ -187,7 +198,6 @@ impl Loader { .process_events( FileType::IotBeaconIngestReport, &self.ingest_store, - gateway_cache, after, before, Some(&xor_data), @@ -226,7 +236,6 @@ impl Loader { .process_events( FileType::IotWitnessIngestReport, &self.ingest_store, - gateway_cache, after - self.ingestor_rollup_time, before + self.ingestor_rollup_time, None, @@ -248,7 +257,6 @@ impl Loader { &self, file_type: FileType, store: &FileStore, - gateway_cache: &GatewayCache, after: chrono::DateTime, before: chrono::DateTime, xor_data: Option<&Mutex>>, @@ -267,13 +275,7 @@ impl Loader { stream::iter(infos) .for_each_concurrent(10, |file_info| async move { match self - .process_file( - store, - file_info.clone(), - gateway_cache, - xor_data, - xor_filter, - ) + .process_file(store, file_info.clone(), xor_data, xor_filter) .await { Ok(()) => tracing::debug!( @@ -297,7 +299,6 @@ impl Loader { &self, store: &FileStore, file_info: FileInfo, - gateway_cache: &GatewayCache, xor_data: Option<&Mutex>>, xor_filter: Option<&Xor16>, ) -> anyhow::Result<()> { @@ -316,7 +317,7 @@ impl Loader { tracing::warn!("skipping report of type {file_type} due to error {err:?}") } Ok(buf) => { - match self.handle_report(file_type, &buf, gateway_cache, xor_data, xor_filter, &metrics).await + match self.handle_report(file_type, &buf, xor_data, xor_filter, &metrics).await { Ok(Some(bindings)) => inserts.push(bindings), Ok(None) => (), @@ -344,7 +345,6 @@ impl Loader { &self, file_type: FileType, buf: &[u8], - gateway_cache: &GatewayCache, xor_data: Option<&Mutex>>, xor_filter: Option<&Xor16>, metrics: &LoaderMetricTracker, @@ -354,10 +354,7 @@ impl Loader { let beacon = IotBeaconIngestReport::decode(buf)?; tracing::debug!("beacon report from ingestor: {:?}", &beacon); let packet_data = beacon.report.data.clone(); - match self - .check_valid_gateway(&beacon.report.pub_key, gateway_cache) - .await - { + match self.check_valid_gateway(&beacon.report.pub_key).await { ValidGatewayResult::Valid => { let res = InsertBindings { id: beacon.ingest_id(), @@ -392,34 +389,29 @@ impl Loader { let packet_data = witness.report.data.clone(); if let Some(filter) = xor_filter { match verify_witness_packet_data(&packet_data, filter) { - true => { - match self - .check_valid_gateway(&witness.report.pub_key, gateway_cache) - .await - { - ValidGatewayResult::Valid => { - let res = InsertBindings { - id: witness.ingest_id(), - remote_entropy: Vec::::with_capacity(0), - packet_data, - buf: buf.to_vec(), - received_ts: witness.received_timestamp, - report_type: ReportType::Witness, - status: IotStatus::Ready, - }; - metrics.increment_witnesses(); - Ok(Some(res)) - } - ValidGatewayResult::Denied => { - metrics.increment_witnesses_denied(); - Ok(None) - } - ValidGatewayResult::Unknown => { - metrics.increment_witnesses_unknown(); - Ok(None) - } + true => match self.check_valid_gateway(&witness.report.pub_key).await { + ValidGatewayResult::Valid => { + let res = InsertBindings { + id: witness.ingest_id(), + remote_entropy: Vec::::with_capacity(0), + packet_data, + buf: buf.to_vec(), + received_ts: witness.received_timestamp, + report_type: ReportType::Witness, + status: IotStatus::Ready, + }; + metrics.increment_witnesses(); + Ok(Some(res)) } - } + ValidGatewayResult::Denied => { + metrics.increment_witnesses_denied(); + Ok(None) + } + ValidGatewayResult::Unknown => { + metrics.increment_witnesses_unknown(); + Ok(None) + } + }, false => { tracing::debug!( "dropping witness report as no associated beacon data: {:?}", @@ -440,28 +432,23 @@ impl Loader { } } - async fn check_valid_gateway( - &self, - pub_key: &PublicKeyBinary, - gateway_cache: &GatewayCache, - ) -> ValidGatewayResult { + async fn check_valid_gateway(&self, pub_key: &PublicKeyBinary) -> ValidGatewayResult { if self.check_gw_denied(pub_key).await { tracing::debug!("dropping denied gateway : {:?}", &pub_key); return ValidGatewayResult::Denied; } - if self.check_unknown_gw(pub_key, gateway_cache).await { + if self.check_unknown_gw(pub_key).await { tracing::debug!("dropping unknown gateway: {:?}", &pub_key); return ValidGatewayResult::Unknown; } ValidGatewayResult::Valid } - async fn check_unknown_gw( - &self, - pub_key: &PublicKeyBinary, - gateway_cache: &GatewayCache, - ) -> bool { - gateway_cache.resolve_gateway_info(pub_key).await.is_err() + async fn check_unknown_gw(&self, pub_key: &PublicKeyBinary) -> bool { + self.gateway_cache + .resolve_gateway_info(pub_key) + .await + .is_err() } async fn check_gw_denied(&self, pub_key: &PublicKeyBinary) -> bool { diff --git a/iot_verifier/src/main.rs b/iot_verifier/src/main.rs index a80f820ef..c5d751c16 100644 --- a/iot_verifier/src/main.rs +++ b/iot_verifier/src/main.rs @@ -1,11 +1,11 @@ use crate::entropy_loader::EntropyLoader; -use anyhow::{Error, Result}; +use anyhow::Result; +use chrono::Duration as ChronoDuration; use clap::Parser; use file_store::{ entropy_report::EntropyReport, file_info_poller::LookbackBehavior, file_sink, file_source, file_upload, iot_packet::IotValidPacket, FileStore, FileType, }; -use futures::TryFutureExt; use iot_config::client::Client as IotConfigClient; use iot_verifier::{ entropy_loader, gateway_cache::GatewayCache, gateway_updater::GatewayUpdater, loader, @@ -14,6 +14,7 @@ use iot_verifier::{ }; use price::PriceTracker; use std::path; +use task_manager::TaskManager; use tokio::signal; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; @@ -65,61 +66,80 @@ impl Server { // Install the prometheus metrics exporter poc_metrics::start_metrics(&settings.metrics)?; - // configure shutdown trigger - let (shutdown_trigger, shutdown) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); - // Create database pool and run migrations - let (pool, db_join_handle) = settings - .database - .connect(env!("CARGO_PKG_NAME"), shutdown.clone()) - .await?; + let pool = settings.database.connect(env!("CARGO_PKG_NAME")).await?; sqlx::migrate!().run(&pool).await?; telemetry::initialize(&pool).await?; let iot_config_client = IotConfigClient::from_settings(&settings.iot_config_client)?; - let (gateway_updater_receiver, gateway_updater) = + let (gateway_updater_receiver, gateway_updater_server) = GatewayUpdater::from_settings(settings, iot_config_client.clone()).await?; + let gateway_cache = GatewayCache::new(gateway_updater_receiver.clone()); let region_cache = RegionCache::from_settings(settings, iot_config_client.clone())?; - let (file_upload_tx, file_upload_rx) = file_upload::message_channel(); - let file_upload = - file_upload::FileUpload::from_settings(&settings.output, file_upload_rx).await?; + let (file_upload, file_upload_server) = + file_upload::FileUpload::from_settings(&settings.output).await?; let store_base_path = std::path::Path::new(&settings.cache); + + // * + // setup the price tracker requirements + // * + // todo: update price tracker to not require shutdown listener to be passed in + let (shutdown_trigger, shutdown) = triggered::trigger(); + let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; + tokio::spawn(async move { + tokio::select! { + _ = sigterm.recv() => shutdown_trigger.trigger(), + _ = signal::ctrl_c() => shutdown_trigger.trigger(), + } + }); + let (price_tracker, _price_receiver) = + PriceTracker::start(&settings.price_tracker, shutdown.clone()).await?; + + // * + // setup the loader requirements + // * + let loader = + loader::Loader::from_settings(settings, pool.clone(), gateway_cache.clone()).await?; + + // * + // setup the density scaler requirements + // * + let density_scaler = + DensityScaler::from_settings(settings, pool.clone(), gateway_updater_receiver.clone()) + .await?; + + // * + // setup the rewarder requirements + // * + // Gateway reward shares sink - let (rewards_sink, mut gateway_rewards_server) = file_sink::FileSinkBuilder::new( + let (rewards_sink, gateway_rewards_sink_server) = file_sink::FileSinkBuilder::new( FileType::IotRewardShare, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_gateway_reward_shares"), - shutdown.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(false) .create() .await?; // Reward manifest - let (reward_manifests_sink, mut reward_manifests_server) = file_sink::FileSinkBuilder::new( - FileType::RewardManifest, - store_base_path, - concat!(env!("CARGO_PKG_NAME"), "_iot_reward_manifest"), - shutdown.clone(), - ) - .deposits(Some(file_upload_tx.clone())) - .auto_commit(false) - .create() - .await?; + let (reward_manifests_sink, reward_manifests_sink_server) = + file_sink::FileSinkBuilder::new( + FileType::RewardManifest, + store_base_path, + concat!(env!("CARGO_PKG_NAME"), "_iot_reward_manifest"), + ) + .file_upload(Some(file_upload.clone())) + .auto_commit(false) + .create() + .await?; let rewarder = Rewarder { pool: pool.clone(), @@ -127,14 +147,16 @@ impl Server { reward_manifests_sink, reward_period_hours: settings.rewards, reward_offset: settings.reward_offset_duration(), + price_tracker, }; - // setup the entropy loader continious source + // * + // setup entropy requirements + // * let max_lookback_age = settings.loader_window_max_lookback_age(); - let mut entropy_loader = EntropyLoader { pool: pool.clone() }; let entropy_store = FileStore::from_settings(&settings.entropy).await?; let entropy_interval = settings.entropy_interval(); - let (entropy_loader_receiver, entropy_loader_source_join_handle) = + let (entropy_loader_receiver, entropy_loader_server) = file_source::continuous_source::() .db(pool.clone()) .store(entropy_store.clone()) @@ -142,15 +164,31 @@ impl Server { .lookback(LookbackBehavior::Max(max_lookback_age)) .poll_duration(entropy_interval) .offset(entropy_interval * 2) - .build()? - .start(shutdown.clone()) - .await?; + .create()?; + + let entropy_loader = EntropyLoader { + pool: pool.clone(), + file_receiver: entropy_loader_receiver, + }; + + // * + // setup the packet loader requirements + // * + + let (non_rewardable_packet_sink, non_rewardable_packet_sink_server) = + file_sink::FileSinkBuilder::new( + FileType::NonRewardablePacket, + store_base_path, + concat!(env!("CARGO_PKG_NAME"), "_non_rewardable_packet"), + ) + .file_upload(Some(file_upload.clone())) + .roll_time(ChronoDuration::minutes(5)) + .create() + .await?; - // setup the packet loader continious source - let packet_loader = packet_loader::PacketLoader::from_settings(settings, pool.clone()); let packet_store = FileStore::from_settings(&settings.packet_ingest).await?; let packet_interval = settings.packet_interval(); - let (pk_loader_receiver, pk_loader_source_join_handle) = + let (pk_loader_receiver, pk_loader_server) = file_source::continuous_source::() .db(pool.clone()) .store(packet_store.clone()) @@ -158,48 +196,120 @@ impl Server { .lookback(LookbackBehavior::Max(max_lookback_age)) .poll_duration(packet_interval) .offset(packet_interval * 2) - .build()? - .start(shutdown.clone()) - .await?; + .create()?; - // init da processes - let mut loader = loader::Loader::from_settings(settings, pool.clone()).await?; - let mut runner = runner::Runner::from_settings(settings, pool.clone()).await?; - let purger = purger::Purger::from_settings(settings, pool.clone()).await?; - let mut density_scaler = - DensityScaler::from_settings(settings, pool, gateway_updater_receiver.clone()).await?; - let (price_tracker, price_receiver) = - PriceTracker::start(&settings.price_tracker, shutdown.clone()).await?; + let packet_loader = packet_loader::PacketLoader::from_settings( + settings, + pool.clone(), + gateway_cache.clone(), + pk_loader_receiver, + non_rewardable_packet_sink, + ); + + // * + // setup the purger requirements + // * + + let (purger_invalid_beacon_sink, purger_invalid_beacon_sink_server) = + file_sink::FileSinkBuilder::new( + FileType::IotInvalidBeaconReport, + store_base_path, + concat!(env!("CARGO_PKG_NAME"), "_invalid_beacon"), + ) + .file_upload(Some(file_upload.clone())) + .auto_commit(false) + .create() + .await?; + + let (purger_invalid_witness_sink, purger_invalid_witness_sink_server) = + file_sink::FileSinkBuilder::new( + FileType::IotInvalidWitnessReport, + store_base_path, + concat!(env!("CARGO_PKG_NAME"), "_invalid_witness_report"), + ) + .file_upload(Some(file_upload.clone())) + .auto_commit(false) + .create() + .await?; - tokio::try_join!( - db_join_handle.map_err(Error::from), - gateway_updater.run(&shutdown).map_err(Error::from), - gateway_rewards_server.run().map_err(Error::from), - reward_manifests_server.run().map_err(Error::from), - file_upload.run(&shutdown).map_err(Error::from), - runner.run( - file_upload_tx.clone(), - &gateway_cache, - ®ion_cache, - density_scaler.hex_density_map(), - &shutdown - ), - entropy_loader.run(entropy_loader_receiver, &shutdown), - loader.run(&shutdown, &gateway_cache), - packet_loader.run( - pk_loader_receiver, - &shutdown, - &gateway_cache, - file_upload_tx.clone() - ), - purger.run(&shutdown), - rewarder.run(price_tracker, &shutdown), - density_scaler.run(&shutdown).map_err(Error::from), - price_receiver.map_err(Error::from), - entropy_loader_source_join_handle.map_err(anyhow::Error::from), - pk_loader_source_join_handle.map_err(anyhow::Error::from), + let purger = purger::Purger::from_settings( + settings, + pool.clone(), + purger_invalid_beacon_sink, + purger_invalid_witness_sink, ) - .map(|_| ()) + .await?; + + // * + // setup the runner requirements + // * + + let (runner_invalid_beacon_sink, runner_invalid_beacon_sink_server) = + file_sink::FileSinkBuilder::new( + FileType::IotInvalidBeaconReport, + store_base_path, + concat!(env!("CARGO_PKG_NAME"), "_invalid_beacon_report"), + ) + .file_upload(Some(file_upload.clone())) + .roll_time(ChronoDuration::minutes(5)) + .create() + .await?; + + let (runner_invalid_witness_sink, runner_invalid_witness_sink_server) = + file_sink::FileSinkBuilder::new( + FileType::IotInvalidWitnessReport, + store_base_path, + concat!(env!("CARGO_PKG_NAME"), "_invalid_witness_report"), + ) + .file_upload(Some(file_upload.clone())) + .roll_time(ChronoDuration::minutes(5)) + .create() + .await?; + + let (runner_poc_sink, runner_poc_sink_server) = file_sink::FileSinkBuilder::new( + FileType::IotPoc, + store_base_path, + concat!(env!("CARGO_PKG_NAME"), "_valid_poc"), + ) + .file_upload(Some(file_upload.clone())) + .roll_time(ChronoDuration::minutes(2)) + .create() + .await?; + + let runner = runner::Runner::from_settings( + settings, + pool.clone(), + gateway_cache.clone(), + region_cache.clone(), + runner_invalid_beacon_sink, + runner_invalid_witness_sink, + runner_poc_sink, + density_scaler.hex_density_map.clone(), + ) + .await?; + + TaskManager::builder() + .add_task(file_upload_server) + .add_task(gateway_rewards_sink_server) + .add_task(reward_manifests_sink_server) + .add_task(non_rewardable_packet_sink_server) + .add_task(pk_loader_server) + .add_task(purger_invalid_beacon_sink_server) + .add_task(purger_invalid_witness_sink_server) + .add_task(runner_invalid_beacon_sink_server) + .add_task(runner_invalid_witness_sink_server) + .add_task(runner_poc_sink_server) + .add_task(density_scaler) + .add_task(gateway_updater_server) + .add_task(entropy_loader_server) + .add_task(entropy_loader) + .add_task(packet_loader) + .add_task(loader) + .add_task(runner) + .add_task(rewarder) + .add_task(purger) + .start() + .await } } diff --git a/iot_verifier/src/packet_loader.rs b/iot_verifier/src/packet_loader.rs index c0a30fe28..eb634c719 100644 --- a/iot_verifier/src/packet_loader.rs +++ b/iot_verifier/src/packet_loader.rs @@ -2,21 +2,21 @@ use crate::{ gateway_cache::GatewayCache, reward_share::GatewayDCShare, telemetry::LoaderMetricTracker, Settings, }; -use chrono::{Duration as ChronoDuration, Utc}; -use file_store::{ - file_info_poller::FileInfoStream, file_sink, file_sink::FileSinkClient, - file_upload::MessageSender as FileUploadSender, iot_packet::IotValidPacket, FileType, -}; -use futures::{StreamExt, TryStreamExt}; +use chrono::Utc; +use file_store::{file_info_poller::FileInfoStream, file_sink, iot_packet::IotValidPacket}; +use futures::{future::LocalBoxFuture, StreamExt, TryStreamExt}; use helium_proto::services::packet_verifier::ValidPacket; use helium_proto::services::poc_lora::{NonRewardablePacket, NonRewardablePacketReason}; use sqlx::PgPool; -use std::path::Path; +use task_manager::ManagedTask; use tokio::sync::mpsc::Receiver; pub struct PacketLoader { pub pool: PgPool, pub cache: String, + gateway_cache: GatewayCache, + file_receiver: Receiver>, + file_sink: file_sink::FileSinkClient, } #[derive(thiserror::Error, Debug)] @@ -27,34 +27,49 @@ pub enum NewLoaderError { DbStoreError(#[from] db_store::Error), } +impl ManagedTask for PacketLoader { + fn start_task( + self: Box, + shutdown_listener: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown_listener)) + } +} + impl PacketLoader { - pub fn from_settings(settings: &Settings, pool: PgPool) -> Self { + pub fn from_settings( + settings: &Settings, + pool: PgPool, + gateway_cache: GatewayCache, + file_receiver: Receiver>, + file_sink: file_sink::FileSinkClient, + ) -> Self { tracing::info!("from_settings packet loader"); let cache = settings.cache.clone(); - Self { pool, cache } + Self { + pool, + cache, + gateway_cache, + file_receiver, + file_sink, + } } - pub async fn run( - &self, - mut receiver: Receiver>, - shutdown: &triggered::Listener, - gateway_cache: &GatewayCache, - file_upload_tx: FileUploadSender, - ) -> anyhow::Result<()> { - tracing::info!("starting verifier iot packet loader"); - let store_base_path = Path::new(&self.cache); - let (non_rewardable_packet_sink, mut non_rewardable_packet_server) = - file_sink::FileSinkBuilder::new( - FileType::NonRewardablePacket, - store_base_path, - concat!(env!("CARGO_PKG_NAME"), "_non_rewardable_packet"), - shutdown.clone(), - ) - .deposits(Some(file_upload_tx.clone())) - .roll_time(ChronoDuration::minutes(5)) - .create() - .await?; - tokio::spawn(async move { non_rewardable_packet_server.run().await }); + pub async fn run(mut self, shutdown: triggered::Listener) -> anyhow::Result<()> { + tracing::info!("starting iot packet loader"); + // let store_base_path = Path::new(&self.cache); + // let (non_rewardable_packet_sink, mut non_rewardable_packet_server) = + // file_sink::FileSinkBuilder::new( + // FileType::NonRewardablePacket, + // store_base_path, + // concat!(env!("CARGO_PKG_NAME"), "_non_rewardable_packet"), + // shutdown.clone(), + // ) + // .deposits(Some(file_upload_tx.clone())) + // .roll_time(ChronoDuration::minutes(5)) + // .create() + // .await?; + // tokio::spawn(async move { non_rewardable_packet_server.run().await }); loop { if shutdown.is_triggered() { @@ -62,14 +77,12 @@ impl PacketLoader { } tokio::select! { _ = shutdown.clone() => break, - msg = receiver.recv() => if let Some(stream) = msg { + msg = self.file_receiver.recv() => if let Some(stream) = msg { let metrics = LoaderMetricTracker::new(); - match self.handle_packet_file(stream, gateway_cache, &non_rewardable_packet_sink, &metrics).await { + match self.handle_packet_file(stream, &metrics).await { Ok(()) => { - // todo: maybe two actions below can occur in handle_packet - // but wasnt able to get it to work ? metrics.record_metrics(); - non_rewardable_packet_sink.commit().await?; + self.file_sink.commit().await?; }, Err(err) => { return Err(err)} @@ -77,15 +90,13 @@ impl PacketLoader { } } } - tracing::info!("stopping verifier iot packet loader"); + tracing::info!("stopping iot packet loader"); Ok(()) } async fn handle_packet_file( &self, file_info_stream: FileInfoStream, - gateway_cache: &GatewayCache, - non_rewardable_packet_sink: &FileSinkClient, metrics: &LoaderMetricTracker, ) -> anyhow::Result<()> { let mut transaction = self.pool.begin().await?; @@ -103,7 +114,8 @@ impl PacketLoader { .try_fold( transaction, |mut transaction, (valid_packet, reward_share)| async move { - if gateway_cache + if self + .gateway_cache .resolve_gateway_info(&reward_share.hotspot_key) .await .is_ok() @@ -120,7 +132,7 @@ impl PacketLoader { reason: reason as i32, timestamp, }; - non_rewardable_packet_sink + self.file_sink .write( non_rewardable_packet_proto, &[("reason", reason.as_str_name())], diff --git a/iot_verifier/src/poc.rs b/iot_verifier/src/poc.rs index c823bf51a..22f6555a9 100644 --- a/iot_verifier/src/poc.rs +++ b/iot_verifier/src/poc.rs @@ -101,9 +101,9 @@ impl Poc { pub async fn verify_beacon( &mut self, - hex_density_map: impl HexDensityMap, - gateway_cache: &GatewayCache, - region_cache: &RegionCache, + hex_density_map: HexDensityMap, + gateway_cache: GatewayCache, + region_cache: RegionCache, pool: &PgPool, beacon_interval: Duration, beacon_interval_tolerance: Duration, @@ -161,8 +161,8 @@ impl Poc { pub async fn verify_witnesses( &mut self, beacon_info: &GatewayInfo, - hex_density_map: impl HexDensityMap, - gateway_cache: &GatewayCache, + hex_density_map: HexDensityMap, + gateway_cache: GatewayCache, ) -> Result { let mut verified_witnesses: Vec = Vec::new(); let mut failed_witnesses: Vec = Vec::new(); @@ -178,8 +178,8 @@ impl Poc { .verify_witness( &witness_report, beacon_info, - gateway_cache, - &hex_density_map, + gateway_cache.clone(), + hex_density_map.clone(), ) .await { @@ -216,8 +216,8 @@ impl Poc { &mut self, witness_report: &IotWitnessIngestReport, beaconer_info: &GatewayInfo, - gateway_cache: &GatewayCache, - hex_density_map: &impl HexDensityMap, + gateway_cache: GatewayCache, + hex_density_map: HexDensityMap, ) -> Result { let witness = &witness_report.report; let witness_pub_key = witness.pub_key.clone(); @@ -277,6 +277,7 @@ impl Poc { ) { Ok(()) => { let tx_scale = hex_density_map + .clone() .get(beaconer_metadata.location) .await .unwrap_or(*DEFAULT_TX_SCALE); diff --git a/iot_verifier/src/purger.rs b/iot_verifier/src/purger.rs index b565fbef4..c6a169815 100644 --- a/iot_verifier/src/purger.rs +++ b/iot_verifier/src/purger.rs @@ -1,22 +1,24 @@ use crate::{entropy::Entropy, poc_report::Report, telemetry, Settings}; use chrono::Duration; use file_store::{ - file_sink::{self, FileSinkClient}, - file_upload, + file_sink::FileSinkClient, iot_beacon_report::IotBeaconIngestReport, iot_invalid_poc::IotInvalidBeaconReport, iot_invalid_poc::IotInvalidWitnessReport, iot_witness_report::IotWitnessIngestReport, traits::{IngestId, MsgDecode}, - FileType, }; -use futures::stream::{self, StreamExt}; +use futures::{ + future::LocalBoxFuture, + stream::{self, StreamExt}, +}; use helium_proto::services::poc_lora::{ InvalidParticipantSide, InvalidReason, LoraInvalidBeaconReportV1, LoraInvalidWitnessReportV1, }; use lazy_static::lazy_static; use sqlx::{PgPool, Postgres}; -use std::{ops::DerefMut, path::Path}; +use std::ops::DerefMut; +use task_manager::ManagedTask; use tokio::{ sync::Mutex, time::{self, MissedTickBehavior}, @@ -36,68 +38,46 @@ lazy_static! { pub struct Purger { pool: PgPool, - cache: String, - output: file_store::Settings, base_stale_period: Duration, + invalid_beacon_sink: FileSinkClient, + invalid_witness_sink: FileSinkClient, } #[derive(thiserror::Error, Debug)] #[error("error creating purger: {0}")] pub struct NewPurgerError(#[from] db_store::Error); +impl ManagedTask for Purger { + fn start_task( + self: Box, + shutdown_listener: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown_listener)) + } +} + impl Purger { - pub async fn from_settings(settings: &Settings, pool: PgPool) -> Result { - let cache = settings.cache.clone(); - let output = settings.output.clone(); + pub async fn from_settings( + settings: &Settings, + pool: PgPool, + invalid_beacon_sink: FileSinkClient, + invalid_witness_sink: FileSinkClient, + ) -> Result { let base_stale_period = settings.base_stale_period(); Ok(Self { pool, - cache, - output, base_stale_period, + invalid_beacon_sink, + invalid_witness_sink, }) } - pub async fn run(&self, shutdown: &triggered::Listener) -> anyhow::Result<()> { + pub async fn run(self, shutdown: triggered::Listener) -> anyhow::Result<()> { tracing::info!("starting purger"); let mut db_timer = time::interval(DB_POLL_TIME); db_timer.set_missed_tick_behavior(MissedTickBehavior::Skip); - let store_base_path = Path::new(&self.cache); - let (file_upload_tx, file_upload_rx) = file_upload::message_channel(); - let file_upload = - file_upload::FileUpload::from_settings(&self.output, file_upload_rx).await?; - - let (invalid_beacon_sink, mut invalid_beacon_sink_server) = - file_sink::FileSinkBuilder::new( - FileType::IotInvalidBeaconReport, - store_base_path, - concat!(env!("CARGO_PKG_NAME"), "_invalid_beacon"), - shutdown.clone(), - ) - .deposits(Some(file_upload_tx.clone())) - .auto_commit(false) - .create() - .await?; - - let (invalid_witness_sink, mut invalid_witness_sink_server) = - file_sink::FileSinkBuilder::new( - FileType::IotInvalidWitnessReport, - store_base_path, - concat!(env!("CARGO_PKG_NAME"), "_invalid_witness_report"), - shutdown.clone(), - ) - .deposits(Some(file_upload_tx.clone())) - .auto_commit(false) - .create() - .await?; - - let upload_shutdown = shutdown.clone(); - tokio::spawn(async move { invalid_beacon_sink_server.run().await }); - tokio::spawn(async move { invalid_witness_sink_server.run().await }); - tokio::spawn(async move { file_upload.run(&upload_shutdown).await }); - loop { if shutdown.is_triggered() { break; @@ -105,7 +85,7 @@ impl Purger { tokio::select! { _ = shutdown.clone() => break, _ = db_timer.tick() => - match self.handle_db_tick(&invalid_beacon_sink, &invalid_witness_sink).await { + match self.handle_db_tick().await { Ok(()) => (), Err(err) => { tracing::error!("fatal purger error: {err:?}"); @@ -117,11 +97,7 @@ impl Purger { Ok(()) } - async fn handle_db_tick( - &self, - invalid_beacon_sink: &FileSinkClient, - invalid_witness_sink: &FileSinkClient, - ) -> anyhow::Result<()> { + async fn handle_db_tick(&self) -> anyhow::Result<()> { // pull stale beacons and witnesses // for each we have to write out an invalid report to S3 // as these wont have previously resulted in a file going to s3 @@ -137,10 +113,7 @@ impl Purger { let tx = Mutex::new(self.pool.begin().await?); stream::iter(stale_beacons) .for_each_concurrent(PURGER_WORKERS, |report| async { - match self - .handle_purged_beacon(&tx, report, invalid_beacon_sink) - .await - { + match self.handle_purged_beacon(&tx, report).await { Ok(()) => (), Err(err) => { tracing::warn!("failed to purge beacon: {err:?}") @@ -148,7 +121,7 @@ impl Purger { } }) .await; - invalid_beacon_sink.commit().await?; + self.invalid_beacon_sink.commit().await?; tx.into_inner().commit().await?; let witness_stale_period = self.base_stale_period + *WITNESS_STALE_PERIOD; @@ -163,10 +136,7 @@ impl Purger { let tx = Mutex::new(self.pool.begin().await?); stream::iter(stale_witnesses) .for_each_concurrent(PURGER_WORKERS, |report| async { - match self - .handle_purged_witness(&tx, report, invalid_witness_sink) - .await - { + match self.handle_purged_witness(&tx, report).await { Ok(()) => (), Err(err) => { tracing::warn!("failed to purge witness: {err:?}") @@ -174,7 +144,7 @@ impl Purger { } }) .await; - invalid_witness_sink.commit().await?; + self.invalid_witness_sink.commit().await?; tx.into_inner().commit().await?; tracing::info!("completed purging {num_stale_witnesses} stale witnesses"); @@ -187,7 +157,6 @@ impl Purger { &self, tx: &Mutex>, db_beacon: Report, - invalid_beacon_sink: &FileSinkClient, ) -> anyhow::Result<()> { let beacon_buf: &[u8] = &db_beacon.report_data; let beacon_report = IotBeaconIngestReport::decode(beacon_buf)?; @@ -201,7 +170,7 @@ impl Purger { } .into(); - invalid_beacon_sink + self.invalid_beacon_sink .write( invalid_beacon_proto, &[("reason", InvalidReason::Stale.as_str_name())], @@ -217,7 +186,6 @@ impl Purger { &self, tx: &Mutex>, db_witness: Report, - invalid_witness_sink: &FileSinkClient, ) -> anyhow::Result<()> { let witness_buf: &[u8] = &db_witness.report_data; let witness_report = IotWitnessIngestReport::decode(witness_buf)?; @@ -231,7 +199,7 @@ impl Purger { } .into(); - invalid_witness_sink + self.invalid_witness_sink .write( invalid_witness_report_proto, &[("reason", InvalidReason::Stale.as_str_name())], diff --git a/iot_verifier/src/region_cache.rs b/iot_verifier/src/region_cache.rs index 021102a17..0ca18ee29 100644 --- a/iot_verifier/src/region_cache.rs +++ b/iot_verifier/src/region_cache.rs @@ -10,6 +10,7 @@ use std::{sync::Arc, time::Duration}; /// how often to evict expired items from the cache ( every 5 mins) const CACHE_EVICTION_FREQUENCY: Duration = Duration::from_secs(60 * 5); +#[derive(Clone)] pub struct RegionCache { pub iot_config_client: IotConfigClient, pub cache: Arc>, diff --git a/iot_verifier/src/rewarder.rs b/iot_verifier/src/rewarder.rs index 80cd6cb92..ca9ed1641 100644 --- a/iot_verifier/src/rewarder.rs +++ b/iot_verifier/src/rewarder.rs @@ -5,12 +5,14 @@ use crate::{ use chrono::{DateTime, Duration, TimeZone, Utc}; use db_store::meta; use file_store::{file_sink, traits::TimestampEncode}; +use futures::future::LocalBoxFuture; use helium_proto::RewardManifest; use price::PriceTracker; use reward_scheduler::Scheduler; use rust_decimal::prelude::*; -use sqlx::{PgExecutor, Pool, Postgres}; +use sqlx::{PgExecutor, PgPool, Pool, Postgres}; use std::ops::Range; +use task_manager::ManagedTask; use tokio::time::sleep; const REWARDS_NOT_CURRENT_DELAY_PERIOD: i64 = 5; @@ -21,15 +23,39 @@ pub struct Rewarder { pub reward_manifests_sink: file_sink::FileSinkClient, pub reward_period_hours: i64, pub reward_offset: Duration, + pub price_tracker: PriceTracker, +} + +impl ManagedTask for Rewarder { + fn start_task( + self: Box, + shutdown_listener: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown_listener)) + } } impl Rewarder { - pub async fn run( - mut self, + pub async fn new( + pool: PgPool, + rewards_sink: file_sink::FileSinkClient, + reward_manifests_sink: file_sink::FileSinkClient, + reward_period_hours: i64, + reward_offset: Duration, price_tracker: PriceTracker, - shutdown: &triggered::Listener, - ) -> anyhow::Result<()> { - tracing::info!("Starting iot verifier rewarder"); + ) -> Self { + Self { + pool, + rewards_sink, + reward_manifests_sink, + reward_period_hours, + reward_offset, + price_tracker, + } + } + + pub async fn run(mut self, shutdown: triggered::Listener) -> anyhow::Result<()> { + tracing::info!("Starting rewarder"); let reward_period_length = Duration::hours(self.reward_period_hours); @@ -44,7 +70,8 @@ impl Rewarder { ); let sleep_duration = if scheduler.should_reward(now) { - let iot_price = price_tracker + let iot_price = self + .price_tracker .price(&helium_proto::BlockchainTokenTypeV1::Iot) .await?; tracing::info!( @@ -71,10 +98,12 @@ impl Rewarder { let shutdown = shutdown.clone(); tokio::select! { - _ = shutdown => return Ok(()), + _ = shutdown => break, _ = sleep(sleep_duration) => (), } } + tracing::info!("stopping rewarder"); + Ok(()) } pub async fn reward( diff --git a/iot_verifier/src/runner.rs b/iot_verifier/src/runner.rs index 3f94661ae..a3e646530 100644 --- a/iot_verifier/src/runner.rs +++ b/iot_verifier/src/runner.rs @@ -5,17 +5,15 @@ use crate::{ }; use chrono::{Duration as ChronoDuration, Utc}; use file_store::{ - file_sink, file_sink::FileSinkClient, - file_upload::MessageSender as FileUploadSender, iot_beacon_report::IotBeaconIngestReport, iot_invalid_poc::{IotInvalidBeaconReport, IotInvalidWitnessReport}, iot_valid_poc::{IotPoc, IotValidBeaconReport, IotVerifiedWitnessReport}, iot_witness_report::IotWitnessIngestReport, traits::{IngestId, MsgDecode, ReportId}, - FileType, SCALING_PRECISION, + SCALING_PRECISION, }; -use futures::stream::{self, StreamExt}; +use futures::{future::LocalBoxFuture, stream, StreamExt}; use helium_proto::services::poc_lora::{ InvalidParticipantSide, InvalidReason, LoraInvalidBeaconReportV1, LoraInvalidWitnessReportV1, LoraPocV1, VerificationStatus, @@ -23,7 +21,7 @@ use helium_proto::services::poc_lora::{ use rust_decimal::{Decimal, MathematicalOps}; use rust_decimal_macros::dec; use sqlx::PgPool; -use std::path::Path; +use task_manager::ManagedTask; use tokio::time::{self, MissedTickBehavior}; /// the cadence in seconds at which the DB is polled for ready POCs @@ -36,12 +34,17 @@ const HIP15_TX_REWARD_UNIT_CAP: Decimal = Decimal::TWO; pub struct Runner { pool: PgPool, - cache: String, beacon_interval: ChronoDuration, beacon_interval_tolerance: ChronoDuration, max_witnesses_per_poc: u64, beacon_max_retries: u64, witness_max_retries: u64, + gateway_cache: GatewayCache, + region_cache: RegionCache, + invalid_beacon_sink: FileSinkClient, + invalid_witness_sink: FileSinkClient, + poc_sink: FileSinkClient, + hex_density_map: HexDensityMap, // TODO: fix this } #[derive(thiserror::Error, Debug)] @@ -58,9 +61,28 @@ pub enum FilterStatus { Drop, Include, } + +impl ManagedTask for Runner { + fn start_task( + self: Box, + shutdown_listener: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown_listener)) + } +} + impl Runner { - pub async fn from_settings(settings: &Settings, pool: PgPool) -> Result { - let cache = settings.cache.clone(); + #[allow(clippy::too_many_arguments)] + pub async fn from_settings( + settings: &Settings, + pool: PgPool, + gateway_cache: GatewayCache, + region_cache: RegionCache, + invalid_beacon_sink: FileSinkClient, + invalid_witness_sink: FileSinkClient, + poc_sink: FileSinkClient, + hex_density_map: HexDensityMap, + ) -> Result { let beacon_interval = settings.beacon_interval(); let beacon_interval_tolerance = settings.beacon_interval_tolerance(); let max_witnesses_per_poc = settings.max_witnesses_per_poc; @@ -68,68 +90,66 @@ impl Runner { let witness_max_retries = settings.witness_max_retries; Ok(Self { pool, - cache, beacon_interval, beacon_interval_tolerance, max_witnesses_per_poc, + gateway_cache, + region_cache, beacon_max_retries, witness_max_retries, + invalid_beacon_sink, + invalid_witness_sink, + poc_sink, + hex_density_map, }) } - pub async fn run( - &mut self, - file_upload_tx: FileUploadSender, - gateway_cache: &GatewayCache, - region_cache: &RegionCache, - hex_density_map: impl HexDensityMap, - shutdown: &triggered::Listener, - ) -> anyhow::Result<()> { + pub async fn run(self, shutdown: triggered::Listener) -> anyhow::Result<()> { tracing::info!("starting runner"); let mut db_timer = time::interval(DB_POLL_TIME); db_timer.set_missed_tick_behavior(MissedTickBehavior::Skip); - let store_base_path = Path::new(&self.cache); - - let (iot_invalid_beacon_sink, mut iot_invalid_beacon_sink_server) = - file_sink::FileSinkBuilder::new( - FileType::IotInvalidBeaconReport, - store_base_path, - concat!(env!("CARGO_PKG_NAME"), "_invalid_beacon_report"), - shutdown.clone(), - ) - .deposits(Some(file_upload_tx.clone())) - .roll_time(ChronoDuration::minutes(5)) - .create() - .await?; - - let (iot_invalid_witness_sink, mut iot_invalid_witness_sink_server) = - file_sink::FileSinkBuilder::new( - FileType::IotInvalidWitnessReport, - store_base_path, - concat!(env!("CARGO_PKG_NAME"), "_invalid_witness_report"), - shutdown.clone(), - ) - .deposits(Some(file_upload_tx.clone())) - .roll_time(ChronoDuration::minutes(5)) - .create() - .await?; - - let (iot_poc_sink, mut iot_poc_sink_server) = file_sink::FileSinkBuilder::new( - FileType::IotPoc, - store_base_path, - concat!(env!("CARGO_PKG_NAME"), "_valid_poc"), - shutdown.clone(), - ) - .deposits(Some(file_upload_tx.clone())) - .roll_time(ChronoDuration::minutes(2)) - .create() - .await?; - - tokio::spawn(async move { iot_invalid_beacon_sink_server.run().await }); - tokio::spawn(async move { iot_invalid_witness_sink_server.run().await }); - tokio::spawn(async move { iot_poc_sink_server.run().await }); + // let store_base_path = Path::new(&self.cache); + + // let (iot_invalid_beacon_sink, mut iot_invalid_beacon_sink_server) = + // file_sink::FileSinkBuilder::new( + // FileType::IotInvalidBeaconReport, + // store_base_path, + // concat!(env!("CARGO_PKG_NAME"), "_invalid_beacon_report"), + // shutdown.clone(), + // ) + // .deposits(Some(file_upload_tx.clone())) + // .roll_time(ChronoDuration::minutes(5)) + // .create() + // .await?; + + // let (iot_invalid_witness_sink, mut iot_invalid_witness_sink_server) = + // file_sink::FileSinkBuilder::new( + // FileType::IotInvalidWitnessReport, + // store_base_path, + // concat!(env!("CARGO_PKG_NAME"), "_invalid_witness_report"), + // shutdown.clone(), + // ) + // .deposits(Some(file_upload_tx.clone())) + // .roll_time(ChronoDuration::minutes(5)) + // .create() + // .await?; + + // let (iot_poc_sink, mut iot_poc_sink_server) = file_sink::FileSinkBuilder::new( + // FileType::IotPoc, + // store_base_path, + // concat!(env!("CARGO_PKG_NAME"), "_valid_poc"), + // shutdown.clone(), + // ) + // .deposits(Some(file_upload_tx.clone())) + // .roll_time(ChronoDuration::minutes(2)) + // .create() + // .await?; + + // tokio::spawn(async move { iot_invalid_beacon_sink_server.run().await }); + // tokio::spawn(async move { iot_invalid_witness_sink_server.run().await }); + // tokio::spawn(async move { iot_poc_sink_server.run().await }); loop { if shutdown.is_triggered() { @@ -138,13 +158,7 @@ impl Runner { tokio::select! { _ = shutdown.clone() => break, _ = db_timer.tick() => - match self.handle_db_tick( shutdown.clone(), - &iot_invalid_beacon_sink, - &iot_invalid_witness_sink, - &iot_poc_sink, - gateway_cache, - region_cache, - hex_density_map.clone()).await { + match self.handle_db_tick().await { Ok(()) => (), Err(err) => { tracing::error!("fatal db runner error: {err:?}"); @@ -157,16 +171,7 @@ impl Runner { } #[allow(clippy::too_many_arguments)] - async fn handle_db_tick( - &self, - _shutdown: triggered::Listener, - iot_invalid_beacon_sink: &FileSinkClient, - iot_invalid_witness_sink: &FileSinkClient, - iot_poc_sink: &FileSinkClient, - gateway_cache: &GatewayCache, - region_cache: &RegionCache, - hex_density_map: impl HexDensityMap, - ) -> anyhow::Result<()> { + async fn handle_db_tick(&self) -> anyhow::Result<()> { tracing::info!("starting query get_next_beacons"); let db_beacon_reports = Report::get_next_beacons(&self.pool, self.beacon_max_retries).await?; @@ -184,27 +189,13 @@ impl Runner { tracing::info!("{beacon_len} beacons ready for verification"); stream::iter(db_beacon_reports) - .for_each_concurrent(BEACON_WORKERS, |db_beacon| { - let hdm = hex_density_map.clone(); - async move { - let beacon_id = db_beacon.id.clone(); - match self - .handle_beacon_report( - db_beacon, - iot_invalid_beacon_sink, - iot_invalid_witness_sink, - iot_poc_sink, - gateway_cache, - region_cache, - hdm, - ) - .await - { - Ok(()) => (), - Err(err) => { - tracing::warn!("failed to handle beacon: {err:?}"); - _ = Report::update_attempts(&self.pool, &beacon_id, Utc::now()).await; - } + .for_each_concurrent(BEACON_WORKERS, |db_beacon| async move { + let beacon_id = db_beacon.id.clone(); + match self.handle_beacon_report(db_beacon).await { + Ok(()) => (), + Err(err) => { + tracing::warn!("failed to handle beacon: {err:?}"); + _ = Report::update_attempts(&self.pool, &beacon_id, Utc::now()).await; } } }) @@ -214,16 +205,7 @@ impl Runner { } #[allow(clippy::too_many_arguments)] - async fn handle_beacon_report( - &self, - db_beacon: Report, - iot_invalid_beacon_sink: &FileSinkClient, - iot_invalid_witness_sink: &FileSinkClient, - iot_poc_sink: &FileSinkClient, - gateway_cache: &GatewayCache, - region_cache: &RegionCache, - hex_density_map: impl HexDensityMap, - ) -> anyhow::Result<()> { + async fn handle_beacon_report(&self, db_beacon: Report) -> anyhow::Result<()> { let entropy_start_time = match db_beacon.timestamp { Some(v) => v, None => return Ok(()), @@ -264,9 +246,9 @@ impl Runner { // verify POC beacon let beacon_verify_result = poc .verify_beacon( - hex_density_map.clone(), - gateway_cache, - region_cache, + self.hex_density_map.clone(), + self.gateway_cache.clone(), + self.region_cache.clone(), &self.pool, self.beacon_interval, self.beacon_interval_tolerance, @@ -277,7 +259,11 @@ impl Runner { // beacon is valid, verify the POC witnesses if let Some(beacon_info) = beacon_verify_result.gateway_info { let verified_witnesses_result = poc - .verify_witnesses(&beacon_info, hex_density_map, gateway_cache) + .verify_witnesses( + &beacon_info, + self.hex_density_map.clone(), + self.gateway_cache.clone(), + ) .await?; // check if there are any failed witnesses // if so update the DB attempts count @@ -358,7 +344,6 @@ impl Runner { valid_beacon_report, selected_witnesses, unselected_witnesses, - iot_poc_sink, ) .await?; } @@ -369,8 +354,6 @@ impl Runner { &beacon_report, witnesses, beacon_verify_result.invalid_reason, - iot_invalid_beacon_sink, - iot_invalid_witness_sink, ) .await?; } @@ -383,8 +366,6 @@ impl Runner { beacon_report: &IotBeaconIngestReport, witness_reports: Vec, invalid_reason: InvalidReason, - iot_invalid_beacon_sink: &FileSinkClient, - iot_invalid_witness_sink: &FileSinkClient, ) -> anyhow::Result<()> { // the beacon is invalid, which in turn renders all witnesses invalid let beacon = &beacon_report.report; @@ -398,7 +379,8 @@ impl Runner { let invalid_poc_proto: LoraInvalidBeaconReportV1 = invalid_poc.into(); // save invalid poc to s3, if write fails update attempts and go no further // allow the poc to be reprocessed next tick - match iot_invalid_beacon_sink + match self + .invalid_beacon_sink .write( invalid_poc_proto, &[("reason", invalid_reason.as_str_name())], @@ -426,7 +408,8 @@ impl Runner { }; let invalid_witness_report_proto: LoraInvalidWitnessReportV1 = invalid_witness_report.into(); - match iot_invalid_witness_sink + match self + .invalid_witness_sink .write( invalid_witness_report_proto, &[("reason", invalid_reason.as_str_name())], @@ -450,7 +433,6 @@ impl Runner { valid_beacon_report: IotValidBeaconReport, selected_witnesses: Vec, unselected_witnesses: Vec, - iot_poc_sink: &FileSinkClient, ) -> anyhow::Result<()> { let received_timestamp = valid_beacon_report.received_timestamp; let pub_key = valid_beacon_report.report.pub_key.clone(); @@ -474,7 +456,7 @@ impl Runner { let poc_proto: LoraPocV1 = iot_poc.into(); // save the poc to s3, if write fails update attempts and go no further // allow the poc to be reprocessed next tick - match iot_poc_sink.write(poc_proto, []).await { + match self.poc_sink.write(poc_proto, []).await { Ok(_) => (), Err(err) => { tracing::error!("failed to save invalid_witness_report to s3, {err}"); diff --git a/iot_verifier/src/tx_scaler.rs b/iot_verifier/src/tx_scaler.rs index 6289d5c04..146987c73 100644 --- a/iot_verifier/src/tx_scaler.rs +++ b/iot_verifier/src/tx_scaler.rs @@ -1,20 +1,22 @@ use crate::{ gateway_updater::MessageReceiver, - hex_density::{compute_hex_density_map, GlobalHexMap, HexDensityMap, SharedHexDensityMap}, + hex_density::{compute_hex_density_map, GlobalHexMap, HexDensityMap}, last_beacon::LastBeacon, Settings, }; use chrono::{DateTime, Duration, Utc}; +use futures::future::LocalBoxFuture; use helium_crypto::PublicKeyBinary; use sqlx::PgPool; use std::collections::HashMap; +use task_manager::ManagedTask; // The number in minutes within which the gateway has registered a beacon // to the oracle for inclusion in transmit scaling density calculations const HIP_17_INTERACTIVITY_LIMIT: i64 = 3600; pub struct Server { - hex_density_map: SharedHexDensityMap, + pub hex_density_map: HexDensityMap, pool: PgPool, refresh_offset: Duration, gateway_cache_receiver: MessageReceiver, @@ -28,14 +30,23 @@ pub enum TxScalerError { RecentActivity(#[from] sqlx::Error), } +impl ManagedTask for Server { + fn start_task( + self: Box, + shutdown_listener: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown_listener)) + } +} + impl Server { pub async fn from_settings( settings: &Settings, pool: PgPool, gateway_cache_receiver: MessageReceiver, - ) -> Result { + ) -> anyhow::Result { let mut server = Self { - hex_density_map: SharedHexDensityMap::new(), + hex_density_map: HexDensityMap::new(), pool, refresh_offset: settings.loader_window_max_lookback_age(), gateway_cache_receiver, @@ -46,11 +57,7 @@ impl Server { Ok(server) } - pub fn hex_density_map(&self) -> impl HexDensityMap { - self.hex_density_map.clone() - } - - pub async fn run(&mut self, shutdown: &triggered::Listener) -> Result<(), TxScalerError> { + pub async fn run(mut self, shutdown: triggered::Listener) -> anyhow::Result<()> { tracing::info!("density_scaler: starting transmit scaler process"); loop { @@ -61,12 +68,15 @@ impl Server { tokio::select! { _ = self.gateway_cache_receiver.changed() => self.refresh_scaling_map().await?, - _ = shutdown.clone() => return Ok(()), + _ = shutdown.clone() => break, } } + + tracing::info!("stopping transmit scaler process"); + Ok(()) } - pub async fn refresh_scaling_map(&mut self) -> Result<(), TxScalerError> { + pub async fn refresh_scaling_map(&mut self) -> anyhow::Result<()> { let refresh_start = Utc::now() - self.refresh_offset; tracing::info!("density_scaler: generating hex scaling map, starting at {refresh_start:?}"); let mut global_map = GlobalHexMap::new(); diff --git a/mobile_config/Cargo.toml b/mobile_config/Cargo.toml index d68bee618..a9c81e249 100644 --- a/mobile_config/Cargo.toml +++ b/mobile_config/Cargo.toml @@ -38,3 +38,5 @@ tonic = {workspace = true} tracing = {workspace = true} tracing-subscriber = {workspace = true} triggered = {workspace = true} +task-manager = { path = "../task_manager" } +tokio-util = { workspace = true } diff --git a/mobile_config/src/main.rs b/mobile_config/src/main.rs index 81ad76729..24796ef42 100644 --- a/mobile_config/src/main.rs +++ b/mobile_config/src/main.rs @@ -1,5 +1,6 @@ use anyhow::{Error, Result}; use clap::Parser; +use futures::future::LocalBoxFuture; use futures_util::TryFutureExt; use helium_proto::services::mobile_config::{ AdminServer, AuthorizationServer, EntityServer, GatewayServer, @@ -9,8 +10,8 @@ use mobile_config::{ entity_service::EntityService, gateway_service::GatewayService, key_cache::KeyCache, settings::Settings, }; -use std::{path::PathBuf, time::Duration}; -use tokio::signal; +use std::{net::SocketAddr, path::PathBuf, time::Duration}; +use task_manager::{ManagedTask, TaskManager}; use tonic::transport; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; @@ -58,31 +59,15 @@ impl Daemon { .with(tracing_subscriber::fmt::layer()) .init(); - // Configure shutdown trigger - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); - // Install prometheus metrics exporter poc_metrics::start_metrics(&settings.metrics)?; // Create database pool - let (pool, pool_handle) = settings - .database - .connect("mobile-config-store", shutdown_listener.clone()) - .await?; + let pool = settings.database.connect("mobile-config-store").await?; sqlx::migrate!().run(&pool).await?; // Create on-chain metadata pool - let (metadata_pool, md_pool_handle) = settings - .metadata - .connect("mobile-config-metadata", shutdown_listener.clone()) - .await?; + let metadata_pool = settings.metadata.connect("mobile-config-metadata").await?; let listen_addr = settings.listen_addr()?; @@ -102,23 +87,43 @@ impl Daemon { settings.signing_keypair()?, ); - let server = transport::Server::builder() - .http2_keepalive_interval(Some(Duration::from_secs(250))) - .http2_keepalive_timeout(Some(Duration::from_secs(60))) - .add_service(AdminServer::new(admin_svc)) - .add_service(GatewayServer::new(gateway_svc)) - .add_service(AuthorizationServer::new(auth_svc)) - .add_service(EntityServer::new(entity_svc)) - .serve_with_shutdown(listen_addr, shutdown_listener) - .map_err(Error::from); - - tokio::try_join!( - pool_handle.map_err(Error::from), - md_pool_handle.map_err(Error::from), - server, - )?; - - Ok(()) + let grpc_server = GrpcServer { + listen_addr, + admin_svc, + gateway_svc, + auth_svc, + entity_svc, + }; + + TaskManager::builder().add_task(grpc_server).start().await + } +} + +pub struct GrpcServer { + listen_addr: SocketAddr, + admin_svc: AdminService, + gateway_svc: GatewayService, + auth_svc: AuthorizationService, + entity_svc: EntityService, +} + +impl ManagedTask for GrpcServer { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(async move { + transport::Server::builder() + .http2_keepalive_interval(Some(Duration::from_secs(250))) + .http2_keepalive_timeout(Some(Duration::from_secs(60))) + .add_service(AdminServer::new(self.admin_svc)) + .add_service(GatewayServer::new(self.gateway_svc)) + .add_service(AuthorizationServer::new(self.auth_svc)) + .add_service(EntityServer::new(self.entity_svc)) + .serve_with_shutdown(self.listen_addr, shutdown) + .map_err(Error::from) + .await + }) } } diff --git a/mobile_packet_verifier/Cargo.toml b/mobile_packet_verifier/Cargo.toml index 60f44e47d..7a6d57a3e 100644 --- a/mobile_packet_verifier/Cargo.toml +++ b/mobile_packet_verifier/Cargo.toml @@ -34,3 +34,5 @@ triggered = {workspace = true} http = {workspace = true} http-serde = {workspace = true} sha2 = {workspace = true} +tokio-util = { workspace = true } +task-manager = { path = "../task_manager" } diff --git a/mobile_packet_verifier/src/daemon.rs b/mobile_packet_verifier/src/daemon.rs index 22780e854..128ba4fc2 100644 --- a/mobile_packet_verifier/src/daemon.rs +++ b/mobile_packet_verifier/src/daemon.rs @@ -1,5 +1,5 @@ use crate::{burner::Burner, settings::Settings}; -use anyhow::{bail, Error, Result}; +use anyhow::{bail, Result}; use chrono::{TimeZone, Utc}; use file_store::{ file_info_poller::{FileInfoStream, LookbackBehavior}, @@ -8,12 +8,12 @@ use file_store::{ mobile_session::DataTransferSessionIngestReport, FileSinkBuilder, FileStore, FileType, }; -use futures_util::TryFutureExt; +use futures::future::LocalBoxFuture; use mobile_config::{client::AuthorizationClient, GatewayClient}; -use solana::{SolanaNetwork, SolanaRpc}; +use solana::{balance_monitor::BalanceMonitor, SolanaNetwork, SolanaRpc}; use sqlx::{Pool, Postgres}; +use task_manager::{ManagedTask, TaskManager}; use tokio::{ - signal, sync::mpsc::Receiver, time::{sleep_until, Duration, Instant}, }; @@ -28,6 +28,18 @@ pub struct Daemon { invalid_data_session_report_sink: FileSinkClient, } +impl ManagedTask for Daemon +where + S: SolanaNetwork, +{ + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } +} + impl Daemon { pub fn new( settings: &Settings, @@ -54,7 +66,8 @@ impl Daemon where S: SolanaNetwork, { - pub async fn run(mut self, shutdown: &triggered::Listener) -> Result<()> { + pub async fn run(mut self, shutdown: triggered::Listener) -> Result<()> { + tracing::info!("starting daemon"); // Set the initial burn period to one minute let mut burn_time = Instant::now() + Duration::from_secs(60); loop { @@ -76,9 +89,11 @@ where self.burner.burn(&self.pool).await?; burn_time = Instant::now() + self.burn_period; } - _ = shutdown.clone() => return Ok(()), + _ = shutdown.clone() => break, } } + tracing::info!("stopping daemon"); + Ok(()) } } @@ -89,20 +104,8 @@ impl Cmd { pub async fn run(self, settings: &Settings) -> Result<()> { poc_metrics::start_metrics(&settings.metrics)?; - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); - // Set up the postgres pool: - let (pool, conn_handler) = settings - .database - .connect("mobile-packet-verifier", shutdown_listener.clone()) - .await?; + let pool = settings.database.connect("mobile-packet-verifier").await?; sqlx::migrate!().run(&pool).await?; // Set up the solana network: @@ -116,37 +119,29 @@ impl Cmd { None }; - let sol_balance_monitor = solana::balance_monitor::start( - env!("CARGO_PKG_NAME"), - solana.clone(), - shutdown_listener.clone(), - ) - .await?; + let sol_balance_monitor = BalanceMonitor::new(env!("CARGO_PKG_NAME"), solana.clone())?; - let (file_upload_tx, file_upload_rx) = file_upload::message_channel(); - let file_upload = - file_upload::FileUpload::from_settings(&settings.output, file_upload_rx).await?; + let (file_upload, file_upload_server) = + file_upload::FileUpload::from_settings(&settings.output).await?; let store_base_path = std::path::Path::new(&settings.cache); - let (valid_sessions, mut valid_sessions_server) = FileSinkBuilder::new( + let (valid_sessions, valid_sessions_server) = FileSinkBuilder::new( FileType::ValidDataTransferSession, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_valid_data_transfer_session"), - shutdown_listener.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(true) .create() .await?; - let (invalid_sessions, mut invalid_sessions_server) = FileSinkBuilder::new( + let (invalid_sessions, invalid_sessions_server) = FileSinkBuilder::new( FileType::InvalidDataTransferSessionIngestReport, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_invalid_data_transfer_session"), - shutdown_listener.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(false) .create() .await?; @@ -155,7 +150,7 @@ impl Cmd { let file_store = FileStore::from_settings(&settings.ingest).await?; - let (reports, source_join_handle) = + let (reports, reports_server) = file_source::continuous_source::() .db(pool.clone()) .store(file_store) @@ -164,9 +159,7 @@ impl Cmd { )) .file_type(FileType::DataTransferSessionIngestReport) .lookback(LookbackBehavior::StartAfter(settings.start_after())) - .build()? - .start(shutdown_listener.clone()) - .await?; + .create()?; let gateway_client = GatewayClient::from_settings(&settings.config_client)?; let auth_client = AuthorizationClient::from_settings(&settings.config_client)?; @@ -181,16 +174,14 @@ impl Cmd { invalid_sessions, ); - tokio::try_join!( - source_join_handle.map_err(Error::from), - valid_sessions_server.run().map_err(Error::from), - invalid_sessions_server.run().map_err(Error::from), - file_upload.run(&shutdown_listener).map_err(Error::from), - daemon.run(&shutdown_listener).map_err(Error::from), - conn_handler.map_err(Error::from), - sol_balance_monitor.map_err(Error::from), - )?; - - Ok(()) + TaskManager::builder() + .add_task(file_upload_server) + .add_task(valid_sessions_server) + .add_task(invalid_sessions_server) + .add_task(reports_server) + .add_task(sol_balance_monitor) + .add_task(daemon) + .start() + .await } } diff --git a/mobile_verifier/Cargo.toml b/mobile_verifier/Cargo.toml index a7fa49bee..9574e3eb3 100644 --- a/mobile_verifier/Cargo.toml +++ b/mobile_verifier/Cargo.toml @@ -43,4 +43,6 @@ reward-scheduler = {path = "../reward_scheduler"} price = {path = "../price"} rand = {workspace = true} async-trait = {workspace = true} -retainer = {workspace = true} \ No newline at end of file +retainer = {workspace = true} +tokio-util = { workspace = true } +task-manager = { path = "../task_manager" } diff --git a/mobile_verifier/src/cli/reward_from_db.rs b/mobile_verifier/src/cli/reward_from_db.rs index b3b29a322..9a6fb7151 100644 --- a/mobile_verifier/src/cli/reward_from_db.rs +++ b/mobile_verifier/src/cli/reward_from_db.rs @@ -32,11 +32,7 @@ impl Cmd { let epoch = start..end; let expected_rewards = get_scheduled_tokens_for_poc_and_dc(epoch.end - epoch.start); - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let (pool, _join_handle) = settings - .database - .connect(env!("CARGO_PKG_NAME"), shutdown_listener) - .await?; + let pool = settings.database.connect(env!("CARGO_PKG_NAME")).await?; let heartbeats = HeartbeatReward::validated(&pool, &epoch); let speedtests = SpeedtestAverages::validated(&pool, epoch.end).await?; @@ -80,7 +76,6 @@ impl Cmd { }))? ); - shutdown_trigger.trigger(); Ok(()) } } diff --git a/mobile_verifier/src/cli/server.rs b/mobile_verifier/src/cli/server.rs index 6d5065787..4a5416ed5 100644 --- a/mobile_verifier/src/cli/server.rs +++ b/mobile_verifier/src/cli/server.rs @@ -3,7 +3,7 @@ use crate::{ speedtests::SpeedtestDaemon, subscriber_location::SubscriberLocationIngestor, telemetry, Settings, }; -use anyhow::{Error, Result}; +use anyhow::Result; use chrono::Duration; use file_store::{ file_info_poller::LookbackBehavior, file_sink, file_source, file_upload, @@ -12,9 +12,9 @@ use file_store::{ FileType, }; -use futures_util::TryFutureExt; use mobile_config::client::{AuthorizationClient, EntityClient, GatewayClient}; use price::PriceTracker; +use task_manager::TaskManager; use tokio::signal; #[derive(Debug, clap::Args)] @@ -24,26 +24,13 @@ impl Cmd { pub async fn run(self, settings: &Settings) -> Result<()> { poc_metrics::start_metrics(&settings.metrics)?; - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); - - let (pool, db_join_handle) = settings - .database - .connect(env!("CARGO_PKG_NAME"), shutdown_listener.clone()) - .await?; + let pool = settings.database.connect(env!("CARGO_PKG_NAME")).await?; sqlx::migrate!().run(&pool).await?; telemetry::initialize(&pool).await?; - let (file_upload_tx, file_upload_rx) = file_upload::message_channel(); - let file_upload = - file_upload::FileUpload::from_settings(&settings.output, file_upload_rx).await?; + let (file_upload, file_upload_server) = + file_upload::FileUpload::from_settings(&settings.output).await?; let store_base_path = std::path::Path::new(&settings.cache); @@ -55,28 +42,34 @@ impl Cmd { let auth_client = AuthorizationClient::from_settings(&settings.config_client)?; let entity_client = EntityClient::from_settings(&settings.config_client)?; + // todo: update price tracker to not require shutdown listener to be passed in // price tracker - let (price_tracker, tracker_process) = - PriceTracker::start(&settings.price_tracker, shutdown_listener.clone()).await?; + let (shutdown_trigger, shutdown) = triggered::trigger(); + let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; + tokio::spawn(async move { + tokio::select! { + _ = sigterm.recv() => shutdown_trigger.trigger(), + _ = signal::ctrl_c() => shutdown_trigger.trigger(), + } + }); + let (price_tracker, _price_receiver) = + PriceTracker::start(&settings.price_tracker, shutdown.clone()).await?; // Heartbeats - let (heartbeats, heartbeats_join_handle) = + let (heartbeats, heartbeats_ingest_server) = file_source::continuous_source::() .db(pool.clone()) .store(report_ingest.clone()) .lookback(LookbackBehavior::StartAfter(settings.start_after())) .file_type(FileType::CellHeartbeatIngestReport) - .build()? - .start(shutdown_listener.clone()) - .await?; + .create()?; - let (valid_heartbeats, mut valid_heartbeats_server) = file_sink::FileSinkBuilder::new( + let (valid_heartbeats, valid_heartbeats_server) = file_sink::FileSinkBuilder::new( FileType::ValidatedHeartbeat, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_heartbeat"), - shutdown_listener.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(false) .roll_time(Duration::minutes(15)) .create() @@ -90,23 +83,20 @@ impl Cmd { ); // Speedtests - let (speedtests, speedtests_join_handle) = + let (speedtests, speedtests_ingest_server) = file_source::continuous_source::() .db(pool.clone()) .store(report_ingest.clone()) .lookback(LookbackBehavior::StartAfter(settings.start_after())) .file_type(FileType::CellSpeedtestIngestReport) - .build()? - .start(shutdown_listener.clone()) - .await?; + .create()?; - let (valid_speedtests, mut valid_speedtests_server) = file_sink::FileSinkBuilder::new( + let (valid_speedtests, valid_speedtests_server) = file_sink::FileSinkBuilder::new( FileType::SpeedtestAvg, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_speedtest_average"), - shutdown_listener.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(false) .roll_time(Duration::minutes(15)) .create() @@ -121,24 +111,22 @@ impl Cmd { // Mobile rewards let reward_period_hours = settings.rewards; - let (mobile_rewards, mut mobile_rewards_server) = file_sink::FileSinkBuilder::new( + let (mobile_rewards, mobile_rewards_server) = file_sink::FileSinkBuilder::new( FileType::MobileRewardShare, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_radio_reward_shares"), - shutdown_listener.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(false) .create() .await?; - let (reward_manifests, mut reward_manifests_server) = file_sink::FileSinkBuilder::new( + let (reward_manifests, reward_manifests_server) = file_sink::FileSinkBuilder::new( FileType::RewardManifest, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_reward_manifest"), - shutdown_listener.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(false) .create() .await?; @@ -154,24 +142,21 @@ impl Cmd { ); // subscriber location - let (subscriber_location_ingest, subscriber_location_ingest_join_handle) = + let (subscriber_location_ingest, subscriber_location_ingest_server) = file_source::continuous_source::() .db(pool.clone()) .store(report_ingest.clone()) .lookback(LookbackBehavior::StartAfter(settings.start_after())) .file_type(FileType::SubscriberLocationIngestReport) - .build()? - .start(shutdown_listener.clone()) - .await?; + .create()?; - let (verified_subscriber_location, mut verified_subscriber_location_server) = + let (verified_subscriber_location, verified_subscriber_location_server) = file_sink::FileSinkBuilder::new( FileType::VerifiedSubscriberLocationIngestReport, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_verified_subscriber_location"), - shutdown_listener.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .auto_commit(false) .create() .await?; @@ -185,46 +170,33 @@ impl Cmd { ); // data transfers - let (data_session_ingest, data_session_ingest_join_handle) = + let (data_session_ingest, data_session_ingest_server) = file_source::continuous_source::() .db(pool.clone()) .store(data_transfer_ingest.clone()) .lookback(LookbackBehavior::StartAfter(settings.start_after())) .file_type(FileType::ValidDataTransferSession) - .build()? - .start(shutdown_listener.clone()) - .await?; - - let data_session_ingestor = DataSessionIngestor::new(pool.clone()); - - tokio::try_join!( - db_join_handle.map_err(Error::from), - valid_heartbeats_server.run().map_err(Error::from), - valid_speedtests_server.run().map_err(Error::from), - mobile_rewards_server.run().map_err(Error::from), - file_upload.run(&shutdown_listener).map_err(Error::from), - reward_manifests_server.run().map_err(Error::from), - verified_subscriber_location_server - .run() - .map_err(Error::from), - subscriber_location_ingestor - .run(&shutdown_listener) - .map_err(Error::from), - data_session_ingestor - .run(data_session_ingest, shutdown_listener.clone()) - .map_err(Error::from), - tracker_process.map_err(Error::from), - heartbeats_join_handle.map_err(Error::from), - speedtests_join_handle.map_err(Error::from), - heartbeat_daemon.run(shutdown_listener.clone()), - speedtest_daemon.run(shutdown_listener.clone()), - rewarder.run(shutdown_listener.clone()), - subscriber_location_ingest_join_handle.map_err(anyhow::Error::from), - data_session_ingest_join_handle.map_err(anyhow::Error::from), - )?; - - tracing::info!("Shutting down verifier server"); - - Ok(()) + .create()?; + + let data_session_ingestor = DataSessionIngestor::new(pool.clone(), data_session_ingest); + + TaskManager::builder() + .add_task(file_upload_server) + .add_task(valid_heartbeats_server) + .add_task(valid_speedtests_server) + .add_task(mobile_rewards_server) + .add_task(reward_manifests_server) + .add_task(verified_subscriber_location_server) + .add_task(heartbeats_ingest_server) + .add_task(speedtests_ingest_server) + .add_task(subscriber_location_ingest_server) + .add_task(data_session_ingest_server) + .add_task(subscriber_location_ingestor) + .add_task(data_session_ingestor) + .add_task(rewarder) + .add_task(heartbeat_daemon) + .add_task(speedtest_daemon) + .start() + .await } } diff --git a/mobile_verifier/src/data_session.rs b/mobile_verifier/src/data_session.rs index 49f1b6d0f..b8f70414d 100644 --- a/mobile_verifier/src/data_session.rs +++ b/mobile_verifier/src/data_session.rs @@ -1,30 +1,41 @@ use chrono::{DateTime, Utc}; use file_store::{file_info_poller::FileInfoStream, mobile_transfer::ValidDataTransferSession}; use futures::{ + future::LocalBoxFuture, stream::{Stream, StreamExt, TryStreamExt}, TryFutureExt, }; use helium_crypto::PublicKeyBinary; use sqlx::{PgPool, Postgres, Transaction}; use std::{collections::HashMap, ops::Range}; +use task_manager::ManagedTask; use tokio::sync::mpsc::Receiver; pub struct DataSessionIngestor { pub pool: PgPool, + pub receiver: Receiver>, +} + +impl ManagedTask for DataSessionIngestor { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } } pub type HotspotMap = HashMap; impl DataSessionIngestor { - pub fn new(pool: sqlx::Pool) -> Self { - Self { pool } + pub fn new( + pool: sqlx::Pool, + receiver: Receiver>, + ) -> Self { + Self { pool, receiver } } - pub async fn run( - self, - mut receiver: Receiver>, - shutdown: triggered::Listener, - ) -> anyhow::Result<()> { + pub async fn run(mut self, shutdown: triggered::Listener) -> anyhow::Result<()> { tracing::info!("starting DataSessionIngestor"); tokio::spawn(async move { loop { @@ -33,7 +44,7 @@ impl DataSessionIngestor { tracing::info!("DataSessionIngestor shutting down"); break; } - Some(file) = receiver.recv() => self.process_file(file).await?, + Some(file) = self.receiver.recv() => self.process_file(file).await?, } } diff --git a/mobile_verifier/src/heartbeats.rs b/mobile_verifier/src/heartbeats.rs index 3bcc89353..1cf334fcd 100644 --- a/mobile_verifier/src/heartbeats.rs +++ b/mobile_verifier/src/heartbeats.rs @@ -7,6 +7,7 @@ use file_store::{ heartbeat::CellHeartbeatIngestReport, }; use futures::{ + future::LocalBoxFuture, stream::{Stream, StreamExt, TryStreamExt}, TryFutureExt, }; @@ -17,6 +18,7 @@ use retainer::Cache; use rust_decimal::{prelude::ToPrimitive, Decimal}; use sqlx::{Postgres, Transaction}; use std::{ops::Range, pin::pin, sync::Arc, time}; +use task_manager::ManagedTask; use tokio::sync::mpsc::Receiver; #[derive(Debug, Clone, PartialEq, Eq, Hash, sqlx::FromRow)] @@ -49,6 +51,15 @@ pub struct HeartbeatDaemon { file_sink: FileSinkClient, } +impl ManagedTask for HeartbeatDaemon { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } +} + impl HeartbeatDaemon { pub fn new( pool: sqlx::Pool, diff --git a/mobile_verifier/src/rewarder.rs b/mobile_verifier/src/rewarder.rs index bbabb36dd..fa87302c4 100644 --- a/mobile_verifier/src/rewarder.rs +++ b/mobile_verifier/src/rewarder.rs @@ -9,6 +9,7 @@ use anyhow::bail; use chrono::{DateTime, Duration, TimeZone, Utc}; use db_store::meta; use file_store::{file_sink::FileSinkClient, traits::TimestampEncode}; +use futures::future::LocalBoxFuture; use helium_proto::services::poc_mobile::mobile_reward_share::Reward as ProtoReward; use helium_proto::RewardManifest; use price::PriceTracker; @@ -17,6 +18,7 @@ use rust_decimal::{prelude::ToPrimitive, Decimal}; use rust_decimal_macros::dec; use sqlx::{PgExecutor, Pool, Postgres}; use std::ops::Range; +use task_manager::ManagedTask; use tokio::time::sleep; const REWARDS_NOT_CURRENT_DELAY_PERIOD: i64 = 5; @@ -31,6 +33,15 @@ pub struct Rewarder { disable_discovery_loc_rewards_to_s3: bool, } +impl ManagedTask for Rewarder { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } +} + impl Rewarder { pub fn new( pool: Pool, diff --git a/mobile_verifier/src/speedtests.rs b/mobile_verifier/src/speedtests.rs index db70e7fae..d7ae20a0e 100644 --- a/mobile_verifier/src/speedtests.rs +++ b/mobile_verifier/src/speedtests.rs @@ -6,6 +6,7 @@ use file_store::{ traits::TimestampEncode, }; use futures::{ + future::LocalBoxFuture, stream::{Stream, StreamExt, TryStreamExt}, TryFutureExt, }; @@ -22,6 +23,7 @@ use std::{ collections::{HashMap, VecDeque}, pin::pin, }; +use task_manager::ManagedTask; use tokio::sync::mpsc::Receiver; const SPEEDTEST_AVG_MAX_DATA_POINTS: usize = 6; @@ -60,6 +62,15 @@ pub struct SpeedtestDaemon { file_sink: FileSinkClient, } +impl ManagedTask for SpeedtestDaemon { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } +} + impl SpeedtestDaemon { pub fn new( pool: sqlx::Pool, diff --git a/mobile_verifier/src/subscriber_location.rs b/mobile_verifier/src/subscriber_location.rs index 2a9fbba61..1787dca93 100644 --- a/mobile_verifier/src/subscriber_location.rs +++ b/mobile_verifier/src/subscriber_location.rs @@ -7,7 +7,7 @@ use file_store::{ VerifiedSubscriberLocationIngestReport, }, }; -use futures::{StreamExt, TryStreamExt}; +use futures::{future::LocalBoxFuture, StreamExt, TryStreamExt}; use helium_crypto::PublicKeyBinary; use helium_proto::services::mobile_config::NetworkKeyRole; use helium_proto::services::poc_mobile::{ @@ -16,6 +16,7 @@ use helium_proto::services::poc_mobile::{ use mobile_config::client::{AuthorizationClient, EntityClient}; use sqlx::{PgPool, Postgres, Transaction}; use std::ops::Range; +use task_manager::ManagedTask; use tokio::sync::mpsc::Receiver; pub type SubscriberValidatedLocations = Vec>; @@ -28,6 +29,15 @@ pub struct SubscriberLocationIngestor { verified_report_sink: FileSinkClient, } +impl ManagedTask for SubscriberLocationIngestor { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } +} + impl SubscriberLocationIngestor { pub fn new( pool: sqlx::Pool, @@ -44,7 +54,7 @@ impl SubscriberLocationIngestor { verified_report_sink, } } - pub async fn run(mut self, shutdown: &triggered::Listener) -> anyhow::Result<()> { + pub async fn run(mut self, shutdown: triggered::Listener) -> anyhow::Result<()> { loop { if shutdown.is_triggered() { break; diff --git a/poc_entropy/Cargo.toml b/poc_entropy/Cargo.toml index c2c1322c6..90a4a61a8 100644 --- a/poc_entropy/Cargo.toml +++ b/poc_entropy/Cargo.toml @@ -30,8 +30,10 @@ tracing-subscriber = { workspace = true } metrics = {workspace = true } metrics-exporter-prometheus = { workspace = true } tokio = { workspace = true } +tokio-util = { workspace = true } chrono = { workspace = true } helium-proto = { workspace = true } helium-crypto = { workspace = true } file-store = { path = "../file_store" } poc-metrics = { path = "../metrics" } +task-manager = { path = "../task_manager" } diff --git a/poc_entropy/src/entropy_generator.rs b/poc_entropy/src/entropy_generator.rs index aec84bd3d..4891d4ccc 100644 --- a/poc_entropy/src/entropy_generator.rs +++ b/poc_entropy/src/entropy_generator.rs @@ -1,7 +1,7 @@ use base64::Engine; use chrono::Utc; use file_store::file_sink; -use futures::TryFutureExt; +use futures::{future::LocalBoxFuture, TryFutureExt}; use helium_proto::EntropyReportV1; use jsonrpsee::{ core::client::ClientT, @@ -10,6 +10,7 @@ use jsonrpsee::{ }; use serde::{Deserialize, Serialize}; use serde_json::json; +use task_manager::ManagedTask; use tokio::{sync::watch, time}; pub const ENTROPY_TICK_TIME: time::Duration = time::Duration::from_secs(60); @@ -76,6 +77,16 @@ pub struct EntropyGenerator { client: HttpClient, sender: MessageSender, + file_sink: file_sink::FileSinkClient, +} + +impl ManagedTask for EntropyGenerator { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } } #[derive(thiserror::Error, Debug)] @@ -89,7 +100,10 @@ pub enum GetEntropyError { } impl EntropyGenerator { - pub async fn new(url: impl AsRef) -> Result { + pub async fn new( + url: impl AsRef, + file_sink: file_sink::FileSinkClient, + ) -> Result { let client = HttpClientBuilder::default() .request_timeout(ENTROPY_TIMEOUT) .build(url)?; @@ -112,14 +126,11 @@ impl EntropyGenerator { client, receiver, sender, + file_sink, }) } - pub async fn run( - &mut self, - file_sink: file_sink::FileSinkClient, - shutdown: &triggered::Listener, - ) -> anyhow::Result<()> { + pub async fn run(mut self, shutdown: triggered::Listener) -> anyhow::Result<()> { tracing::info!("started entropy generator"); let mut entropy_timer = time::interval(ENTROPY_TICK_TIME); entropy_timer.set_missed_tick_behavior(time::MissedTickBehavior::Delay); @@ -130,7 +141,7 @@ impl EntropyGenerator { } tokio::select! { _ = shutdown.clone() => break, - _ = entropy_timer.tick() => match self.handle_entropy_tick(&file_sink).await { + _ = entropy_timer.tick() => match self.handle_entropy_tick().await { Ok(()) => (), Err(err) => { tracing::error!("fatal entropy generator error: {err:?}"); @@ -147,10 +158,7 @@ impl EntropyGenerator { self.receiver.clone() } - async fn handle_entropy_tick( - &mut self, - file_sink: &file_sink::FileSinkClient, - ) -> anyhow::Result<()> { + async fn handle_entropy_tick(&mut self) -> anyhow::Result<()> { let source_data = match Self::get_entropy(&self.client).await { Ok(data) => data, Err(err) => { @@ -177,7 +185,9 @@ impl EntropyGenerator { entropy.timestamp ); - file_sink.write(EntropyReportV1::from(entropy), []).await?; + self.file_sink + .write(EntropyReportV1::from(entropy), []) + .await?; Ok(()) } diff --git a/poc_entropy/src/main.rs b/poc_entropy/src/main.rs index 5843d933c..c30de50f2 100644 --- a/poc_entropy/src/main.rs +++ b/poc_entropy/src/main.rs @@ -1,11 +1,10 @@ -use anyhow::{Error, Result}; +use anyhow::Result; use chrono::Duration; use clap::Parser; use file_store::{file_sink, file_upload, FileType}; -use futures_util::TryFutureExt; use poc_entropy::{entropy_generator::EntropyGenerator, server::ApiServer, Settings}; use std::{net::SocketAddr, path}; -use tokio::{self, signal}; +use task_manager::TaskManager; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; const ENTROPY_SINK_ROLL_MINS: i64 = 2; @@ -57,53 +56,38 @@ impl Server { // Install the prometheus metrics exporter poc_metrics::start_metrics(&settings.metrics)?; - // configure shutdown trigger - let (shutdown_trigger, shutdown) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); - // Initialize uploader - let (file_upload_tx, file_upload_rx) = file_upload::message_channel(); - let file_upload = - file_upload::FileUpload::from_settings(&settings.output, file_upload_rx).await?; + let (file_upload, file_upload_server) = + file_upload::FileUpload::from_settings(&settings.output).await?; let store_base_path = path::Path::new(&settings.cache); - // entropy - let mut entropy_generator = EntropyGenerator::new(&settings.source).await?; - let entropy_watch = entropy_generator.receiver(); - - let (entropy_sink, mut entropy_sink_server) = file_sink::FileSinkBuilder::new( + let (entropy_sink, entropy_sink_server) = file_sink::FileSinkBuilder::new( FileType::EntropyReport, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_report_submission"), - shutdown.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .roll_time(Duration::minutes(ENTROPY_SINK_ROLL_MINS)) .create() .await?; + let entropy_generator = EntropyGenerator::new(&settings.source, entropy_sink).await?; + let entropy_watch = entropy_generator.receiver(); + // server let socket_addr: SocketAddr = settings.listen.parse()?; let api_server = ApiServer::new(socket_addr, entropy_watch).await?; tracing::info!("api listening on {}", api_server.socket_addr); - tokio::try_join!( - api_server.run(&shutdown), - entropy_generator - .run(entropy_sink, &shutdown) - .map_err(Error::from), - entropy_sink_server.run().map_err(Error::from), - file_upload.run(&shutdown).map_err(Error::from), - ) - .map(|_| ()) + TaskManager::builder() + .add_task(file_upload_server) + .add_task(entropy_sink_server) + .add_task(entropy_generator) + .add_task(api_server) + .start() + .await } } diff --git a/poc_entropy/src/server.rs b/poc_entropy/src/server.rs index be1843a26..5627f6482 100644 --- a/poc_entropy/src/server.rs +++ b/poc_entropy/src/server.rs @@ -1,9 +1,11 @@ use crate::entropy_generator::MessageReceiver; +use futures::future::LocalBoxFuture; use helium_proto::{ services::poc_entropy::{EntropyReqV1, PocEntropy, Server as GrpcServer}, EntropyReportV1, }; use std::net::SocketAddr; +use task_manager::ManagedTask; use tokio::time::Duration; use tonic::transport; @@ -11,6 +13,15 @@ struct EntropyServer { entropy_watch: MessageReceiver, } +impl ManagedTask for ApiServer { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } +} + #[tonic::async_trait] impl PocEntropy for EntropyServer { async fn entropy( @@ -41,7 +52,7 @@ impl ApiServer { }) } - pub async fn run(self, shutdown: &triggered::Listener) -> anyhow::Result<()> { + pub async fn run(self, shutdown: triggered::Listener) -> anyhow::Result<()> { tracing::info!(listen = self.socket_addr.to_string(), "starting"); transport::Server::builder() .http2_keepalive_interval(Some(Duration::from_secs(250))) diff --git a/price/Cargo.toml b/price/Cargo.toml index fde5da4dc..faebb88ea 100644 --- a/price/Cargo.toml +++ b/price/Cargo.toml @@ -21,6 +21,7 @@ tracing-subscriber = { workspace = true } metrics = {workspace = true } metrics-exporter-prometheus = { workspace = true } tokio = { workspace = true } +tokio-util = { workspace = true } chrono = { workspace = true } helium-proto = { workspace = true } file-store = { path = "../file_store" } @@ -30,3 +31,4 @@ solana-client = {workspace = true} solana-sdk = {workspace = true} price-oracle = {workspace = true} anchor-lang = {workspace = true} +task-manager = {path = "../task_manager"} diff --git a/price/src/main.rs b/price/src/main.rs index b90a05572..999657623 100644 --- a/price/src/main.rs +++ b/price/src/main.rs @@ -1,12 +1,11 @@ -use anyhow::{Error, Result}; +use anyhow::Result; use chrono::Duration; use clap::Parser; use file_store::{file_sink, file_upload, FileType}; -use futures_util::TryFutureExt; use helium_proto::BlockchainTokenTypeV1; use price::{cli::check, PriceGenerator, Settings}; use std::path::{self, PathBuf}; -use tokio::{self, signal}; +use task_manager::TaskManager; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; const PRICE_SINK_ROLL_MINS: i64 = 3; @@ -80,62 +79,42 @@ impl Server { // Install the prometheus metrics exporter poc_metrics::start_metrics(&settings.metrics)?; - // configure shutdown trigger - let (shutdown_trigger, shutdown) = triggered::trigger(); - - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); - // Initialize uploader - let (file_upload_tx, file_upload_rx) = file_upload::message_channel(); - let file_upload = - file_upload::FileUpload::from_settings(&settings.output, file_upload_rx).await?; + let (file_upload, file_upload_server) = + file_upload::FileUpload::from_settings(&settings.output).await?; let store_base_path = path::Path::new(&settings.cache); - // price generators - let mut hnt_price_generator = - PriceGenerator::new(settings, BlockchainTokenTypeV1::Hnt).await?; - let mut mobile_price_generator = - PriceGenerator::new(settings, BlockchainTokenTypeV1::Mobile).await?; - let mut iot_price_generator = - PriceGenerator::new(settings, BlockchainTokenTypeV1::Iot).await?; - let mut hst_price_generator = - PriceGenerator::new(settings, BlockchainTokenTypeV1::Hst).await?; - - let (price_sink, mut price_sink_server) = file_sink::FileSinkBuilder::new( + let (price_sink, price_sink_server) = file_sink::FileSinkBuilder::new( FileType::PriceReport, store_base_path, concat!(env!("CARGO_PKG_NAME"), "_report_submission"), - shutdown.clone(), ) - .deposits(Some(file_upload_tx.clone())) + .file_upload(Some(file_upload.clone())) .roll_time(Duration::minutes(PRICE_SINK_ROLL_MINS)) .create() .await?; - tokio::try_join!( - hnt_price_generator - .run(price_sink.clone(), &shutdown) - .map_err(Error::from), - mobile_price_generator - .run(price_sink.clone(), &shutdown) - .map_err(Error::from), - iot_price_generator - .run(price_sink.clone(), &shutdown) - .map_err(Error::from), - hst_price_generator - .run(price_sink, &shutdown) - .map_err(Error::from), - price_sink_server.run().map_err(Error::from), - file_upload.run(&shutdown).map_err(Error::from), - ) - .map(|_| ()) + // price generators + let hnt_price_generator = + PriceGenerator::new(settings, BlockchainTokenTypeV1::Hnt, price_sink.clone()).await?; + let mobile_price_generator = + PriceGenerator::new(settings, BlockchainTokenTypeV1::Mobile, price_sink.clone()) + .await?; + let iot_price_generator = + PriceGenerator::new(settings, BlockchainTokenTypeV1::Iot, price_sink.clone()).await?; + let hst_price_generator = + PriceGenerator::new(settings, BlockchainTokenTypeV1::Hst, price_sink).await?; + + TaskManager::builder() + .add_task(file_upload_server) + .add_task(price_sink_server) + .add_task(hnt_price_generator) + .add_task(mobile_price_generator) + .add_task(iot_price_generator) + .add_task(hst_price_generator) + .start() + .await } } diff --git a/price/src/price_generator.rs b/price/src/price_generator.rs index c3a5bf805..7fb46c847 100644 --- a/price/src/price_generator.rs +++ b/price/src/price_generator.rs @@ -3,13 +3,14 @@ use anchor_lang::AccountDeserialize; use anyhow::{anyhow, Error, Result}; use chrono::{DateTime, Duration, TimeZone, Utc}; use file_store::file_sink; -use futures::TryFutureExt; +use futures::{future::LocalBoxFuture, TryFutureExt}; use helium_proto::{BlockchainTokenTypeV1, PriceReportV1}; use price_oracle::{calculate_current_price, PriceOracleV0}; use serde::{Deserialize, Serialize}; use solana_client::nonblocking::rpc_client::RpcClient; use solana_sdk::pubkey::Pubkey as SolPubkey; use std::{path::PathBuf, str::FromStr}; +use task_manager::ManagedTask; use tokio::{fs, time}; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -38,6 +39,16 @@ pub struct PriceGenerator { default_price: Option, stale_price_duration: Duration, latest_price_file: PathBuf, + file_sink: file_sink::FileSinkClient, +} + +impl ManagedTask for PriceGenerator { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } } impl From for PriceReportV1 { @@ -68,7 +79,11 @@ impl TryFrom for Price { } impl PriceGenerator { - pub async fn new(settings: &Settings, token_type: BlockchainTokenTypeV1) -> Result { + pub async fn new( + settings: &Settings, + token_type: BlockchainTokenTypeV1, + file_sink: file_sink::FileSinkClient, + ) -> Result { let client = RpcClient::new(settings.source.clone()); Ok(Self { last_price_opt: None, @@ -80,20 +95,14 @@ impl PriceGenerator { stale_price_duration: settings.stale_price_duration(), latest_price_file: PathBuf::from_str(&settings.cache)? .join(format!("{token_type:?}.latest")), + file_sink, }) } - pub async fn run( - &mut self, - file_sink: file_sink::FileSinkClient, - shutdown: &triggered::Listener, - ) -> Result<()> { + pub async fn run(mut self, shutdown: triggered::Listener) -> Result<()> { match (self.key, self.default_price) { - (Some(key), _) => self.run_with_key(key, file_sink, shutdown).await, - (None, Some(defaut_price)) => { - self.run_with_default(defaut_price, file_sink, shutdown) - .await - } + (Some(key), _) => self.run_with_key(key, &shutdown).await, + (None, Some(defaut_price)) => self.run_with_default(defaut_price, &shutdown).await, _ => { tracing::warn!( "stopping price generator for {:?}, not configured", @@ -107,7 +116,6 @@ impl PriceGenerator { async fn run_with_default( &self, default_price: u64, - file_sink: file_sink::FileSinkClient, shutdown: &triggered::Listener, ) -> Result<()> { tracing::info!( @@ -123,7 +131,7 @@ impl PriceGenerator { let price = Price::new(Utc::now(), default_price, self.token_type); let price_report = PriceReportV1::from(price); tracing::info!("updating {:?} with default price: {}", self.token_type, default_price); - file_sink.write(price_report, []).await?; + self.file_sink.write(price_report, []).await?; } } } @@ -132,12 +140,7 @@ impl PriceGenerator { Ok(()) } - async fn run_with_key( - &mut self, - key: SolPubkey, - file_sink: file_sink::FileSinkClient, - shutdown: &triggered::Listener, - ) -> Result<()> { + async fn run_with_key(&mut self, key: SolPubkey, shutdown: &triggered::Listener) -> Result<()> { tracing::info!("starting price generator for {:?}", self.token_type); let mut trigger = time::interval(self.interval_duration); self.last_price_opt = self.read_price_file().await; @@ -145,7 +148,7 @@ impl PriceGenerator { loop { tokio::select! { _ = shutdown.clone() => break, - _ = trigger.tick() => self.handle(&key, &file_sink).await?, + _ = trigger.tick() => self.handle(&key).await?, } } @@ -153,11 +156,7 @@ impl PriceGenerator { Ok(()) } - async fn handle( - &mut self, - key: &SolPubkey, - file_sink: &file_sink::FileSinkClient, - ) -> Result<()> { + async fn handle(&mut self, key: &SolPubkey) -> Result<()> { let price_opt = match get_price(&self.client, key, self.token_type).await { Ok(new_price) => { tracing::info!( @@ -212,7 +211,7 @@ impl PriceGenerator { if let Some(price) = price_opt { let price_report = PriceReportV1::from(price); tracing::debug!("price_report: {:?}", price_report); - file_sink.write(price_report, []).await?; + self.file_sink.write(price_report, []).await?; } Ok(()) diff --git a/reward_index/Cargo.toml b/reward_index/Cargo.toml index 3ba6cb6a9..5c2ca846f 100644 --- a/reward_index/Cargo.toml +++ b/reward_index/Cargo.toml @@ -41,3 +41,5 @@ rust_decimal_macros = {workspace = true} tonic = {workspace = true} rand = {workspace = true} async-trait = {workspace = true} +tokio-util = { workspace = true } +task-manager = { path = "../task_manager" } diff --git a/reward_index/src/indexer.rs b/reward_index/src/indexer.rs index a5b559fa7..89c16a8b0 100644 --- a/reward_index/src/indexer.rs +++ b/reward_index/src/indexer.rs @@ -4,7 +4,7 @@ use chrono::Utc; use file_store::{ file_info_poller::FileInfoStream, reward_manifest::RewardManifest, FileInfo, FileStore, }; -use futures::{stream, StreamExt, TryStreamExt}; +use futures::{future::LocalBoxFuture, stream, StreamExt, TryStreamExt}; use helium_crypto::PublicKeyBinary; use helium_proto::{ services::poc_lora::{iot_reward_share::Reward as IotReward, IotRewardShare}, @@ -14,6 +14,7 @@ use helium_proto::{ use poc_metrics::record_duration; use sqlx::{Pool, Postgres, Transaction}; use std::{collections::HashMap, str::FromStr}; +use task_manager::ManagedTask; use tokio::sync::mpsc::Receiver; pub struct Indexer { @@ -21,6 +22,7 @@ pub struct Indexer { verifier_store: FileStore, mode: settings::Mode, op_fund_key: String, + receiver: Receiver>, } #[derive(sqlx::Type, Debug, Clone, PartialEq, Eq, Hash)] @@ -38,8 +40,21 @@ pub struct RewardKey { reward_type: RewardType, } +impl ManagedTask for Indexer { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + Box::pin(self.run(shutdown)) + } +} + impl Indexer { - pub async fn new(settings: &Settings, pool: Pool) -> Result { + pub async fn new( + settings: &Settings, + pool: Pool, + receiver: Receiver>, + ) -> Result { Ok(Self { mode: settings.mode, verifier_store: FileStore::from_settings(&settings.verifier).await?, @@ -50,14 +65,11 @@ impl Indexer { .ok_or_else(|| anyhow!("operation fund key is required for IOT mode"))?, settings::Mode::Mobile => String::new(), }, + receiver, }) } - pub async fn run( - &mut self, - shutdown: triggered::Listener, - mut receiver: Receiver>, - ) -> Result<()> { + pub async fn run(mut self, shutdown: triggered::Listener) -> Result<()> { tracing::info!(mode = self.mode.to_string(), "starting index"); loop { @@ -66,7 +78,7 @@ impl Indexer { tracing::info!("Indexer shutting down"); return Ok(()); } - msg = receiver.recv() => if let Some(file_info_stream) = msg { + msg = self.receiver.recv() => if let Some(file_info_stream) = msg { let key = &file_info_stream.file_info.key.clone(); tracing::info!(file = %key, "Processing reward file"); let mut txn = self.pool.begin().await?; diff --git a/reward_index/src/main.rs b/reward_index/src/main.rs index 9c33a6849..c20e35161 100644 --- a/reward_index/src/main.rs +++ b/reward_index/src/main.rs @@ -5,10 +5,9 @@ use file_store::{ file_info_poller::LookbackBehavior, file_source, reward_manifest::RewardManifest, FileStore, FileType, }; -use futures_util::TryFutureExt; use reward_index::{settings::Settings, telemetry, Indexer}; use std::path::PathBuf; -use tokio::signal; +use task_manager::TaskManager; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; #[derive(Debug, clap::Parser)] @@ -57,30 +56,17 @@ impl Server { // Install the prometheus metrics exporter poc_metrics::start_metrics(&settings.metrics)?; - // - // Configure shutdown trigger - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); // Create database pool let app_name = format!("{}_{}", settings.mode, env!("CARGO_PKG_NAME")); - let (pool, db_join_handle) = settings - .database - .connect(&app_name, shutdown_listener.clone()) - .await?; + let pool = settings.database.connect(&app_name).await?; sqlx::migrate!().run(&pool).await?; telemetry::initialize(&pool).await?; let file_store = FileStore::from_settings(&settings.verifier).await?; - let (receiver, source_join_handle) = file_source::continuous_source::() + let (receiver, receiver_server) = file_source::continuous_source::() .db(pool.clone()) .store(file_store) .file_type(FileType::RewardManifest) @@ -91,20 +77,16 @@ impl Server { )) .poll_duration(settings.interval()) .offset(settings.interval() * 2) - .build()? - .start(shutdown_listener.clone()) - .await?; + .create()?; // Reward server - let mut indexer = Indexer::new(settings, pool).await?; + let indexer = Indexer::new(settings, pool, receiver).await?; - tokio::try_join!( - db_join_handle.map_err(anyhow::Error::from), - source_join_handle.map_err(anyhow::Error::from), - indexer.run(shutdown_listener, receiver), - )?; - - Ok(()) + TaskManager::builder() + .add_task(receiver_server) + .add_task(indexer) + .start() + .await } } diff --git a/solana/Cargo.toml b/solana/Cargo.toml index 3ee0f32be..cc41f1800 100644 --- a/solana/Cargo.toml +++ b/solana/Cargo.toml @@ -7,6 +7,7 @@ authors.workspace = true license.workspace = true [dependencies] +anyhow = { workspace = true } async-trait = {workspace = true} anchor-lang = {workspace = true} anchor-client = {workspace = true} @@ -22,7 +23,9 @@ solana-client = {workspace = true} solana-program = {workspace = true} solana-sdk = {workspace = true} spl-token = {workspace = true} +task-manager = { path = "../task_manager" } thiserror = {workspace = true} tokio = {workspace = true} +tokio-util = { workspace = true } tracing = {workspace = true} triggered = {workspace = true} diff --git a/solana/src/balance_monitor.rs b/solana/src/balance_monitor.rs index 257370518..8e388ca4c 100644 --- a/solana/src/balance_monitor.rs +++ b/solana/src/balance_monitor.rs @@ -1,30 +1,54 @@ use crate::{SolanaRpc, SolanaRpcError}; +use futures::{future::LocalBoxFuture, TryFutureExt}; use solana_sdk::{pubkey::Pubkey, signature::Keypair, signer::Signer}; use std::{sync::Arc, time::Duration}; +use task_manager::ManagedTask; // Check balance every 12 hours const DURATION: Duration = Duration::from_secs(43_200); -pub async fn start( - app_account: &str, - solana: Option>, - shutdown: triggered::Listener, -) -> Result>, SolanaRpcError> -{ - Ok(match solana { - None => Box::pin(async move { Ok(()) }), - Some(rpc_client) => { - let Ok(keypair) = Keypair::from_bytes(&rpc_client.keypair) else { - tracing::error!("sol monitor: keypair failed to deserialize"); - return Err(SolanaRpcError::InvalidKeypair) - }; - let app_metric_name = format!("{app_account}-sol-balance"); - let handle = tokio::spawn(async move { - run(app_metric_name, rpc_client, keypair.pubkey(), shutdown).await - }); - Box::pin(handle) +pub enum BalanceMonitor { + Solana(String, Arc, Pubkey), + Noop, +} + +impl BalanceMonitor { + pub fn new( + app_account: &str, + solana: Option>, + ) -> Result> { + match solana { + None => Ok(BalanceMonitor::Noop), + Some(rpc_client) => { + let Ok(keypair) = Keypair::from_bytes(&rpc_client.keypair) else { + tracing::error!("sol monitor: keypair failed to deserialize"); + return Err(Box::new(SolanaRpcError::InvalidKeypair)) + }; + let app_metric_name = format!("{app_account}-sol-balance"); + + Ok(BalanceMonitor::Solana( + app_metric_name, + rpc_client, + keypair.pubkey(), + )) + } } - }) + } +} + +impl ManagedTask for BalanceMonitor { + fn start_task( + self: Box, + shutdown: triggered::Listener, + ) -> LocalBoxFuture<'static, anyhow::Result<()>> { + match *self { + Self::Noop => Box::pin(async move { Ok(()) }), + Self::Solana(metric, solana, pubkey) => { + let handle = tokio::spawn(run(metric, solana, pubkey, shutdown)); + Box::pin(handle.map_err(anyhow::Error::from)) + } + } + } } async fn run( @@ -33,16 +57,15 @@ async fn run( service_pubkey: Pubkey, shutdown: triggered::Listener, ) { + tracing::info!("starting sol monitor"); + let mut trigger = tokio::time::interval(DURATION); loop { let shutdown = shutdown.clone(); tokio::select! { - _ = shutdown => { - tracing::info!("sol monitor: shutting down"); - break - } + _ = shutdown => break, _ = trigger.tick() => { match solana.provider.get_balance(&service_pubkey).await { Ok(balance) => metrics::gauge!(metric_name.clone(), balance as f64), @@ -51,4 +74,5 @@ async fn run( } } } + tracing::info!("stopping sol monitor") } diff --git a/task_manager/src/lib.rs b/task_manager/src/lib.rs index 1336b9419..b2402010d 100644 --- a/task_manager/src/lib.rs +++ b/task_manager/src/lib.rs @@ -22,7 +22,7 @@ pub struct TaskManagerBuilder { } pub struct StopableLocalFuture { - shutdown_listener: triggered::Listener, + shutdown_trigger: triggered::Trigger, future: LocalBoxFuture<'static, anyhow::Result<()>>, } @@ -50,6 +50,12 @@ where } } +impl Default for TaskManager { + fn default() -> Self { + Self::new() + } +} + impl TaskManager { pub fn new() -> Self { Self { tasks: Vec::new() } @@ -66,29 +72,22 @@ impl TaskManager { pub async fn start(self) -> anyhow::Result<()> { let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate()).unwrap(); - let (shutdown_trigger, shutdown_listener) = triggered::trigger(); - let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())?; - tokio::spawn(async move { - tokio::select! { - _ = sigterm.recv() => shutdown_trigger.trigger(), - _ = signal::ctrl_c() => shutdown_trigger.trigger(), - } - }); + let shutdown_triggers = create_triggers(self.tasks.len()); - let mut futures = start_futures(shutdown_listener.clone(), self.tasks); + let mut futures = start_futures(shutdown_triggers.clone(), self.tasks); - // let mut shutdown_listener = - // futures::future::select(Box::pin(sigterm.recv()), Box::pin(signal::ctrl_c())); + let mut shutdown = + futures::future::select(Box::pin(sigterm.recv()), Box::pin(signal::ctrl_c())); loop { - if futures.len() == 0 { + if futures.is_empty() { break; } let mut select = select_all(futures.into_iter()); tokio::select! { - _ = &mut shutdown_listener => { + _ = &mut shutdown => { return stop_all(select.into_inner()).await; } (result, _index, remaining) = &mut select => match result { @@ -108,7 +107,7 @@ impl TaskManager { } impl TaskManagerBuilder { - pub fn add(mut self, task: impl ManagedTask + 'static) -> Self { + pub fn add_task(mut self, task: impl ManagedTask + 'static) -> Self { self.tasks.push(Box::new(task)); self } @@ -120,151 +119,162 @@ impl TaskManagerBuilder { } fn start_futures( - shutdown_listener: triggered::Listener, + shutdown_triggers: Vec<(triggered::Trigger, triggered::Listener)>, tasks: Vec>, ) -> Vec { - tasks + shutdown_triggers .into_iter() - .map(|task| StopableLocalFuture { - shutdown_listener: shutdown_listener.clone(), - future: task.start_task(shutdown_listener), - }) + .zip(tasks.into_iter()) + .map( + |((shutdown_trigger, shutdown_listener), task)| StopableLocalFuture { + shutdown_trigger, + future: task.start_task(shutdown_listener), + }, + ) .collect() } async fn stop_all(futures: Vec) -> anyhow::Result<()> { futures::stream::iter(futures.into_iter().rev()) .fold(Ok(()), |last_result, local| async move { - local.shutdown_listener.tri.cancel(); + local.shutdown_trigger.trigger(); let result = local.future.await; last_result.and(result) }) .await } -#[cfg(test)] -mod tests { - use super::*; - use anyhow::anyhow; - use futures::TryFutureExt; - use tokio::sync::mpsc; - - struct TestTask { - id: u64, - delay: u64, - result: anyhow::Result<()>, - sender: mpsc::Sender, - } - - impl ManagedTask for TestTask { - fn start_task( - self: Box, - shutdown_listener: triggered::Listener, - ) -> LocalBoxFuture<'static, anyhow::Result<()>> { - let handle = tokio::spawn(async move { - tokio::select! { - _ = shutdown_listener.cancelled() => (), - _ = tokio::time::sleep(std::time::Duration::from_millis(self.delay)) => (), - } - - self.sender.send(self.id).await.expect("unable to send"); - self.result - }); - - Box::pin( - handle - .map_err(|err| err.into()) - .and_then(|result| async move { result }), - ) - } - } - - #[tokio::test] - async fn stop_when_all_tasks_have_completed() { - let (sender, mut receiver) = mpsc::channel(5); - - let result = TaskManager::builder() - .add(TestTask { - id: 1, - delay: 50, - result: Ok(()), - sender: sender.clone(), - }) - .add(TestTask { - id: 2, - delay: 100, - result: Ok(()), - sender: sender.clone(), - }) - .start() - .await; - - assert_eq!(Some(1), receiver.recv().await); - assert_eq!(Some(2), receiver.recv().await); - assert!(result.is_ok()); - } - - #[tokio::test] - async fn will_stop_all_in_reverse_order_after_error() { - let (sender, mut receiver) = mpsc::channel(5); - - let result = TaskManager::builder() - .add(TestTask { - id: 1, - delay: 1000, - result: Ok(()), - sender: sender.clone(), - }) - .add(TestTask { - id: 2, - delay: 50, - result: Err(anyhow!("error")), - sender: sender.clone(), - }) - .add(TestTask { - id: 3, - delay: 1000, - result: Ok(()), - sender: sender.clone(), - }) - .start() - .await; - - assert_eq!(Some(2), receiver.recv().await); - assert_eq!(Some(3), receiver.recv().await); - assert_eq!(Some(1), receiver.recv().await); - assert_eq!("error", result.unwrap_err().to_string()); - } - - #[tokio::test] - async fn will_return_first_error_returned() { - let (sender, mut receiver) = mpsc::channel(5); - - let result = TaskManager::builder() - .add(TestTask { - id: 1, - delay: 1000, - result: Ok(()), - sender: sender.clone(), - }) - .add(TestTask { - id: 2, - delay: 50, - result: Err(anyhow!("error")), - sender: sender.clone(), - }) - .add(TestTask { - id: 3, - delay: 200, - result: Err(anyhow!("second")), - sender: sender.clone(), - }) - .start() - .await; - - assert_eq!(Some(2), receiver.recv().await); - assert_eq!(Some(3), receiver.recv().await); - assert_eq!(Some(1), receiver.recv().await); - assert_eq!("error", result.unwrap_err().to_string()); - } +fn create_triggers(n: usize) -> Vec<(triggered::Trigger, triggered::Listener)> { + let (shutdown_trigger, shutdown_listener) = triggered::trigger(); + (0..n).fold(Vec::new(), |mut vec, _| { + vec.push((shutdown_trigger.clone(), shutdown_listener.clone())); + vec + }) } + +// #[cfg(test)] +// mod tests { +// use super::*; +// use anyhow::anyhow; +// use futures::TryFutureExt; +// use tokio::sync::mpsc; + +// struct TestTask { +// id: u64, +// delay: u64, +// result: anyhow::Result<()>, +// sender: mpsc::Sender, +// } + +// impl ManagedTask for TestTask { +// fn start_task( +// self: Box, +// shutdown_listener: triggered::Listener, +// ) -> LocalBoxFuture<'static, anyhow::Result<()>> { +// let handle = tokio::spawn(async move { +// tokio::select! { +// _ = shutdown_listener.cancelled() => (), +// _ = tokio::time::sleep(std::time::Duration::from_millis(self.delay)) => (), +// } + +// self.sender.send(self.id).await.expect("unable to send"); +// self.result +// }); + +// Box::pin( +// handle +// .map_err(|err| err.into()) +// .and_then(|result| async move { result }), +// ) +// } +// } + +// #[tokio::test] +// async fn stop_when_all_tasks_have_completed() { +// let (sender, mut receiver) = mpsc::channel(5); + +// let result = TaskManager::builder() +// .add_task(TestTask { +// id: 1, +// delay: 50, +// result: Ok(()), +// sender: sender.clone(), +// }) +// .add_task(TestTask { +// id: 2, +// delay: 100, +// result: Ok(()), +// sender: sender.clone(), +// }) +// .start() +// .await; + +// assert_eq!(Some(1), receiver.recv().await); +// assert_eq!(Some(2), receiver.recv().await); +// assert!(result.is_ok()); +// } + +// #[tokio::test] +// async fn will_stop_all_in_reverse_order_after_error() { +// let (sender, mut receiver) = mpsc::channel(5); + +// let result = TaskManager::builder() +// .add_task(TestTask { +// id: 1, +// delay: 1000, +// result: Ok(()), +// sender: sender.clone(), +// }) +// .add_task(TestTask { +// id: 2, +// delay: 50, +// result: Err(anyhow!("error")), +// sender: sender.clone(), +// }) +// .add_task(TestTask { +// id: 3, +// delay: 1000, +// result: Ok(()), +// sender: sender.clone(), +// }) +// .start() +// .await; + +// assert_eq!(Some(2), receiver.recv().await); +// assert_eq!(Some(3), receiver.recv().await); +// assert_eq!(Some(1), receiver.recv().await); +// assert_eq!("error", result.unwrap_err().to_string()); +// } + +// #[tokio::test] +// async fn will_return_first_error_returned() { +// let (sender, mut receiver) = mpsc::channel(5); + +// let result = TaskManager::builder() +// .add_task(TestTask { +// id: 1, +// delay: 1000, +// result: Ok(()), +// sender: sender.clone(), +// }) +// .add_task(TestTask { +// id: 2, +// delay: 50, +// result: Err(anyhow!("error")), +// sender: sender.clone(), +// }) +// .add_task(TestTask { +// id: 3, +// delay: 200, +// result: Err(anyhow!("second")), +// sender: sender.clone(), +// }) +// .start() +// .await; + +// assert_eq!(Some(2), receiver.recv().await); +// assert_eq!(Some(3), receiver.recv().await); +// assert_eq!(Some(1), receiver.recv().await); +// assert_eq!("error", result.unwrap_err().to_string()); +// } +// }