refactor!: Migrate to tokio's AbortOnDropHandle (#2701)

## Description

This migrates us to tokio's AbortOnDropHandle from the various
versions we had for this ourselves.  Only for iroh::node::Node we keep
a local version around since it is both Clone *and* awaits the tasks.
Generally it's probably better to not make this kind of design though,
if possible.

## Breaking Changes

Removes the `iroh_net::util` module.  This only contained utilities
that should have been internal in the first place.  If you were using
one of the tools to abort tasks on drop, we recommend migrating to
`tokio_util::task::AbortOnDropHandle`.

## Notes & open questions

<!-- Any notes, remarks or open questions you have to make about the PR.
-->

## Change checklist

- [x] Self-review.
- [x] Documentation updates following the [style
guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text),
if relevant.
- ~~[ ] Tests if relevant.~~
- [x] All breaking changes documented.

Author: flub

Cargo.lock

@@ -61,6 +61,7 @@ tempfile = "3.10.1"
 thiserror = "1.0.58"
 time = { version = "0.3", features = ["formatting"] }
 tokio = { version = "1.36.0", features = ["full"] }
+tokio-util = { version = "0.7.12", features = ["rt"] }
 toml = { version = "0.8.12", features = ["preserve_order"] }
 tracing = "0.1.40"
 tracing-appender = "0.2.3"

iroh-cli/Cargo.toml

@@ -35,7 +35,6 @@ use iroh::{
         netcheck, portmapper,
         relay::{RelayMap, RelayMode, RelayUrl},
         ticket::NodeTicket,
-        util::CancelOnDrop,
         Endpoint, NodeAddr, NodeId,
     },
     util::{path::IrohPaths, progress::ProgressWriter},
@@ -44,6 +43,7 @@ use portable_atomic::AtomicU64;
 use postcard::experimental::max_size::MaxSize;
 use serde::{Deserialize, Serialize};
 use tokio::{io::AsyncWriteExt, sync};
+use tokio_util::task::AbortOnDropHandle;
 
 use iroh::net::metrics::MagicsockMetrics;
 use iroh_metrics::core::Core;
@@ -372,7 +372,7 @@ struct Gui {
     recv_pb: ProgressBar,
     echo_pb: ProgressBar,
     #[allow(dead_code)]
-    counter_task: Option<CancelOnDrop>,
+    counter_task: Option<AbortOnDropHandle<()>>,
 }
 
 impl Gui {
@@ -412,10 +412,7 @@ impl Gui {
             send_pb,
             recv_pb,
             echo_pb,
-            counter_task: Some(CancelOnDrop::new(
-                "counter_task",
-                counter_task.abort_handle(),
-            )),
+            counter_task: Some(AbortOnDropHandle::new(counter_task)),
         }
     }
 

iroh-cli/src/commands/doctor.rs

@@ -44,7 +44,7 @@ tempfile = { version = "3.4" }
 thiserror = "1"
 tokio = { version = "1", features = ["sync", "rt", "time", "macros"] }
 tokio-stream = { version = "0.1", optional = true, features = ["sync"]}
-tokio-util = { version = "0.7", optional = true, features = ["codec", "io-util", "io"] }
+tokio-util = { version = "0.7.12", optional = true, features = ["codec", "io-util", "io", "rt"] }
 tracing = "0.1"
 
 [dev-dependencies]

iroh-docs/Cargo.toml

@@ -14,10 +14,10 @@ use futures_lite::{Stream, StreamExt};
 use iroh_blobs::downloader::Downloader;
 use iroh_blobs::{store::EntryStatus, Hash};
 use iroh_gossip::net::Gossip;
-use iroh_net::util::SharedAbortingJoinHandle;
 use iroh_net::{key::PublicKey, Endpoint, NodeAddr};
 use serde::{Deserialize, Serialize};
 use tokio::sync::{mpsc, oneshot};
+use tokio_util::task::AbortOnDropHandle;
 use tracing::{error, error_span, Instrument};
 
 use crate::{actor::SyncHandle, ContentStatus, ContentStatusCallback, Entry, NamespaceId};
