From d2633d5868f2d63aa58d9ff7322a3716883cefbd Mon Sep 17 00:00:00 2001
From: Bartosz Nowak <bartekn337@gmail.com>
Date: Sun, 14 Apr 2024 14:26:32 +0200
Subject: [PATCH] gracefull drop

---
 Cargo.toml               |  1 +
 crates/peer/Cargo.toml   |  1 +
 crates/peer/src/swarm.rs | 14 ++++++++++++--
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 083d97e..23d6cb4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -36,6 +36,7 @@ starknet = "0.9.0"
 tempfile = "3.10.1"
 thiserror = "1.0.58"
 tokio = { version = "1.36", features = ["full"] }
+tokio-util = "0.7.10"
 tracing = "0.1.37"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 
diff --git a/crates/peer/Cargo.toml b/crates/peer/Cargo.toml
index d0f47bf..88dbf99 100644
--- a/crates/peer/Cargo.toml
+++ b/crates/peer/Cargo.toml
@@ -15,6 +15,7 @@ futures.workspace = true
 libp2p.workspace = true
 sharp-p2p-common.workspace = true
 starknet.workspace = true
+tokio-util.workspace = true
 tokio.workspace = true
 tracing-subscriber.workspace = true
 tracing.workspace = true
\ No newline at end of file
diff --git a/crates/peer/src/swarm.rs b/crates/peer/src/swarm.rs
index 78d9c3e..aea6100 100644
--- a/crates/peer/src/swarm.rs
+++ b/crates/peer/src/swarm.rs
@@ -8,6 +8,7 @@ use libp2p::{mdns, noise, tcp, yamux, Swarm, SwarmBuilder};
 use std::error::Error;
 use std::pin::Pin;
 use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
 use tracing::{debug, error};
 
 #[derive(NetworkBehaviour)]
@@ -18,6 +19,7 @@ struct PeerBehaviour {
 
 pub struct SwarmRunner {
     swarm: Swarm<PeerBehaviour>,
+    cancellation_token: CancellationToken,
 }
 
 impl SwarmRunner {
@@ -43,7 +45,7 @@ impl SwarmRunner {
         swarm.listen_on("/ip4/0.0.0.0/udp/0/quic-v1".parse()?)?;
         swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
 
-        Ok(SwarmRunner { swarm })
+        Ok(SwarmRunner { swarm, cancellation_token: CancellationToken::new() })
     }
 
     fn init_gossip(
@@ -106,10 +108,18 @@ impl SwarmRunner {
                         }
                         _ => {}
                     },
-                    else => break
+                    _ = self.cancellation_token.cancelled() => {
+                        break
+                    }
                 }
             }
         };
         Box::pin(stream)
     }
 }
+
+impl Drop for SwarmRunner {
+    fn drop(&mut self) {
+        self.cancellation_token.cancel();
+    }
+}