1 files changed, 44 insertions, 50 deletions
diff --git a/unique_ids/src/main.rs b/unique_ids/src/main.rs
index 3787033..f899e2b 100644
--- a/unique_ids/src/main.rs
+++ b/unique_ids/src/main.rs
@@ -1,83 +1,77 @@
+#![feature(return_position_impl_trait_in_trait)]
 use std::{
-    io,
+    future::Future,
     time::{SystemTime, UNIX_EPOCH},
 };
 
-use common::{run_with, Handler};
-use rand::{rngs::StdRng, Rng, SeedableRng};
+use common::{
+    msg::{MessageHeader, Output},
+    msg_id::{gen_msg_id, MessageID},
+    run_server, Handler,
+};
+use rand::{thread_rng, Rng};
 use serde::{Deserialize, Serialize};
 
 fn main() {
-    run_with::<UniqueIdsHandler>(io::stdin(), io::stdout())
+    run_server::<UniqueIdsHandler>()
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "type")]
 enum UniqueIdsBody {
     #[serde(rename = "generate")]
-    Generate { msg_id: usize },
+    Generate { msg_id: MessageID },
 
     #[serde(rename = "generate_ok")]
     GenerateOk {
-        msg_id: usize,
-        in_reply_to: usize,
+        msg_id: MessageID,
+        in_reply_to: MessageID,
         id: u128,
     },
 }
 
-type RngMethod = StdRng;
-
 struct UniqueIdsHandler {
-    rng: RngMethod,
-    next_msg_id: usize,
+    output: Output<UniqueIdsBody>,
 }
 
 impl Handler for UniqueIdsHandler {
     type Body = UniqueIdsBody;
 
-    fn init(_node_id: String, _node_ids: Vec<String>, _msg_id: usize) -> Self {
-        Self {
-            rng: RngMethod::from_entropy(), // TODO: This could be seeded from the node ID
-            next_msg_id: 1,
-        }
+    fn init(_node_id: String, _node_ids: Vec<String>, output: Output<Self::Body>) -> Self {
+        Self { output }
     }
 
-    fn handle(
-        &mut self,
-        header: common::msg::MessageHeader,
-        body: Self::Body,
-        writer: &mut common::MsgWriter<impl std::io::Write>,
-    ) -> () {
-        match body {
-            UniqueIdsBody::Generate { msg_id } => {
-                let id = self.gen_id();
-                writer.write(
-                    header.src,
-                    &UniqueIdsBody::GenerateOk {
-                        msg_id: self.next_msg_id,
-                        in_reply_to: msg_id,
-                        id,
-                    },
-                );
-
-                self.next_msg_id += 1;
-            }
-            UniqueIdsBody::GenerateOk { .. } => (),
-        };
+    fn handle(&self, header: MessageHeader, body: Self::Body) -> impl Future<Output = ()> + Send {
+        async move {
+            match body {
+                UniqueIdsBody::Generate { msg_id } => {
+                    let id = gen_id();
+                    self.output
+                        .send(
+                            &header.src,
+                            &UniqueIdsBody::GenerateOk {
+                                msg_id: gen_msg_id(),
+                                in_reply_to: msg_id,
+                                id,
+                            },
+                        )
+                        .await;
+                }
+                UniqueIdsBody::GenerateOk { .. } => (),
+            };
+        }
     }
 }
 
-impl UniqueIdsHandler {
-    fn gen_id(&mut self) -> u128 {
-        // Time since UNIX epoch in milliseconds
-        let now = SystemTime::now();
-        let time_millis: u128 = now.duration_since(UNIX_EPOCH).unwrap().as_millis();
+fn gen_id() -> u128 {
+    // Time since UNIX epoch in milliseconds
+    let now = SystemTime::now();
+    let time_millis: u128 = now.duration_since(UNIX_EPOCH).unwrap().as_millis();
 
-        // 80 bits of randomness
-        let rand1: u16 = self.rng.gen();
-        let rand2: u64 = self.rng.gen();
-        let rand: u128 = rand1 as u128 | ((rand2 as u128) << 64);
+    // 80 bits of randomness
+    let rand1: u16 = thread_rng().gen();
+    let rand2: u64 = thread_rng().gen();
+    let rand: u128 = rand1 as u128 | ((rand2 as u128) << 64);
 
-        (time_millis << 80) | rand
-    }
+    (time_millis << 80) | rand
 }