chore: replace CancellationToken with AbortableJoinHandle

Satellite-im · Dec 4, 2024 · df3edf9 · df3edf9
1 parent f7992d7
commit df3edf9
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 62 deletions.
diff --git a/extensions/warp-ipfs/src/store/event_subscription.rs b/extensions/warp-ipfs/src/store/event_subscription.rs
@@ -1,4 +1,4 @@
-use crate::rt::{Executor, LocalExecutor};
+use crate::rt::{AbortableJoinHandle, Executor, LocalExecutor};
 use futures::{
     channel::{
         mpsc::{channel, Receiver, Sender},
@@ -7,11 +7,9 @@ use futures::{
     stream::BoxStream,
     SinkExt, StreamExt,
 };
+use std::collections::VecDeque;
 use std::fmt::Debug;
 use std::task::{Poll, Waker};
-use std::{collections::VecDeque, sync::Arc};
-use tokio::select;
-use tokio_util::sync::{CancellationToken, DropGuard};
 use warp::error::Error;
 
 #[allow(clippy::large_enum_variant)]
@@ -27,7 +25,7 @@ enum Command<T: Clone + Debug + Send + 'static> {
 #[derive(Clone, Debug)]
 pub struct EventSubscription<T: Clone + Debug + Send + 'static> {
     tx: Sender<Command<T>>,
-    _task_cancellation: Arc<DropGuard>,
+    _handle: AbortableJoinHandle<()>,
 }
 
 impl<T: Clone + Debug + Send + 'static> EventSubscription<T> {
@@ -43,19 +41,9 @@ impl<T: Clone + Debug + Send + 'static> EventSubscription<T> {
             rx,
         };
 
-        let token = CancellationToken::new();
-        let drop_guard = token.clone().drop_guard();
-        executor.dispatch(async move {
-            select! {
-                _ = token.cancelled() => {}
-                _ = task.run() => {}
-            }
-        });
+        let _handle = executor.spawn_abortable(async move { task.run().await });
 
