1 files changed, 34 insertions, 32 deletions

diff --git a/echo/src/main.rs b/echo/src/main.rs
index 5ae72f3..fed0059 100644
--- a/echo/src/main.rs
+++ b/echo/src/main.rs
@@ -1,57 +1,59 @@
-use std::io::{self, Write};
-
-use common::{msg::*, run_with, Handler, MsgWriter};
+#![feature(return_position_impl_trait_in_trait)]
+use std::future::Future;
+
+use common::{
+    msg::*,
+    msg_id::{gen_msg_id, MessageID},
+    run_server, Handler,
+};
 use serde::{Deserialize, Serialize};
 
 fn main() {
-    run_with::<EchoHandler>(io::stdin(), io::stdout());
+    run_server::<EchoHandler>();
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "type")]
 pub enum EchoBody {
     #[serde(rename = "echo")]
-    Echo { msg_id: usize, echo: String },
+    Echo { msg_id: MessageID, echo: String },
 
     #[serde(rename = "echo_ok")]
     EchoOk {
-        msg_id: usize,
-        in_reply_to: usize,
+        msg_id: MessageID,
+        in_reply_to: MessageID,
         echo: String,
     },
 }
 
 pub struct EchoHandler {
-    next_msg_id: usize,
+    output: Output<EchoBody>,
 }
 
 impl Handler for EchoHandler {
     type Body = EchoBody;
 
-    fn init(_node_id: String, _node_ids: Vec<String>, _msg_id: usize) -> Self {
-        EchoHandler { next_msg_id: 1 }
+    fn init(_node_id: String, _node_ids: Vec<String>, output: Output<Self::Body>) -> Self {
+        EchoHandler { output }
     }
 
-    fn handle(
-        &mut self,
-        header: MessageHeader,
-        body: Self::Body,
-        writer: &mut MsgWriter<impl Write>,
-    ) {
-        match body {
-            EchoBody::Echo { msg_id, echo } => {
-                writer.write(
-                    header.src,
-                    &EchoBody::EchoOk {
-                        msg_id: self.next_msg_id,
-                        in_reply_to: msg_id,
-                        echo,
-                    },
-                );
-
-                self.next_msg_id += 1;
-            }
-            EchoBody::EchoOk { .. } => (),
-        };
+    fn handle(&self, header: MessageHeader, body: Self::Body) -> impl Future<Output = ()> + Send {
+        async move {
+            match body {
+                EchoBody::Echo { msg_id, echo } => {
+                    self.output
+                        .send(
+                            &header.src,
+                            &EchoBody::EchoOk {
+                                msg_id: gen_msg_id(),
+                                in_reply_to: msg_id,
+                                echo,
+                            },
+                        )
+                        .await;
+                }
+                EchoBody::EchoOk { .. } => (),
+            };
+        }
     }
 }