From 19f70bd58f3254be58467fef4eeaa50846525424 Mon Sep 17 00:00:00 2001
From: dasman <dasman@onionmail.org>
Date: Tue, 30 Jul 2024 19:00:26 +0300
Subject: [PATCH] event_graph: gracefully stop pruning task

---
 bin/darkirc/src/main.rs |  3 +++
 src/event_graph/mod.rs  | 10 ++--------
 2 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/bin/darkirc/src/main.rs b/bin/darkirc/src/main.rs
index 3e2c39819..893e49b60 100644
--- a/bin/darkirc/src/main.rs
+++ b/bin/darkirc/src/main.rs
@@ -260,6 +260,8 @@ async fn realmain(args: Args, ex: Arc<Executor<'static>>) -> Result<()> {
     )
     .await?;
 
+    let prune_task = event_graph.prune_task.get().unwrap();
+
     info!("Registering EventGraph P2P protocol");
     let event_graph_ = Arc::clone(&event_graph);
     let registry = p2p.protocol_registry();
@@ -413,6 +415,7 @@ async fn realmain(args: Args, ex: Arc<Executor<'static>>) -> Result<()> {
 
     info!("Stopping IRC server");
     irc_task.stop().await;
+    prune_task.stop().await;
 
     info!("Flushing sled database...");
     let flushed_bytes = sled_db.flush_async().await?;
diff --git a/src/event_graph/mod.rs b/src/event_graph/mod.rs
index 6824ee5f2..3249125d5 100644
--- a/src/event_graph/mod.rs
+++ b/src/event_graph/mod.rs
@@ -101,7 +101,7 @@ pub struct EventGraph {
     /// `TipRep` message telling peers about our unreferenced tips.
     broadcasted_ids: RwLock<HashSet<blake3::Hash>>,
     /// DAG Pruning Task
-    prune_task: OnceCell<StoppableTaskPtr>,
+    pub prune_task: OnceCell<StoppableTaskPtr>,
     /// Event publisher, this notifies whenever an event is
     /// inserted into the DAG
     pub event_pub: PublisherPtr<Event>,
@@ -167,14 +167,13 @@ impl EventGraph {
 
         // Spawn the DAG pruning task
         if days_rotation > 0 {
-            let self__ = self_.clone();
             let prune_task = StoppableTask::new();
             let _ = self_.prune_task.set(prune_task.clone()).await;
 
             prune_task.clone().start(
                 self_.clone().dag_prune_task(days_rotation),
                 |_| async move {
-                    self__.clone()._handle_stop(sled_db).await;
+                    info!(target: "event_graph::_handle_stop()", "[EVENTGRAPH] Prune task stopped, flushing sled")
                 },
                 Error::DetachedTaskStopped,
                 ex.clone(),
@@ -188,11 +187,6 @@ impl EventGraph {
         self.days_rotation
     }
 
-    async fn _handle_stop(&self, sled_db: sled::Db) {
-        info!(target: "event_graph::_handle_stop()", "[EVENTGRAPH] Prune task stopped, flushing sled");
-        sled_db.flush_async().await.unwrap();
-    }
-
     /// Sync the DAG from connected peers
     pub async fn dag_sync(&self) -> Result<()> {
         // We do an optimistic sync where we ask all our connected peers for