feat(node-agent): implement gRPC server

Signed-off-by: Martin Moreira de Jesus <martin.moreiradj@icloud.com>
Signed-off-by: Noé Tarbouriech <noe.tarbouriech@etu.umontpellier.fr>
Co-Authored-By: Noé Tarbouriech <noe.tarbouriech@etu.umontpellier.fr>

Author: Martin Moreira de Jesus

node-agent/src/args.rs

@@ -9,9 +9,6 @@ pub struct CliArguments {
     #[arg(long, default_value_t = 3, env)]
     pub lifecycle_retries: i32,
 
-    /// The port to start the node agent on.
-    #[arg(long, default_value_t = 50052, env)]
-    pub node_agent_port: u16,
 
     /// The address of the scheduler to connect the node agent to.
     #[arg(long, default_value = "[::]", env)]
@@ -21,6 +18,14 @@ pub struct CliArguments {
     #[arg(long, default_value_t = 50051, env)]
     pub scheduler_port: u16,
 
+    /// The address of the node agent for the scheduler to connect to.
+    #[arg(long, default_value = "[::]", env)]
+    pub node_agent_address: String,
+
+    /// The port of the node agent for the scheduler to connect to.
+    #[arg(long, default_value_t = 50052, env)]
+    pub node_agent_port: u16,
+
     /// Verbosity level.
     #[command(flatten)]
     pub verbose: Verbosity<InfoLevel>,

node-agent/src/main.rs

@@ -1,15 +1,18 @@
 mod args;
 mod workload_manager;
 
+use std::process::exit;
+
 use clap::Parser;
-use orka_proto::{node_agent::Workload, scheduler_agent::{lifecycle_service_client::LifecycleServiceClient, ConnectionRequest, status_update_service_client::StatusUpdateServiceClient}};
-use tracing::{info, error};
+use orka_proto::scheduler_agent::{lifecycle_service_client::LifecycleServiceClient, ConnectionRequest, status_update_service_client::StatusUpdateServiceClient, DisconnectionNotice};
+use tokio::sync::mpsc::{Sender, Receiver, self};
+use tracing::{info, error, warn};
 use uuid::Uuid;
-use workload_manager::container::client::ContainerClient;
 use anyhow::Result;
 use tracing_log::AsTrace;
 
-use crate::{workload_manager::container::metrics::metrics::any_to_resource, args::CliArguments};
+use crate::workload_manager::grpc::server::GrpcServer;
+use crate::args::CliArguments;
 use crate::workload_manager::node::metrics::stream_node_status;
 
 async fn execute_node_lifecycle(
@@ -73,53 +76,8 @@ async fn execute_node_lifecycle(
     }
 }
 
-const CID: &str = "nginx";
-
 #[tokio::main]
 async fn main() {
-    let mut workload_manager = ContainerClient::new("/var/run/containerd/containerd.sock").await.unwrap();
-
-    let workload = Workload {
-        instance_id: CID.to_string(),
-        image: "docker.io/library/nginx:latest".to_string(),
-        environment: vec!["FOO=BAR".to_string()],
-        ..Default::default()
-    };
-
-    workload_manager.create(&workload).await.unwrap();
-
-    let response = workload_manager.info(CID).await.unwrap();
-
-    println!("{:?}", response);
-
-    let response = workload_manager.status(CID).await.unwrap();
-
-    println!("{:?}", response);
-
-    for _ in 0..10 {
-        let response = workload_manager.metrics(CID).await.unwrap();
-
-        for metric in response.into_inner().metrics {
-            let data = metric.data.unwrap();
-
-            println!("{:?}", any_to_resource(&data));
-        }
-
-        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
-    }
-
-    let signal = 9;
-    workload_manager.kill(CID, signal).await.unwrap();
-
-    workload_manager.wait(CID).await.unwrap();
-
-    workload_manager.cleanup(CID).await.unwrap();
-
-    match workload_manager.status(CID).await {
-        Ok(_) => panic!("Workload should not exist"),
-        Err(_) => println!("Workload does not exist"),
-    };
-
     let args = CliArguments::parse();
 
     tracing_subscriber::fmt()
@@ -139,22 +97,97 @@ async fn main() {
 
     let node_id = Uuid::new_v4();
 
-    let retries = args.lifecycle_retries;
+    info!("Node ID: {}", node_id);
 
-    for _ in 0..retries {
-        match execute_node_lifecycle(
-            node_id,
-            args.node_agent_port,
-            scheduler_connection_string.clone(),
-        )
-        .await
-        {
-            Ok(_) => {
-                // will never be reached
+    let (tx, mut rx): (Sender<i32>, Receiver<i32>) = mpsc::channel(1);
+    
+    let lifecycle_tx = tx.clone();
+    let lifecycle_connection_string = scheduler_connection_string.clone();
+
+    // join scheduler and stream node status to scheduler, retrying on failure
+    tokio::spawn(async move {
+        let retries = args.lifecycle_retries;
+
+        for _ in 0..retries {
+            match execute_node_lifecycle(
+                node_id,
+                args.node_agent_port,
+                lifecycle_connection_string.clone(),
+            )
+            .await
+            {
+                Ok(_) => {
+                    // will never be reached
+                }
+                Err(e) => {
+                    error!("Failed to execute node lifecycle: {:?}", e);
+                }
             }
+        }
+
+        error!("
+            Node lifecycle failed, initiating graceful shutdown"
+        );
+
+        let _ = lifecycle_tx.send(1).await;
+    });
+
+    // start grpc server
+    tokio::spawn(async move {
+        error!("Starting gRPC server on {}:{}",
+            args.node_agent_address,
+            args.node_agent_port
+        );
+
+        let grpc = GrpcServer::new(args.node_agent_address, args.node_agent_port);
+
+        let server = match grpc.map_err(|e| {
+            error!("Failed to create gRPC server: {:?}", e);
+        }) {
+            Ok(server) => server,
             Err(e) => {
-                error!("Failed to execute node lifecycle: {:?}", e);
+                error!("Failed to create gRPC server: {:?}", e);
+                let _ = tx.send(1).await;
+                exit(1);
             }
+        };
+
+        match server.start_server().await {
+            Ok(_) => {}
+            Err(e) => {
+                error!("Failed to start gRPC server: {:?}", e);
+                let _ = tx.send(1).await;
+                exit(1);
+            }
+        };
+    });
+
+    let _ = rx.recv().await;
+
+    info!("Exiting");
+
+    info!("Trying to quit cluster");
+
+    match LifecycleServiceClient::connect(scheduler_connection_string).await {
+        Err(error) => {
+            warn!("Failed connecting to scheduler when exiting {:?}", error)
+        }
+        Ok(mut client) => {
+            match client
+                .leave_cluster(DisconnectionNotice {
+                    id: node_id.to_string(),
+                })
+                .await
+            {
+                Ok(_) => {
+                    info!("Successfully left cluster")
+                }
+                Err(error) => {
+                    warn!("Failed existing cluster {:?}", error)
+                }
+            };
         }
     }
+
+    exit(1);
 }

node-agent/src/workload_manager/grpc/mod.rs

@@ -0,0 +1,2 @@
+pub mod scheduler_workload_service;
+pub mod server;

node-agent/src/workload_manager/grpc/scheduler_workload_service.rs

@@ -0,0 +1,164 @@
+use crate::workload_manager::container::client::ContainerClient;
+use crate::workload_manager::container::error::into_tonic_status;
+use crate::workload_manager::container::metrics::metrics::any_to_resource;
+use orka_proto::node_agent::workload_service_server::WorkloadService;
+use orka_proto::node_agent::workload_signal::Signal;
+use orka_proto::node_agent::{workload_status, Empty, Workload, WorkloadSignal, WorkloadStatus};
+use tokio_stream::wrappers::ReceiverStream;
+use tonic::{Request, Response, Result, Status};
+use tracing::debug;
+
+pub struct WorkloadSvc {}
+
+impl WorkloadSvc {
+    /// Create a new `WorkloadService` gRPC service manager.
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+#[tonic::async_trait]
+impl WorkloadService for WorkloadSvc {
+    type CreateStream = ReceiverStream<Result<WorkloadStatus>>;
+
+    async fn create(
+        &self,
+        request: Request<Workload>,
+    ) -> Result<Response<Self::CreateStream>, Status> {
+        let workload = request.into_inner();
+        let mut client = match ContainerClient::new("/var/run/containerd/containerd.sock").await {
+            Ok(x) => x,
+            Err(e) => {
+                return Err(Status::new(
+                    tonic::Code::Internal,
+                    format!("Failed to create container client{:?}", e),
+                ))
+            }
+        };
+
+        let container = match client.create(&workload).await {
+            Ok(x) => x,
+            Err(e) => {
+                return Err(Status::new(
+                    tonic::Code::Internal,
+                    format!("Failed to create container: {:?}", e),
+                ))
+            }
+        }
+        .into_inner();
+
+        let (tx, rx) = tokio::sync::mpsc::channel(4);
+        tokio::spawn(async move {
+            loop {
+                let metrics = match client.metrics(&container.container_id).await {
+                    Ok(x) => x,
+                    Err(e) => {
+                        return Status::new(
+                            tonic::Code::Internal,
+                            format!("Failed to get container info: {:?}", e),
+                        )
+                    }
+                }
+                .into_inner()
+                .metrics;
+
+                let status = match client.status(&container.container_id).await {
+                    Ok(x) => x,
+                    Err(e) => {
+                        return Status::new(
+                            tonic::Code::Internal,
+                            format!("Failed to get container info: {:?}", e),
+                        )
+                    }
+                }
+                .into_inner()
+                .process;
+
+                for metric in metrics {
+                    let data = metric.data.unwrap();
+                    let resource = any_to_resource(&data);
+                    let status_code = containerd_client::types::v1::Status::from_i32(
+                        status.clone().unwrap().status,
+                    )
+                    .unwrap();
+                    let workload_status = WorkloadStatus {
+                        instance_id: workload.instance_id.clone(),
+                        status: Some(workload_status::Status {
+                            code: match status_code {
+                                containerd_client::types::v1::Status::Running => {
+                                    workload_status::status::StatusCode::Running as i32
+                                }
+                                containerd_client::types::v1::Status::Stopped => {
+                                    workload_status::status::StatusCode::Terminated as i32
+                                }
+                                _ => workload_status::status::StatusCode::Waiting as i32,
+                            },
+                            ..Default::default()
+                        }),
+                        resource_usage: resource.ok(),
+                    };
+                    match tx.send(Ok(workload_status)).await {
+                        Ok(x) => x,
+                        Err(e) => {
+                            return Status::new(
+                                tonic::Code::Internal,
+                                format!("Failed to send workload status: {:?}", e),
+                            )
+                        }
+                    };
+                }
+                tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+            }
+        });
+
+        Ok(Response::new(ReceiverStream::new(rx)))
+    }
+
+    async fn signal(&self, request: Request<WorkloadSignal>) -> Result<Response<Empty>, Status> {
+        let workload_signal = request.into_inner();
+
+        let mut client = ContainerClient::new("/var/run/containerd/containerd.sock")
+            .await
+            .map_err(|error| {
+                Status::new(
+                    tonic::Code::Internal,
+                    format!("Failed to create container client{:?}", error),
+                )
+            })?;
+
+        let signal = if workload_signal.signal == Signal::Stop as i32 {
+            sysinfo::Signal::Quit
+        } else {
+            sysinfo::Signal::Kill
+        };
+
+        // kill task
+        client
+            .kill(&workload_signal.instance_id, signal as u32)
+            .await
+            .map_err(|error| {
+                debug!("Failed to kill container {:?}", error);
+                into_tonic_status(error)
+            })?;
+
+        // wait for container to stop
+        client
+            .wait(&workload_signal.instance_id)
+            .await
+            .map_err(|error| {
+                debug!("Failed to wait for container {:?}", error);
+                into_tonic_status(error)
+            })?;
+
+        // cleanup container
+        client
+            .cleanup(&workload_signal.instance_id)
+            .await
+            .map_err(|error| {
+                debug!("Failed to cleanup container {:?}", error);
+                into_tonic_status(error)
+            })?;
+
+        Ok(Response::new(Empty {}))
+    }
+}