Working on client / relayer integration

crates/relayer/src/connection.rs

@@ -12,6 +12,7 @@ use tokio::{
         mpsc::{channel, Receiver, Sender},
         RwLock,
     },
+    task::JoinHandle,
 };
 #[allow(unused_imports)]
 use tokio_tungstenite::{accept_async, tungstenite::Message, WebSocketStream};
@@ -22,10 +23,19 @@ pub struct Connection {
     pub(crate) conn_id: u128,
     sender: Sender<Response>,
     subscriptions: RwLock<HashMap<String, u128>>,
+    handler: Option<JoinHandle<()>>,
 }
 
 const MAX_SUBSCRIPTIONS_BUFFER: usize = 100;
 
+impl Drop for Connection {
+    fn drop(&mut self) {
+        if let Some(handler) = self.handler.take() {
+            let _ = handler.abort();
+        }
+    }
+}
+
 impl Connection {
     #[cfg(test)]
     pub fn new_for_test() -> (Self, Receiver<Response>) {
@@ -35,12 +45,13 @@ impl Connection {
                 conn_id: 0,
                 sender,
                 subscriptions: RwLock::new(HashMap::new()),
+                handler: None,
             },
             receiver,
         )
     }
 
-    pub async fn new(
+    pub async fn new_connection(
         broadcast_request: Sender<(u128, Request)>,
         disconnection_notify: Option<Sender<u128>>,
         stream: TcpStream,
@@ -48,29 +59,28 @@ impl Connection {
         let websocket = accept_async(stream).await?;
         let conn_id = get_id();
         let (sender, receiver) = channel(MAX_SUBSCRIPTIONS_BUFFER);
-        Self::spawn(
-            broadcast_request,
-            websocket,
-            receiver,
-            disconnection_notify,
-            conn_id,
-        );
         let _ = sender.send(Auth::default().into()).await;
         Ok(Self {
             conn_id,
             sender,
             subscriptions: RwLock::new(HashMap::new()),
+            handler: Some(Self::spawn(
+                broadcast_request,
+                websocket,
+                receiver,
+                disconnection_notify,
+                conn_id,
+            )),
         })
     }
 
-    #[allow(unused)]
     fn spawn(
-        broadcast_request: Sender<(u128, Request)>,
+        publish_message_from_client: Sender<(u128, Request)>,
         websocket: WebSocketStream<TcpStream>,
         mut receiver: Receiver<Response>,
         disconnection_notify: Option<Sender<u128>>,
         conn_id: u128,
-    ) {
+    ) -> JoinHandle<()> {
         tokio::spawn(async move {
             let mut _subscriptions: HashMap<String, (u128, Receiver<Response>)> = HashMap::new();
             let (mut writer, mut reader) = websocket.split();
@@ -92,7 +102,7 @@ impl Connection {
                             let msg: Result<Request, _> = serde_json::from_str(&msg);
                             match msg {
                                 Ok(msg) => {
-                                    let _ = broadcast_request.send((conn_id, msg)).await;
+                                    let _ = publish_message_from_client.send((conn_id, msg)).await;
                                 },
                                 Err(err) => {
                                     log::error!("Error parsing message from client: {}", err);
@@ -123,7 +133,7 @@ impl Connection {
             if let Some(disconnection_notify) = disconnection_notify {
                 let _ = disconnection_notify.try_send(conn_id);
             }
-        });
+        })
     }
 
     #[inline]

crates/relayer/src/relayer.rs

@@ -7,7 +7,7 @@ use nostr_rs_types::{
     types::{Event, SubscriptionId},
     Request, Response,
 };
-use std::{collections::HashMap, ops::Deref, sync::Arc};
+use std::{collections::HashMap, ops::Deref};
 use tokio::{
     net::TcpStream,
     sync::{
@@ -60,10 +60,10 @@ impl<T: Storage> Relayer<T> {
     pub fn new(
         storage: Option<T>,
         client_pool: Option<Pool>,
-    ) -> Result<(Arc<Self>, Receiver<(u128, Request)>), Error> {
+    ) -> Result<(Self, Receiver<(u128, Request)>), Error> {
         let (sender, receiver) = channel(100_000);
         Ok((
-            Arc::new(Self {
+            Self {
                 storage,
                 sender: sender.clone(),
                 subscriptions: Default::default(),
@@ -74,11 +74,16 @@ impl<T: Storage> Relayer<T> {
                 } else {
                     None
                 },
-            }),
+            },
             receiver,
         ))
     }
 
+    pub fn main(self) -> JoinHandle<()> {
+        let relayer = self;
+        tokio::spawn(async move { todo!() })
+    }
+
     fn handle_client_pool(
         client_pool: Pool,
         sender: Sender<(u128, Request)>,
@@ -111,10 +116,10 @@ impl<T: Storage> Relayer<T> {
         disconnection_notify: Option<mpsc::Sender<u128>>,
         stream: TcpStream,
     ) -> Result<u128, Error> {
-        let client = Connection::new(self.sender.clone(), disconnection_notify, stream).await?;
+        let client =
+            Connection::new_connection(self.sender.clone(), disconnection_notify, stream).await?;
         let id = client.conn_id;
-        let mut clients = self.clients.write().await;
-        clients.insert(id, client);
+        self.clients.write().await.insert(id, client);
 
         Ok(id)
     }