1 files changed, 24 insertions, 20 deletions

diff --git a/p2p/examples/chat.rs b/p2p/examples/chat.rs
index 4358362..907ba06 100644
--- a/p2p/examples/chat.rs
+++ b/p2p/examples/chat.rs
@@ -1,9 +1,11 @@
+mod shared;
+
 use std::sync::Arc;
 
 use async_std::io;
 use async_trait::async_trait;
 use clap::Parser;
-use smol::{channel, future, Executor};
+use smol::{channel, Executor};
 
 use karyons_net::{Endpoint, Port};
 
@@ -12,6 +14,8 @@ use karyons_p2p::{
     ArcPeer, Backend, Config, P2pError, PeerID, Version,
 };
 
+use shared::run_executor;
+
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
 struct Cli {
@@ -109,33 +113,33 @@ fn main() {
         ..Default::default()
     };
 
+    // Create a new Executor
+    let ex = Arc::new(Executor::new());
+
     // Create a new Backend
-    let backend = Backend::new(peer_id, config);
+    let backend = Backend::new(peer_id, config, ex.clone());
 
     let (ctrlc_s, ctrlc_r) = channel::unbounded();
     let handle = move || ctrlc_s.try_send(()).unwrap();
     ctrlc::set_handler(handle).unwrap();
 
-    // Create a new Executor
-    let ex = Arc::new(Executor::new());
-
-    let ex_cloned = ex.clone();
-    let task = ex.spawn(async {
-        let username = cli.username;
-
-        // Attach the ChatProtocol
-        let c = move |peer| ChatProtocol::new(&username, peer);
-        backend.attach_protocol::<ChatProtocol>(c).await.unwrap();
+    run_executor(
+        async {
+            let username = cli.username;
 
-        // Run the backend
-        backend.run(ex_cloned).await.unwrap();
+            // Attach the ChatProtocol
+            let c = move |peer| ChatProtocol::new(&username, peer);
+            backend.attach_protocol::<ChatProtocol>(c).await.unwrap();
 
-        // Wait for ctrlc signal
-        ctrlc_r.recv().await.unwrap();
+            // Run the backend
+            backend.run().await.unwrap();
 
-        // Shutdown the backend
-        backend.shutdown().await;
-    });
+            // Wait for ctrlc signal
+            ctrlc_r.recv().await.unwrap();
 
-    future::block_on(ex.run(task));
+            // Shutdown the backend
+            backend.shutdown().await;
+        },
+        ex,
+    );
 }