Merge pull request #114 from lunatic-solutions/distributed_quic

Replace tcp with quic for distributed nodes

Author: bkolobara

Cargo.lock

@@ -18,6 +18,12 @@ path = "src/lib.rs"
 name = "lunatic"
 path = "src/main.rs"
 
+[features]
+default = ["distributed-quinn"]
+distributed-tcp = ["lunatic-distributed/tcp"]
+distributed-s2n = ["lunatic-distributed/quic-s2n"]
+distributed-quinn = ["lunatic-distributed/quic-quinn"]
+
 [dependencies]
 anyhow = "^1.0"
 clap = { version = "^3.0", features = ["cargo"] }
@@ -33,6 +39,7 @@ serde = "^1.0"
 bincode = "^1.3"
 dashmap = "^5.3.4"
 regex = "^1.5"
+uuid = { version = "^1.1", features = ["v4"] }
 hash-map-id = { version = "^0.10", path = "crates/hash-map-id" }
 lunatic-stdout-capture = { version = "^0.10", path = "crates/lunatic-stdout-capture" }
 lunatic-process = { version = "^0.10", path = "crates/lunatic-process" }

Cargo.toml

@@ -122,7 +122,6 @@ where
             .data(&caller)
             .get(params_ptr as usize..(params_ptr + params_len) as usize)
             .or_trap("lunatic::distributed::spawn::params")?;
