From e4727affe6053a3ae82833876ffb8f9c7b7c0aed Mon Sep 17 00:00:00 2001
From: trivernis <trivernis@protonmail.com>
Date: Sun, 10 Oct 2021 13:45:31 +0200
Subject: [PATCH] Add stop function to context

Signed-off-by: trivernis <trivernis@protonmail.com>
---
 Cargo.lock         |  2 +-
 Cargo.toml         |  2 +-
 src/error.rs       |  3 +++
 src/ipc/client.rs  |  9 ++++++++-
 src/ipc/context.rs | 23 +++++++++++++++++++++--
 src/ipc/server.rs  |  2 +-
 6 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 02abeb37..795eed0c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -148,7 +148,7 @@ dependencies = [
 
 [[package]]
 name = "rmp-ipc"
-version = "0.4.2"
+version = "0.4.3"
 dependencies = [
  "lazy_static",
  "log",
diff --git a/Cargo.toml b/Cargo.toml
index 06c40691..614bd7ad 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "rmp-ipc"
-version = "0.4.2"
+version = "0.4.3"
 authors = ["trivernis <trivernis@protonmail.com>"]
 edition = "2018"
 readme = "README.md"
diff --git a/src/error.rs b/src/error.rs
index ab3cb2ce..933b39af 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -22,6 +22,9 @@ pub enum Error {
 
     #[error("Channel Error: {0}")]
     ReceiveError(#[from] oneshot::error::RecvError),
+
+    #[error("Send Error")]
+    SendError,
 }
 
 impl From<String> for Error {
diff --git a/src/ipc/client.rs b/src/ipc/client.rs
index b1795984..adaae1f2 100644
--- a/src/ipc/client.rs
+++ b/src/ipc/client.rs
@@ -7,6 +7,7 @@ use crate::namespaces::namespace::Namespace;
 use std::collections::HashMap;
 use std::sync::Arc;
 use tokio::net::TcpStream;
+use tokio::sync::oneshot;
 use tokio::sync::RwLock;
 use typemap_rev::TypeMap;
 
@@ -26,20 +27,26 @@ impl IPCClient {
         let stream = TcpStream::connect(address).await?;
         let (read_half, write_half) = stream.into_split();
         let emitter = StreamEmitter::new(write_half);
+        let (tx, rx) = oneshot::channel();
         let ctx = Context::new(
             StreamEmitter::clone(&emitter),
             Arc::new(RwLock::new(self.data)),
+            Some(tx),
         );
         let handler = Arc::new(self.handler);
         let namespaces = Arc::new(self.namespaces);
         log::debug!("IPC client connected to {}", address);
 
-        tokio::spawn({
+        let handle = tokio::spawn({
             let ctx = Context::clone(&ctx);
             async move {
                 handle_connection(namespaces, handler, read_half, ctx).await;
             }
         });
+        tokio::spawn(async move {
+            let _ = rx.await;
+            handle.abort();
+        });
 
         Ok(ctx)
     }
diff --git a/src/ipc/context.rs b/src/ipc/context.rs
index f942dcf4..4cd10e7f 100644
--- a/src/ipc/context.rs
+++ b/src/ipc/context.rs
@@ -1,8 +1,10 @@
-use crate::error::Result;
+use crate::error::{Error, Result};
 use crate::ipc::stream_emitter::StreamEmitter;
 use crate::Event;
 use std::collections::HashMap;
+use std::mem;
 use std::sync::Arc;
+use tokio::sync::oneshot::Sender;
 use tokio::sync::{oneshot, Mutex, RwLock};
 use typemap_rev::TypeMap;
 
@@ -28,15 +30,22 @@ pub struct Context {
     /// Field to store additional context data
     pub data: Arc<RwLock<TypeMap>>,
 
+    stop_sender: Arc<Mutex<Option<Sender<()>>>>,
+
     reply_listeners: Arc<Mutex<HashMap<u64, oneshot::Sender<Event>>>>,
 }
 
 impl Context {
-    pub(crate) fn new(emitter: StreamEmitter, data: Arc<RwLock<TypeMap>>) -> Self {
+    pub(crate) fn new(
+        emitter: StreamEmitter,
+        data: Arc<RwLock<TypeMap>>,
+        stop_sender: Option<Sender<()>>,
+    ) -> Self {
         Self {
             emitter,
             reply_listeners: Arc::new(Mutex::new(HashMap::new())),
             data,
+            stop_sender: Arc::new(Mutex::new(stop_sender)),
         }
     }
 
@@ -52,6 +61,16 @@ impl Context {
         Ok(event)
     }
 
+    /// Stops the listener and closes the connection
+    pub async fn stop(self) -> Result<()> {
+        let mut sender = self.stop_sender.lock().await;
+        if let Some(sender) = mem::take(&mut *sender) {
+            sender.send(()).map_err(|_| Error::SendError)?;
+        }
+
+        Ok(())
+    }
+
     /// Returns the channel for a reply to the given message id
     pub(crate) async fn get_reply_sender(&self, ref_id: u64) -> Option<oneshot::Sender<Event>> {
         let mut listeners = self.reply_listeners.lock().await;
diff --git a/src/ipc/server.rs b/src/ipc/server.rs
index 1fa3e5ca..6926eb60 100644
--- a/src/ipc/server.rs
+++ b/src/ipc/server.rs
@@ -37,7 +37,7 @@ impl IPCServer {
             tokio::spawn(async {
                 let (read_half, write_half) = stream.into_split();
                 let emitter = StreamEmitter::new(write_half);
-                let ctx = Context::new(StreamEmitter::clone(&emitter), data);
+                let ctx = Context::new(StreamEmitter::clone(&emitter), data, None);
 
                 handle_connection(namespaces, handler, read_half, ctx).await;
             });