@@ -49,7 +49,7 @@ pub struct Engine {
     pub default_author: Arc<DefaultAuthor>,
     to_live_actor: mpsc::Sender<ToLiveActor>,
     #[allow(dead_code)]
-    actor_handle: SharedAbortingJoinHandle<()>,
+    actor_handle: Arc<AbortOnDropHandle<()>>,
     #[debug("ContentStatusCallback")]
     content_status_cb: ContentStatusCallback,
 }
@@ -108,7 +108,7 @@ impl Engine {
             endpoint,
             sync,
             to_live_actor: live_actor_tx,
-            actor_handle: actor_handle.into(),
+            actor_handle: Arc::new(AbortOnDropHandle::new(actor_handle)),
             content_status_cb,
             default_author: Arc::new(default_author),
         })

iroh-docs/src/engine.rs

@@ -35,7 +35,7 @@ futures-concurrency = { version = "7.6.1", optional = true }
 futures-util = { version = "0.3.30", optional = true }
 iroh-net = { path = "../iroh-net", version = "0.24.0", optional = true, default-features = false }
 tokio = { version = "1", optional = true, features = ["io-util", "sync", "rt", "macros", "net", "fs"] }
-tokio-util = { version = "0.7.8", optional = true, features = ["codec"] }
+tokio-util = { version = "0.7.12", optional = true, features = ["codec", "rt"] }
 tracing = "0.1"
 
 [dev-dependencies]

iroh-gossip/Cargo.toml

@@ -13,18 +13,19 @@ use iroh_net::{
     dialer::Dialer,
     endpoint::{get_remote_node_id, Connection},
     key::PublicKey,
-    util::SharedAbortingJoinHandle,
     AddrInfo, Endpoint, NodeAddr, NodeId,
 };
 use rand::rngs::StdRng;
 use rand_core::SeedableRng;
 use std::{
     collections::{BTreeSet, HashMap, HashSet, VecDeque},
     pin::Pin,
+    sync::Arc,
     task::{Context, Poll},
     time::Instant,
 };
 use tokio::{sync::mpsc, task::JoinSet};
+use tokio_util::task::AbortOnDropHandle;
 use tracing::{debug, error_span, trace, warn, Instrument};
 
 use self::util::{read_message, write_message, Timers};
@@ -90,7 +91,7 @@ type ProtoMessage = proto::Message<PublicKey>;
 pub struct Gossip {
     to_actor_tx: mpsc::Sender<ToActor>,
     on_direct_addrs_tx: mpsc::Sender<Vec<iroh_net::endpoint::DirectAddr>>,
-    _actor_handle: SharedAbortingJoinHandle<()>,
+    _actor_handle: Arc<AbortOnDropHandle<()>>,
     max_message_size: usize,
 }
 
@@ -138,7 +139,7 @@ impl Gossip {
         Self {
             to_actor_tx,
             on_direct_addrs_tx: on_endpoints_tx,
-            _actor_handle: actor_handle.into(),
+            _actor_handle: Arc::new(AbortOnDropHandle::new(actor_handle)),
             max_message_size,
         }
     }

iroh-gossip/src/net.rs

