taskman up mobile config

helium · Aug 23, 2023 · 0b953ce · 0b953ce
1 parent 3a3552a
commit 0b953ce
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 24 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/mobile_config/Cargo.toml b/mobile_config/Cargo.toml
@@ -34,7 +34,9 @@ retainer = {workspace = true}
 thiserror = {workspace = true}
 tokio = {workspace = true}
 tokio-stream = {workspace = true}
+tokio-util = { workspace = true }
 tonic = {workspace = true}
 tracing = {workspace = true}
 tracing-subscriber = {workspace = true}
 triggered = {workspace = true}
+task-manager = { path = "../task_manager" }
diff --git 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,16 +59,6 @@ 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)?;
 
@@ -96,18 +87,43 @@ impl Daemon {
             settings.signing_keypair()?,
         );
 
-        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)
-            .await?;
-
-        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<Self>,
+        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
+        })
     }
 }