diff --git a/riklet/src/core.rs b/riklet/src/core.rs index 6c2911c..d1d7275 100644 --- a/riklet/src/core.rs +++ b/riklet/src/core.rs @@ -217,4 +217,18 @@ impl Riklet { network: global_runtime_network, }) } + + pub async fn shutdown(&mut self) -> Result<()> { + info!("Shutting down ..."); + + // Stop and clean all runtime + for (_, runtime) in &mut self.runtimes { + runtime + .down() + .await + .map_err(RikletError::RuntimeManagerError)?; + } + + Ok(()) + } } diff --git a/riklet/src/main.rs b/riklet/src/main.rs index 6a50849..ce1fbbe 100644 --- a/riklet/src/main.rs +++ b/riklet/src/main.rs @@ -8,9 +8,11 @@ mod runtime; mod structs; use crate::core::Riklet; -use anyhow::Result; +use anyhow::{Context, Result}; -use tracing::{error, metadata::LevelFilter}; +use tokio::signal::ctrl_c; +use tokio::signal::unix::{signal, SignalKind}; +use tracing::{error, info, metadata::LevelFilter}; use tracing_subscriber::{ fmt, prelude::__tracing_subscriber_SubscriberExt, util::SubscriberInitExt, EnvFilter, }; @@ -40,6 +42,34 @@ pub fn init_logger() -> Result<()> { Ok(()) } +async fn serve() -> Result<()> { + let mut riklet = Riklet::new().await.unwrap_or_else(|e| { + error!( + "An error occured during the bootstraping process of the Riklet. {}", + e + ); + std::process::exit(2); + }); + + // Stream of SIGTERM signals. + let mut signals = signal(SignalKind::terminate())?; + + tokio::select! { + _ = riklet.run() => {}, + _ = ctrl_c() => { + info!("Receive SIGINT signal."); + }, + _ = signals.recv() => { + info!("Receive SIGTERM signal."); + } + } + + riklet + .shutdown() + .await + .context("Could not graceful shutdown riklet") +} + #[tokio::main] async fn main() -> Result<()> { init_logger()?; @@ -50,17 +80,9 @@ async fn main() -> Result<()> { std::process::exit(1); } - Riklet::new() - .await - .unwrap_or_else(|e| { - error!( - "An error occured during the bootstraping process of the Riklet. {}", - e - ); - std::process::exit(2); - }) - .run() - .await?; + serve().await?; + + info!("Riklet stopped"); Ok(()) } diff --git a/riklet/src/runtime/function_runtime.rs b/riklet/src/runtime/function_runtime.rs index 6ea0d05..f646989 100644 --- a/riklet/src/runtime/function_runtime.rs +++ b/riklet/src/runtime/function_runtime.rs @@ -145,16 +145,11 @@ impl Runtime for FunctionRuntime { } }?; - machine - .stop() - .await - .map_err(RuntimeError::FirecrackerError)?; - debug!("microVM properly stopped"); - machine .kill() .await .map_err(RuntimeError::FirecrackerError)?; + debug!("microVM properly stopped"); debug!("Destroying function runtime network"); self.network diff --git a/riklet/src/runtime/network/function_network.rs b/riklet/src/runtime/network/function_network.rs index d49f811..66e9c3d 100644 --- a/riklet/src/runtime/network/function_network.rs +++ b/riklet/src/runtime/network/function_network.rs @@ -1,7 +1,8 @@ use async_trait::async_trait; +use ipnetwork::Ipv4Network; use proto::worker::InstanceScheduling; use std::net::Ipv4Addr; -use tracing::debug; +use tracing::{debug, error}; use crate::constants::DEFAULT_FIRECRACKER_NETWORK_MASK; use crate::net_utils::{self, get_iptables_riklet_chain}; @@ -120,6 +121,21 @@ impl FunctionRuntimeNetwork { } Ok(()) } + + /// Release allocated IPs + fn release_network(&self) -> Result<()> { + debug!("Release subnet IPs"); + + let subnet = Ipv4Network::new(self.host_ip, 30) + .map_err(|e| NetworkError::Error(format!("Fail to get function subnet {}", e)))?; + + match IP_ALLOCATOR.lock() { + Ok(mut ip_allocator) => ip_allocator.free_subnet(subnet), + Err(e) => error!("Couldn't free subnet {}, reason: {}", subnet, e), + } + + Ok(()) + } } #[async_trait] @@ -160,6 +176,7 @@ impl RuntimeNetwork for FunctionRuntimeNetwork { async fn destroy(&mut self) -> Result<()> { debug!("Destroy function network"); self.down_routing()?; + self.release_network()?; Ok(()) } }