feat(scheduler): add NodeRegistered structure

Signed-off-by: iverly <[email protected]>

scheduler/Cargo.toml

@@ -17,3 +17,4 @@ serde_derive = "1.0.142"
 confy = "0.4.0"
 anyhow = "1.0.62"
 thiserror = "1.0.32"
+async-stream = "0.3.3"

scheduler/src/lib.rs

@@ -1,3 +1,5 @@
+use std::net::IpAddr;
+
 use proto::scheduler::{
     Instance, InstanceStatus, NodeRegisterRequest, NodeRegisterResponse, NodeStatus,
     NodeUnregisterRequest, NodeUnregisterResponse,
@@ -9,7 +11,8 @@ use tonic::Response;
 pub mod config;
 pub mod instance;
 pub mod manager;
-pub mod node_listener;
+pub mod node;
+pub mod parser;
 pub mod storage;
 
 #[derive(Error, Debug)]
@@ -28,6 +31,14 @@ pub enum SchedulerError {
 
 #[derive(Error, Debug)]
 pub enum ProxyError {
+    #[error("an transport error occurred from tonic: {0}")]
+    TonicTransportError(#[from] tonic::transport::Error),
+    #[error("an status error occurred from tonic: {0}")]
+    TonicStatusError(#[from] tonic::Status),
+    #[error("the gRPC client was not found")]
+    GrpcClientNotFound,
+    #[error("the gRPC stream was not found")]
+    GrpcStreamNotFound,
     #[error("an error occurred while sending a message to the channel")]
     ChannelSenderError,
 }
@@ -64,6 +75,7 @@ pub enum Event {
     // Node events
     NodeRegister(
         NodeRegisterRequest,
+        IpAddr,
         oneshot::Sender<Result<Response<NodeRegisterResponse>, tonic::Status>>,
     ),
     NodeUnregister(

scheduler/src/manager.rs

@@ -11,8 +11,9 @@ use tokio::{sync::oneshot, task::JoinHandle};
 use tonic::{transport::Server, Response};
 
 use crate::instance::listener::InstanceListener;
+use crate::node::listener::NodeListener;
 use crate::SchedulerError;
-use crate::{config::Config, node_listener::NodeListener, storage::Storage, Event, Node};
+use crate::{config::Config, storage::Storage, Event, Node};
 
 #[derive(Debug)]
 pub struct Manager {
@@ -139,7 +140,7 @@ impl Manager {
                         info!("received instance destroy event : {:?}", id);
                         tx.send(Ok(Response::new(()))).unwrap();
                     }
-                    Event::NodeRegister(request, tx) => {
+                    Event::NodeRegister(request, _, tx) => {
                         info!("received node register event : {:?}", request);
                         tx.send(Ok(Response::new(NodeRegisterResponse::default())))
                             .unwrap();

scheduler/src/node_listener.rs → scheduler/src/node/listener.rs

@@ -1,15 +1,13 @@
-use log::debug;
 use proto::scheduler::{
     node_service_server::NodeService, NodeRegisterRequest, NodeRegisterResponse, NodeStatus,
     NodeUnregisterRequest, NodeUnregisterResponse,
 };
 use tokio::sync::mpsc;
-use tonic::{Request, Response, Status, Streaming};
+use tonic::{Request, Response, Streaming};
 
 use crate::{manager::Manager, Event};
 
 #[derive(Debug)]
-#[allow(dead_code)]
 pub struct NodeListener {
     sender: mpsc::Sender<Event>,
 }
@@ -25,33 +23,34 @@ impl NodeService for NodeListener {
     async fn status(
         &self,
         request: Request<Streaming<NodeStatus>>,
-    ) -> Result<Response<()>, Status> {
+    ) -> Result<Response<()>, tonic::Status> {
+        log::debug!("received gRPC request: {:?}", request);
+
         let mut stream = request.into_inner();
-        let (tx, mut rx) = Manager::create_mpsc_channel();
 
+        // send each status to the manager
         loop {
+            let (tx, mut rx) = Manager::create_mpsc_channel();
             let message = stream.message().await?;
+
             match message {
                 Some(node_status) => {
-                    debug!("Node status: {:?}", node_status);
                     self.sender
                         .send(Event::NodeStatus(node_status, tx.clone()))
                         .await
                         .unwrap();
 
+                    // wait for the manager to respond
                     if let Some(res) = rx.recv().await {
                         match res {
-                            Ok(()) => {
-                                debug!("Node status updated successfully");
-                            }
-                            Err(err) => {
-                                debug!("Error updating node status: {:?}", err);
-                                return Err(err);
-                            }
+                            Ok(_) => {}
+                            Err(err) => return Err(err),
                         }
                     }
                 }
                 None => {
+                    log::error!("Node status stream closed");
+                    // todo: emit node crash event (get the node id from the first status)
                     return Ok(Response::new(()));
                 }
             }
@@ -61,29 +60,33 @@ impl NodeService for NodeListener {
     async fn register(
         &self,
         request: Request<NodeRegisterRequest>,
-    ) -> Result<Response<NodeRegisterResponse>, Status> {
-        debug!("{:?}", request);
+    ) -> Result<Response<NodeRegisterResponse>, tonic::Status> {
+        log::debug!("received gRPC request: {:?}", request);
+
         let (tx, rx) = Manager::create_oneshot_channel();
+        let remote_addr = request.remote_addr().unwrap().ip();
+        log::debug!("Registering a new node from: {:?}", remote_addr);
 
         match self
             .sender
-            .send(Event::NodeRegister(request.into_inner(), tx))
+            .send(Event::NodeRegister(request.into_inner(), remote_addr, tx))
             .await
         {
             Ok(_) => {
                 return rx.await.unwrap();
             }
             Err(_) => {
-                return Err(Status::internal("could not send event to manager"));
+                return Err(tonic::Status::internal("could not send event to manager"));
             }
         }
     }
 
     async fn unregister(
         &self,
         request: Request<NodeUnregisterRequest>,
-    ) -> Result<Response<NodeUnregisterResponse>, Status> {
-        debug!("{:?}", request);
+    ) -> Result<Response<NodeUnregisterResponse>, tonic::Status> {
+        log::debug!("received gRPC request: {:?}", request);
+
         let (tx, rx) = Manager::create_oneshot_channel();
 
         match self
@@ -95,7 +98,7 @@ impl NodeService for NodeListener {
                 return rx.await.unwrap();
             }
             Err(_) => {
-                return Err(Status::internal("could not send event to manager"));
+                return Err(tonic::Status::internal("could not send event to manager"));
             }
         }
     }

scheduler/src/node/mod.rs

@@ -0,0 +1,11 @@
+use proto::scheduler::{Resource, Status};
+
+pub mod listener;
+pub mod registered;
+
+#[derive(Debug, Clone)]
+pub struct Node {
+    pub id: String,
+    pub status: Status,
+    pub resource: Option<Resource>,
+}