From 01871f975ebb22fadfb22e0c6c0ed2fcca7a359f Mon Sep 17 00:00:00 2001
From: Teo Lemane <teo.lemane@gmail.com>
Date: Mon, 18 Dec 2023 17:26:07 +0100
Subject: [PATCH] fix

---
 include/kmtricks/signals.hpp | 11 +++++++----
 include/kmtricks/state.hpp   |  8 +++++++-
 include/kmtricks/task.hpp    |  2 +-
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/include/kmtricks/signals.hpp b/include/kmtricks/signals.hpp
index e6c4f5c..543cbab 100644
--- a/include/kmtricks/signals.hpp
+++ b/include/kmtricks/signals.hpp
@@ -66,6 +66,8 @@ inline std::string signal_to_string(int signal)
   }
 }
 
+static std::mutex SigLock;
+
 class SignalHandler
 {
  public:
@@ -95,10 +97,11 @@ class SignalHandler
 
   static void state_callback(int signal)
   {
-    state::get().write();
-    spdlog::info("Killed after {}. The run can be resumed by running the same command again.",
-                 signal_to_string(signal));
-    exit(signal);
+    std::unique_lock<std::mutex> _(SigLock);
+
+    state::get().write(
+      fmt::format(
+        "Killed after {}. The run can be resumed by running the same command again.", signal_to_string(signal)));
   }
 
   static void default_callback(int signal)
diff --git a/include/kmtricks/state.hpp b/include/kmtricks/state.hpp
index cd013cf..844ee20 100644
--- a/include/kmtricks/state.hpp
+++ b/include/kmtricks/state.hpp
@@ -48,7 +48,7 @@ namespace km {
         inf.read(reinterpret_cast<char*>(m_merge.data()), sizeof(std::uint8_t) * p);
       }
 
-      void write()
+      void write(const std::string& msg = "")
       {
         std::unique_lock<spinlock> _1(m_lock_1);
         std::unique_lock<spinlock> _2(m_lock_2);
@@ -59,6 +59,12 @@ namespace km {
         inf.write(reinterpret_cast<char*>(m_superk.data()), sizeof(std::uint8_t) * m_nb);
         inf.write(reinterpret_cast<char*>(m_count.data()), sizeof(std::uint8_t) * (m_nb*m_part));
         inf.write(reinterpret_cast<char*>(m_merge.data()), sizeof(std::uint8_t) * m_part);
+
+        if (!msg.empty())
+        {
+          spdlog::info(msg);
+          std::quick_exit(1);
+        }
       }
 
     public:
diff --git a/include/kmtricks/task.hpp b/include/kmtricks/task.hpp
index eb2d502..9665f96 100644
--- a/include/kmtricks/task.hpp
+++ b/include/kmtricks/task.hpp
@@ -403,13 +403,13 @@ class HashCountTask : public ITask
   void preprocess() {}
   void postprocess()
   {
+    state::get().count_done(m_sample_id, m_part_id);
     if (this->m_clear)
     {
       m_superk_storage->closeFile(m_part_id);
       Eraser::get().erase(m_superk_storage->getFileName(m_part_id));
     }
     this->m_finish = true;
-    state::get().count_done(m_sample_id, m_part_id);
     this->exec_callback();
   }