Deregister node on exit and conflicting addr

Author: withtypes

crates/lunatic-distributed/src/control/client.rs

@@ -61,6 +61,7 @@ impl Client {
             node_id,
             signed_cert,
         } = client.send_registration(signing_request).await?;
+
         Ok((node_id, client, signed_cert))
     }
 
@@ -135,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())
     }

crates/lunatic-distributed/src/control/message.rs

@@ -5,6 +5,9 @@ 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),
@@ -14,6 +17,7 @@ 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",
@@ -28,6 +32,7 @@ pub enum Response {
     Module(Option<Vec<u8>>),
     ModuleId(u64),
     Error(String),
+    None,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]

crates/lunatic-distributed/src/control/server.rs

@@ -22,6 +22,7 @@ pub struct Server {
 struct InnerServer {
     next_node_id: AtomicU64,
     nodes: DashMap<u64, Registration>,
+    addr_to_node: DashMap<SocketAddr, u64>,
     next_module_id: AtomicU64,
     modules: DashMap<u64, Vec<u8>>,
     ca_cert: Certificate,
@@ -34,6 +35,7 @@ impl Server {
                 next_node_id: AtomicU64::new(1),
                 next_module_id: AtomicU64::new(1),
                 nodes: DashMap::new(),
+                addr_to_node: DashMap::new(),
                 modules: DashMap::new(),
                 ca_cert,
             }),
@@ -58,7 +60,15 @@ impl Server {
             .and_then(|sign_request| sign_request.serialize_pem_with_signer(&self.inner.ca_cert));
         match signed_cert {
             Ok(signed_cert) => {
+                // Remove another node using the same address. This is temporarily until we define
+                // details of connection status & reconnecting/registering.
+                if let Some(proc_id) = self.inner.addr_to_node.get(&reg.node_address) {
+                    self.inner.nodes.remove(&proc_id);
+                }
+
+                self.inner.addr_to_node.insert(reg.node_address, node_id);
                 self.inner.nodes.insert(node_id, reg);
+
                 Response::Register(Registered {
                     node_id,
                     signed_cert,
@@ -68,6 +78,11 @@ impl Server {
         }
     }
 
+    pub fn deregister(&self, node_id: u64) -> Response {
+        self.inner.nodes.remove(&node_id);
+        Response::None
+    }
+
     pub fn list_nodes(&self) -> Response {
         Response::Nodes(
             self.inner
@@ -170,6 +185,7 @@ pub async fn handle_request(
     use crate::control::message::Request::*;
     let response = match request {
         Register(reg) => server.register(reg),
+        Deregister(node_id) => server.deregister(node_id),
         ListNodes => server.list_nodes(),
         AddModule(bytes) => server.add_module(bytes),
         GetModule(id) => server.get_module(id),

src/mode/execution.rs

@@ -128,60 +128,62 @@ pub(crate) async fn execute() -> Result<()> {
 
     let env = envs.get_or_create(1);
 
-    let distributed_state = if let (Some(node_address), Some(control_address)) =
-        (args.value_of("node"), args.value_of("control"))
-    {
-        // TODO unwrap, better message
-        let node_address = node_address.parse().unwrap();
-        let node_name = Uuid::new_v4().to_string();
-        let control_address = control_address.parse().unwrap();
-        let ca_cert = lunatic_distributed::distributed::server::root_cert(
-            args.is_present("test_ca"),
-            args.value_of("ca_cert"),
-        )
-        .unwrap();
-        let node_cert =
-            lunatic_distributed::distributed::server::gen_node_cert(&node_name).unwrap();
+    let (distributed_state, control_client, node_id) =
+        if let (Some(node_address), Some(control_address)) =
+            (args.value_of("node"), args.value_of("control"))
+        {
+            // TODO unwrap, better message
+            let node_address = node_address.parse().unwrap();
+            let node_name = Uuid::new_v4().to_string();
+            let control_address = control_address.parse().unwrap();
+            let ca_cert = lunatic_distributed::distributed::server::root_cert(
+                args.is_present("test_ca"),
+                args.value_of("ca_cert"),
+            )
+            .unwrap();
+            let node_cert =
+                lunatic_distributed::distributed::server::gen_node_cert(&node_name).unwrap();
 
-        let quic_client = quic::new_quic_client(&ca_cert).unwrap();
+            let quic_client = quic::new_quic_client(&ca_cert).unwrap();
 
-        let (node_id, control_client, signed_cert_pem) = control::Client::register(
-            node_address,
-            node_name.to_string(),
-            control_address,
-            quic_client.clone(),
-            node_cert.serialize_request_pem().unwrap(),
-        )
-        .await?;
+            let (node_id, control_client, signed_cert_pem) = control::Client::register(
+                node_address,
+                node_name.to_string(),
+                control_address,
+                quic_client.clone(),
+                node_cert.serialize_request_pem().unwrap(),
+            )
+            .await?;
 
-        let distributed_client =
-            distributed::Client::new(node_id, control_client.clone(), quic_client.clone()).await?;
+            let distributed_client =
+                distributed::Client::new(node_id, control_client.clone(), quic_client.clone())
+                    .await?;
 
-        let dist = lunatic_distributed::DistributedProcessState::new(
-            node_id,
-            control_client.clone(),
-            distributed_client,
-        )
-        .await?;
-
-        tokio::task::spawn(lunatic_distributed::distributed::server::node_server(
-            ServerCtx {
-                envs,
-                modules: Modules::<DefaultProcessState>::default(),
-                distributed: dist.clone(),
-                runtime: runtime.clone(),
-            },
-            node_address,
-            signed_cert_pem,
-            node_cert.serialize_private_key_pem(),
-        ));
-
-        log::info!("Registration successful, node id {}", node_id);
-
-        Some(dist)
-    } else {
-        None
-    };
+            let dist = lunatic_distributed::DistributedProcessState::new(
+                node_id,
+                control_client.clone(),
+                distributed_client,
+            )
+            .await?;
+
+            tokio::task::spawn(lunatic_distributed::distributed::server::node_server(
+                ServerCtx {
+                    envs,
+                    modules: Modules::<DefaultProcessState>::default(),
+                    distributed: dist.clone(),
+                    runtime: runtime.clone(),
+                },
+                node_address,
+                signed_cert_pem,
+                node_cert.serialize_private_key_pem(),
+            ));
+
+            log::info!("Registration successful, node id {}", node_id);
+
+            (Some(dist), Some(control_client), Some(node_id))
+        } else {
+            (None, None, None)
+        };
 
     let mut config = DefaultProcessConfig::default();
     // Allow initial process to compile modules, create configurations and spawn sub-processes
@@ -244,7 +246,14 @@ pub(crate) async fn execute() -> Result<()> {
                 path.to_string_lossy()
             ))?;
         // Wait on the main process to finish
-        task.await.map(|_| ()).map_err(|e| anyhow!(e.to_string()))
+        let result = task.await.map(|_| ()).map_err(|e| anyhow!(e.to_string()));
+
+        // Until we refactor registration and reconnect authentication, send node id explicitly
+        if let (Some(ctrl), Some(node_id)) = (control_client, node_id) {
+            ctrl.deregister(node_id).await;
+        }
+
+        result
     } else {
         // Block forever
         let (_sender, mut receiver) = channel::<()>(1);