1 files changed, 8 insertions, 6 deletions
diff --git a/p2p/src/protocol.rs b/p2p/src/protocol.rs
index 021844f..249692b 100644
--- a/p2p/src/protocol.rs
+++ b/p2p/src/protocol.rs
@@ -56,11 +56,8 @@ impl EventValue for ProtocolEvent {
 /// #[async_trait]
 /// impl Protocol for NewProtocol {
 ///     async fn start(self: Arc<Self>) -> Result<(), Error> {
-///         let listener = self.peer.register_listener::<Self>().await;
 ///         loop {
-///             let event = listener.recv().await.unwrap();
-///
-///             match event {
+///             match self.peer.recv::<Self>().await.expect("Receive msg") {
 ///                 ProtocolEvent::Message(msg) => {
 ///                     println!("{:?}", msg);
 ///                 }
@@ -69,8 +66,6 @@ impl EventValue for ProtocolEvent {
 ///                 }
 ///             }
 ///         }
-///
-///         listener.cancel().await;
 ///         Ok(())
 ///     }
 ///
@@ -114,3 +109,10 @@ pub trait Protocol: Send + Sync {
     where
         Self: Sized;
 }
+
+#[async_trait]
+pub(crate) trait InitProtocol: Send + Sync {
+    type T;
+    /// Initialize the protocol
+    async fn init(self: Arc<Self>) -> Self::T;
+}