-        Self {
-            tx,
-            _task_cancellation: Arc::new(drop_guard),
-        }
+        Self { tx, _handle }
     }
 
     pub async fn subscribe<'a>(&self) -> Result<BoxStream<'a, T>, Error> {

diff --git a/extensions/warp-ipfs/src/store/files.rs b/extensions/warp-ipfs/src/store/files.rs
@@ -15,7 +15,6 @@ use futures_finally::try_stream::FinallyTryStreamExt;
 
 use rust_ipfs::{unixfs::UnixfsStatus, Ipfs, IpfsPath};
 
-use tokio_util::sync::{CancellationToken, DropGuard};
 use tracing::{Instrument, Span};
 use warp::{
     constellation::{
@@ -31,7 +30,7 @@ use super::{
     document::root::RootDocumentMap, event_subscription::EventSubscription,
     MAX_THUMBNAIL_STREAM_SIZE,
 };
-use crate::rt::{Executor, LocalExecutor};
+use crate::rt::{AbortableJoinHandle, Executor, LocalExecutor};
 use crate::{
     config::{self, Config},
     thumbnail::ThumbnailGenerator,
@@ -44,7 +43,7 @@ pub struct FileStore {
     path: Arc<RwLock<PathBuf>>,
     config: config::Config,
     command_sender: mpsc::Sender<FileTaskCommand>,
-    _guard: Arc<DropGuard>,
+    _handle: AbortableJoinHandle<()>,
 }
 
 impl FileStore {
@@ -55,6 +54,7 @@ impl FileStore {
         constellation_tx: EventSubscription<ConstellationEventKind>,
         span: &Span,
     ) -> Self {
+        let executor = LocalExecutor;
         let config = config.clone();
 
         let index = Directory::new("root");
@@ -92,24 +92,16 @@ impl FileStore {
         let signal = Some(task.signal_tx.clone());
         index.rebuild_paths(&signal);
 
-        let token = CancellationToken::new();
-        let _guard = Arc::new(token.clone().drop_guard());
-
         let span = span.clone();
 
-        LocalExecutor.dispatch(async move {
-            tokio::select! {
-                _ = task.run().instrument(span) => {}
-                _ = token.cancelled() => {}
-            }
-        });
+        let _handle = executor.spawn_abortable(async move { task.run().await }.instrument(span));
 
         FileStore {
             index,
             config,
             path,
             command_sender,
-            _guard,
+            _handle,
         }
     }
 }

diff --git a/extensions/warp-ipfs/src/store/message.rs b/extensions/warp-ipfs/src/store/message.rs
@@ -29,7 +29,6 @@ use ipld_core::cid::Cid;
 use rust_ipfs::{Ipfs, PeerId};
 
 use serde::{Deserialize, Serialize};
-use tokio_util::sync::{CancellationToken, DropGuard};
 use uuid::Uuid;
 
 use super::{document::root::RootDocumentMap, ds_key::DataStoreKey, PeerIdExt};
@@ -79,7 +78,7 @@ pub type DownloadStream = BoxStream<'static, Result<Bytes, std::io::Error>>;
 #[derive(Clone)]
 pub struct MessageStore {
     inner: Arc<tokio::sync::RwLock<ConversationInner>>,
-    _task_cancellation: Arc<DropGuard>,
+    _handle: AbortableJoinHandle<()>,
 }
 
 impl MessageStore {
@@ -94,9 +93,6 @@ impl MessageStore {
         let executor = LocalExecutor;
         tracing::info!("Initializing MessageStore");
 
-        let token = CancellationToken::new();
-        let drop_guard = token.clone().drop_guard();
-
         let root = identity.root_document().clone();
 
         let mut inner = ConversationInner {
@@ -127,19 +123,9 @@ impl MessageStore {
             identity: identity.clone(),
         };
 
-        executor.dispatch({
-            async move {
-                tokio::select! {
-                    _ = token.cancelled() => {}
-                    _ = task.run() => {}
-                }
-            }
-        });
+        let _handle = executor.spawn_abortable(task.run());
 
-        Self {
-            inner,
-            _task_cancellation: Arc::new(drop_guard),
-        }
+        Self { inner, _handle }
     }
 }
 

diff --git a/extensions/warp-ipfs/src/store/queue.rs b/extensions/warp-ipfs/src/store/queue.rs
@@ -1,14 +1,13 @@
 use futures::{channel::mpsc, StreamExt, TryFutureExt};
 
-use crate::rt::{Executor, LocalExecutor};
+use crate::rt::{AbortableJoinHandle, Executor, LocalExecutor};
 use crate::store::{
     ds_key::DataStoreKey, ecdh_encrypt, payload::PayloadBuilder, topics::PeerTopic, PeerIdExt,
 };
 use ipld_core::cid::Cid;
 use rust_ipfs::{Ipfs, Keypair};
 use std::{collections::HashMap, sync::Arc, time::Duration};
 use tokio::sync::RwLock;
-use tokio_util::sync::{CancellationToken, DropGuard};
 use warp::{crypto::DID, error::Error};
 use web_time::Instant;
 
@@ -228,7 +227,7 @@ pub struct QueueEntry {
     recipient: DID,
     keypair: Keypair,
     item: RequestResponsePayload,
-    drop_guard: Arc<RwLock<Option<DropGuard>>>,
+    drop_guard: Arc<RwLock<Option<AbortableJoinHandle<()>>>>,
     executor: LocalExecutor,
 }
 
@@ -249,9 +248,6 @@ impl QueueEntry {
             executor: LocalExecutor,
         };
 
-        let token = CancellationToken::new();
-        let drop_guard = token.clone().drop_guard();
-
         let fut = {
             let entry = entry.clone();
             async move {
@@ -332,16 +328,9 @@ impl QueueEntry {
             }
         };
 
-        entry.executor.dispatch(async move {
-            futures::pin_mut!(fut);
-
-            tokio::select! {
-                _ = token.cancelled() => {}
-                _ = &mut fut => {}
-            }
-        });
+        let _handle = entry.executor.spawn_abortable(fut);
 
-        *entry.drop_guard.write().await = Some(drop_guard);
+        *entry.drop_guard.write().await = Some(_handle);
 
         entry
     }