@@ -66,7 +66,7 @@ tokio = { version = "1", features = ["io-util", "macros", "sync", "rt", "net", "
 tokio-rustls = { version = "0.26", default-features = false, features = ["logging", "ring"] }
 tokio-tungstenite = "0.21"
 tokio-tungstenite-wasm = "0.3"
-tokio-util = { version = "0.7", features = ["io-util", "io", "codec"] }
+tokio-util = { version = "0.7.12", features = ["io-util", "io", "codec", "rt"] }
 tracing = "0.1"
 tungstenite = "0.21"
 url = { version = "2.4", features = ["serde"] }

iroh-net/Cargo.toml

@@ -343,8 +343,9 @@ mod tests {
 
     use parking_lot::Mutex;
     use rand::Rng;
+    use tokio_util::task::AbortOnDropHandle;
 
-    use crate::{key::SecretKey, relay::RelayMode, util::AbortingJoinHandle};
+    use crate::{key::SecretKey, relay::RelayMode};
 
     use super::*;
 
@@ -588,7 +589,7 @@ mod tests {
     async fn new_endpoint(
         secret: SecretKey,
         disco: impl Discovery + 'static,
-    ) -> (Endpoint, AbortingJoinHandle<anyhow::Result<()>>) {
+    ) -> (Endpoint, AbortOnDropHandle<anyhow::Result<()>>) {
         let ep = Endpoint::builder()
             .secret_key(secret)
             .discovery(Box::new(disco))
@@ -611,7 +612,7 @@ mod tests {
             }
         });
 
-        (ep, AbortingJoinHandle::from(handle))
+        (ep, AbortOnDropHandle::new(handle))
     }
 
     fn system_time_now() -> u64 {
@@ -632,6 +633,7 @@ mod test_dns_pkarr {
 
     use anyhow::Result;
     use iroh_base::key::SecretKey;
+    use tokio_util::task::AbortOnDropHandle;
 
     use crate::{
         discovery::pkarr::PkarrPublisher,
@@ -642,7 +644,6 @@ mod test_dns_pkarr {
             pkarr_dns_state::State,
             run_relay_server, DnsPkarrServer,
         },
-        util::AbortingJoinHandle,
         AddrInfo, Endpoint, NodeAddr,
     };
 
@@ -753,7 +754,7 @@ mod test_dns_pkarr {
     async fn ep_with_discovery(
         relay_map: &RelayMap,
         dns_pkarr_server: &DnsPkarrServer,
-    ) -> Result<(Endpoint, AbortingJoinHandle<Result<()>>)> {
+    ) -> Result<(Endpoint, AbortOnDropHandle<Result<()>>)> {
         let secret_key = SecretKey::generate();
         let ep = Endpoint::builder()
             .relay_mode(RelayMode::Custom(relay_map.clone()))
@@ -778,6 +779,6 @@ mod test_dns_pkarr {
             }
         });
 
-        Ok((ep, AbortingJoinHandle::from(handle)))
+        Ok((ep, AbortOnDropHandle::new(handle)))
     }
 }

iroh-net/src/discovery.rs

@@ -17,10 +17,10 @@ use tracing::{debug, error, trace, warn};
 use iroh_base::key::PublicKey;
 use swarm_discovery::{Discoverer, DropGuard, IpClass, Peer};
 use tokio::{sync::mpsc, task::JoinSet};
+use tokio_util::task::AbortOnDropHandle;
 
 use crate::{
     discovery::{Discovery, DiscoveryItem},
-    util::AbortingJoinHandle,
     AddrInfo, Endpoint, NodeId,
 };
 
@@ -37,7 +37,7 @@ const DISCOVERY_DURATION: Duration = Duration::from_secs(10);
 #[derive(Debug)]
 pub struct LocalSwarmDiscovery {
     #[allow(dead_code)]
-    handle: AbortingJoinHandle<()>,
+    handle: AbortOnDropHandle<()>,
     sender: mpsc::Sender<Message>,
 }
 
@@ -192,7 +192,7 @@ impl LocalSwarmDiscovery {
             }
         });
         Ok(Self {
-            handle: handle.into(),
+            handle: AbortOnDropHandle::new(handle),
             sender: send,
         })
     }

iroh-net/src/discovery/local_swarm_discovery.rs

@@ -18,7 +18,6 @@ use crate::{
     },
     dns::node_info::NodeInfo,
     key::SecretKey,
-    util::AbortingJoinHandle,
     AddrInfo, Endpoint, NodeId,
 };
 use futures_lite::StreamExt;
@@ -27,6 +26,7 @@ use pkarr::{
     PkarrClient, PkarrClientAsync, PkarrRelayClient, PkarrRelayClientAsync, PublicKey,
     RelaySettings, SignedPacket,
 };
+use tokio_util::task::AbortOnDropHandle;
 use url::Url;
 
 /// Republish delay for the DHT. This is only for when the info does not change.