-
         let params = params
             .chunks_exact(17)
             .map(|chunk| {

crates/lunatic-distributed-api/src/lib.rs

@@ -6,13 +6,24 @@ description = "Node to node communication"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[features]
+tcp = []
+quic-s2n = ["dep:s2n-quic"]
+quic-quinn = ["dep:quinn", "dep:rustls", "dep:rustls-pemfile", "dep:futures-util"]
+
 [dependencies]
 anyhow = "^1.0"
 tokio = { version = "^1.20", features = ["macros", "rt-multi-thread", "net", "time", "io-util"] }
 serde = { version = "^1.0", features = ["derive"] }
+s2n-quic = { version = "1",  default-features = false, features = ["provider-address-token-default", "provider-tls-rustls"], optional = true}
 wasmtime = "^0.39"
 dashmap = "5.3.4"
 bincode = "^1.3"
 log = "^0.4"
 async_cell = "0.2.1"
-lunatic-process = { version = "^0.10", path = "../lunatic-process" }
+rcgen = { version="^0.9", features=["pem", "x509-parser"] }
+lunatic-process = { version = "^0.10", path = "../lunatic-process" }
+quinn = { version = "^0.8", optional = true}
+rustls = { version = "^0.20", optional = true }
+rustls-pemfile = { version = "^1.0", optional = true }
+futures-util = { version = "^0.3", optional = true}

crates/lunatic-distributed/Cargo.toml

@@ -7,14 +7,15 @@ use std::{
     sync::{atomic, atomic::AtomicU64, Arc, RwLock},
     time::Duration,
 };
-use tokio::net::TcpStream;
 
 use crate::{
-    connection::Connection,
-    control::message::{Registration, Request, Response},
+    control::message::{Registered, Registration, Request, Response},
+    quic::{self, Connection},
     NodeInfo,
 };
 
+use super::server::CTRL_SERVER_NAME;
+
 #[derive(Clone)]
 pub struct Client {
     inner: Arc<InnerClient>,
@@ -23,6 +24,7 @@ pub struct Client {
 pub struct InnerClient {
     next_message_id: AtomicU64,
     node_addr: SocketAddr,
+    node_name: String,
     control_addr: SocketAddr,
     connection: Connection,
     pending_requests: DashMap<u64, Arc<AsyncCell<Response>>>,
@@ -31,13 +33,22 @@ pub struct InnerClient {
 }
 
 impl Client {
-    pub async fn register(node_addr: SocketAddr, control_addr: SocketAddr) -> Result<(u64, Self)> {
+    pub async fn register(
+        node_addr: SocketAddr,
+        node_name: String,
+        control_addr: SocketAddr,
+        quic_client: quic::Client,
+        signing_request: String,
+    ) -> Result<(u64, Self, String)> {
         let client = Client {
             inner: Arc::new(InnerClient {
                 next_message_id: AtomicU64::new(1),
                 control_addr,
                 node_addr,
-                connection: connect(control_addr, 5).await?,
+                node_name,
+                connection: quic_client
+                    .connect(control_addr, CTRL_SERVER_NAME, 5)
+                    .await?,
                 pending_requests: DashMap::new(),
                 nodes: Default::default(),
                 node_ids: Default::default(),
@@ -46,8 +57,12 @@ impl Client {
         // Spawn reader task before register
         tokio::task::spawn(reader_task(client.clone()));
         tokio::task::spawn(refresh_nodes_task(client.clone()));
-        let node_id: u64 = client.send_registration().await?;
-        Ok((node_id, client))
+        let Registered {
+            node_id,
+            signed_cert,
+        } = client.send_registration(signing_request).await?;
+
+        Ok((node_id, client, signed_cert))
     }
 
     pub fn next_message_id(&self) -> u64 {
@@ -78,15 +93,18 @@ impl Client {
         self.inner.connection.receive().await
     }
 
-    async fn send_registration(&self) -> Result<u64> {
+    async fn send_registration(&self, signing_request: String) -> Result<Registered> {
         let reg = Registration {
             node_address: self.inner.node_addr,
+            node_name: self.inner.node_name.clone(),
+            signing_request,
         };
         let resp = self.send(Request::Register(reg)).await?;
-        if let Response::Register(node_id) = resp {
-            return Ok(node_id);
+        match resp {
+            Response::Register(data) => Ok(data),
+            Response::Error(e) => Err(anyhow!("Registration failed. {e}")),
+            _ => Err(anyhow!("Registration failed.")),
         }
-        Err(anyhow!("Registration failed."))
     }
 
     fn process_response(&self, id: u64, resp: Response) {
@@ -106,6 +124,7 @@ impl Client {
                         NodeInfo {
                             id,
                             address: reg.node_address,
+                            name: reg.node_name,
                         },
                     );
                 }
@@ -117,6 +136,10 @@ impl Client {
         Ok(())
     }
 
+    pub async fn deregister(&self, node_id: u64) {
+        self.send(Request::Deregister(node_id)).await.ok();
+    }
+
     pub fn node_info(&self, node_id: u64) -> Option<NodeInfo> {
         self.inner.nodes.get(&node_id).map(|e| e.clone())
     }
@@ -146,17 +169,6 @@ impl Client {
     }
 }
 
-async fn connect(addr: SocketAddr, retry: u32) -> Result<Connection> {
-    for _ in 0..retry {
-        log::info!("Connecting to control {addr}");
-        if let Ok(stream) = TcpStream::connect(addr).await {
-            return Ok(Connection::new(stream));
-        }
-        tokio::time::sleep(Duration::from_secs(2)).await;
-    }
-    Err(anyhow!("Failed to connect to {addr}"))
-}
-
 async fn reader_task(client: Client) -> Result<()> {
     loop {
         if let Ok((id, resp)) = client.recv().await {

crates/lunatic-distributed/src/connection.rs

@@ -5,20 +5,45 @@ use serde::{Deserialize, Serialize};
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub enum Request {
     Register(Registration),
+    // Currently a node will send it's own id. We need to refactor this part: the control server
+    // should always handle registration first and later know which node is sending requests.
+    Deregister(u64),
     ListNodes,
     AddModule(Vec<u8>),
     GetModule(u64),
 }
 
+impl Request {
+    pub fn kind(&self) -> &'static str {
+        match self {
+            Request::Register(_) => "Register",
+            Request::Deregister(_) => "Deregister",
+            Request::ListNodes => "ListNodes",
+            Request::AddModule(_) => "AddModule",
+            Request::GetModule(_) => "GetModule",
+        }
+    }
+}
+
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub enum Response {
-    Register(u64),
+    Register(Registered),
     Nodes(Vec<(u64, Registration)>),
     Module(Option<Vec<u8>>),
     ModuleId(u64),
+    Error(String),
+    None,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct Registration {
     pub node_address: SocketAddr,
+    pub node_name: String,
+    pub signing_request: String,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct Registered {
+    pub node_id: u64,
+    pub signed_cert: String,
 }