[Rust] Add `OwnedBackgroundTaskGuard` for finishing background task automatically

1 files changed, 39 insertions, 4 deletions

diff --git a/rust/src/background_task.rs b/rust/src/background_task.rs
index 1e993355..94cc3368 100644
--- a/rust/src/background_task.rs
+++ b/rust/src/background_task.rs
@@ -24,14 +24,40 @@ use crate::string::*;
 
 pub type Result<R> = result::Result<R, ()>;
 
+/// An RAII guard for [`BackgroundTask`] to finish the task when dropped.
+pub struct OwnedBackgroundTaskGuard {
+    pub(crate) task: Ref<BackgroundTask>,
+}
+
+impl OwnedBackgroundTaskGuard {
+    pub fn cancel(&mut self) {
+        self.task.cancel();
+    }
+
+    pub fn is_cancelled(&self) -> bool {
+        self.task.is_cancelled()
+    }
+
+    pub fn set_progress_text(&mut self, text: &str) {
+        self.task.set_progress_text(text);
+    }
+}
+
+impl Drop for OwnedBackgroundTaskGuard {
+    fn drop(&mut self) {
+        self.task.finish();
+    }
+}
+
 /// A [`BackgroundTask`] does not actually execute any code, only act as a handler, primarily to query
 /// the status of the task, and to cancel the task.
 ///
-/// If you are looking to execute code in the background consider using rusts threading API, or if you
-/// want the core to execute the task on a worker thread, use the [`crate::worker_thread`] API.
+/// If you are looking to execute code in the background, consider using rusts threading API, or if you
+/// want the core to execute the task on a worker thread, instead use the [`crate::worker_thread`] API.
 ///
-/// NOTE: If you do not call [`BackgroundTask::finish`] or [`BackgroundTask::cancel`] the task will
-/// persist even _after_ it has been dropped.
+/// NOTE: If you do not call [`BackgroundTask::finish`] or [`BackgroundTask::cancel`], the task will
+/// persist even _after_ it has been dropped, use [`OwnedBackgroundTaskGuard`] to ensure the task is
+/// finished, see [`BackgroundTask::enter`] for usage.
 #[derive(PartialEq, Eq, Hash)]
 pub struct BackgroundTask {
     pub(crate) handle: *mut BNBackgroundTask,
@@ -52,6 +78,15 @@ impl BackgroundTask {
         unsafe { Ref::new(Self { handle }) }
     }
 
+    /// Creates a [`OwnedBackgroundTaskGuard`] that is responsible for finishing the background task
+    /// once dropped. Because the status of a task does not dictate the underlying objects' lifetime,
+    /// this can be safely done without requiring exclusive ownership.
+    pub fn enter(&self) -> OwnedBackgroundTaskGuard {
+        OwnedBackgroundTaskGuard {
+            task: self.to_owned(),
+        }
+    }
+
     pub fn can_cancel(&self) -> bool {
         unsafe { BNCanCancelBackgroundTask(self.handle) }
     }