@@ -59,8 +59,9 @@ struct Inner {
     #[debug("Option<PkarrRelayClientAsync>")]
     pkarr_relay: Option<PkarrRelayClientAsync>,
     /// The background task that periodically publishes the node address.
-    /// Due to AbortingJoinHandle, this will be aborted when the discovery is dropped.
-    task: Mutex<Option<AbortingJoinHandle<()>>>,
+    ///
+    /// Due to [`AbortOnDropHandle`], this will be aborted when the discovery is dropped.
+    task: Mutex<Option<AbortOnDropHandle<()>>>,
     /// Optional keypair for signing the DNS packets.
     ///
     /// If this is None, the node will not publish its address to the DHT.
@@ -370,7 +371,7 @@ impl Discovery for DhtDiscovery {
         let this = self.clone();
         let curr = tokio::spawn(this.publish_loop(keypair.clone(), signed_packet));
         let mut task = self.0.task.lock().unwrap();
-        *task = Some(curr.into());
+        *task = Some(AbortOnDropHandle::new(curr));
     }
 
     fn resolve(

iroh-net/src/discovery/pkarr/dht.rs

@@ -133,7 +133,7 @@ pub mod relay;
 pub mod stun;
 pub mod ticket;
 pub mod tls;
-pub mod util;
+pub(crate) mod util;
 
 pub use endpoint::{AddrInfo, Endpoint, NodeAddr};
 

iroh-net/src/lib.rs

@@ -2816,11 +2816,11 @@ mod tests {
     use anyhow::Context;
     use iroh_test::CallOnDrop;
     use rand::RngCore;
+    use tokio_util::task::AbortOnDropHandle;
 
     use crate::defaults::staging::EU_RELAY_HOSTNAME;
     use crate::relay::RelayMode;
     use crate::tls;
-    use crate::util::AbortingJoinHandle;
     use crate::Endpoint;
 
     use super::*;
@@ -3761,7 +3761,7 @@ mod tests {
             }
             .instrument(info_span!("ep2.accept, me = node_id_2.fmt_short()"))
         });
-        let _accept_task = AbortingJoinHandle::from(accept_task);
+        let _accept_task = AbortOnDropHandle::new(accept_task);
 
         let node_addr_2 = NodeAddr {
             node_id: node_id_2,
@@ -3826,7 +3826,7 @@ mod tests {
             }
             .instrument(info_span!("ep2.accept", me = node_id_2.fmt_short()))
         });
-        let _accept_task = AbortingJoinHandle::from(accept_task);
+        let _accept_task = AbortOnDropHandle::new(accept_task);
 
         // Add an empty entry in the NodeMap of ep_1
         msock_1.node_map.add_node_addr(

iroh-net/src/magicsock.rs

@@ -17,12 +17,12 @@ use iroh_metrics::inc;
 use tokio::sync::{self, mpsc, oneshot};
 use tokio::time::{Duration, Instant};
 use tokio_util::sync::CancellationToken;
+use tokio_util::task::AbortOnDropHandle;
 use tracing::{debug, error, info_span, trace, warn, Instrument};
 
 use crate::dns::DnsResolver;
 use crate::net::{IpFamily, UdpSocket};
 use crate::relay::RelayUrl;
-use crate::util::CancelOnDrop;
 
 use super::portmapper;
 use super::relay::RelayMap;
@@ -174,7 +174,7 @@ pub struct Client {
     /// the actor will terminate.
     addr: Addr,
     /// Ensures the actor is terminated when the client is dropped.
-    _drop_guard: Arc<CancelOnDrop>,
+    _drop_guard: Arc<AbortOnDropHandle<()>>,
 }
 
 #[derive(Debug)]
@@ -210,7 +210,7 @@ impl Client {
         let addr = actor.addr();
         let task =
             tokio::spawn(async move { actor.run().await }.instrument(info_span!("netcheck.actor")));
-        let drop_guard = CancelOnDrop::new("netcheck actor", task.abort_handle());
+        let drop_guard = AbortOnDropHandle::new(task);
         Ok(Client {
             addr,
             _drop_guard: Arc::new(drop_guard),