feat(riklet): graceful shutdown (#90)

* Trigger graceful on CTRL+C (SIGTERM / SIGINT)
* Clean up iptables on shutdown of riklet

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(())
+    }
 }

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(())
 }

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

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(())
     }
 }