From 6986e2d1d10e00a331ae13fee802a458a8b67d62 Mon Sep 17 00:00:00 2001
From: Matthieu Dorier <mdorier@anl.gov>
Date: Sat, 17 Aug 2024 09:47:30 +0100
Subject: [PATCH] better cleanup

---
 src/margo-core.c      | 19 +++++++------------
 src/margo-hg-config.c | 15 +++------------
 src/margo-hg-config.h |  2 +-
 src/margo-init.c      |  2 +-
 4 files changed, 12 insertions(+), 26 deletions(-)

diff --git a/src/margo-core.c b/src/margo-core.c
index 06a920b8..d13fdb91 100644
--- a/src/margo-core.c
+++ b/src/margo-core.c
@@ -145,7 +145,7 @@ static void margo_cleanup(margo_instance_id mid)
     /* finalize Mercury before anything else because this
      * could trigger some margo_cb for forward operations that
      * have not completed yet (cancelling them) */
-    __margo_hg_destroy(&(mid->hg), false);
+    __margo_hg_destroy(&(mid->hg));
 
     if (mid->abt_profiling_enabled) {
         MARGO_TRACE(mid, "Dumping ABT profile");
@@ -180,11 +180,8 @@ static void margo_cleanup(margo_instance_id mid)
         mid->registered_rpcs = next_rpc;
     }
 
-    if (mid->refcount == 0) {
-        __margo_hg_destroy(&(mid->hg), true);
-        __margo_abt_destroy(&(mid->abt));
-        free(mid);
-    }
+    __margo_abt_destroy(&(mid->abt));
+    free(mid);
 
     MARGO_TRACE(0, "Completed margo_cleanup");
 }
@@ -212,9 +209,7 @@ hg_return_t margo_instance_release(margo_instance_id mid)
         if (!mid->finalize_flag) {
             margo_finalize(mid);
         } else {
-            __margo_hg_destroy(&(mid->hg), true);
-            __margo_abt_destroy(&(mid->abt));
-            free(mid);
+            margo_cleanup(mid);
         }
     }
     return HG_SUCCESS;
@@ -274,7 +269,7 @@ void margo_finalize(margo_instance_id mid)
 
     ABT_mutex_lock(mid->finalize_mutex);
     mid->finalize_flag = true;
-    do_cleanup         = mid->finalize_refcount == 0;
+    do_cleanup         = mid->finalize_refcount == 0 && mid->refcount == 0;
 
     ABT_mutex_unlock(mid->finalize_mutex);
     ABT_cond_broadcast(mid->finalize_cond);
@@ -307,7 +302,7 @@ void margo_finalize_and_wait(margo_instance_id mid)
         ABT_cond_wait(mid->finalize_cond, mid->finalize_mutex);
 
     mid->finalize_refcount--;
-    do_cleanup = mid->finalize_refcount == 0;
+    do_cleanup = mid->finalize_refcount == 0 && mid->refcount == 0;
 
     ABT_mutex_unlock(mid->finalize_mutex);
 
@@ -330,7 +325,7 @@ void margo_wait_for_finalize(margo_instance_id mid)
         ABT_cond_wait(mid->finalize_cond, mid->finalize_mutex);
 
     mid->finalize_refcount--;
-    do_cleanup = mid->finalize_refcount == 0;
+    do_cleanup = mid->finalize_refcount == 0 && mid->refcount == 0;
 
     ABT_mutex_unlock(mid->finalize_mutex);
 
diff --git a/src/margo-hg-config.c b/src/margo-hg-config.c
index 4ccc7261..23dddcd3 100644
--- a/src/margo-hg-config.c
+++ b/src/margo-hg-config.c
@@ -268,7 +268,7 @@ bool __margo_hg_init_from_json(const struct json_object*   json,
     return true;
 
 error:
-    __margo_hg_destroy(hg, true);
+    __margo_hg_destroy(hg);
     return false;
 }
 
@@ -424,7 +424,7 @@ struct json_object* __margo_hg_to_json(const margo_hg_t* hg)
     return json;
 }
 
-void __margo_hg_destroy(margo_hg_t* hg, bool free_class)
+void __margo_hg_destroy(margo_hg_t* hg)
 {
     free((char*)hg->hg_init_info.sm_info_string);
     hg->hg_init_info.sm_info_string = NULL;
@@ -444,19 +444,10 @@ void __margo_hg_destroy(margo_hg_t* hg, bool free_class)
         hg->hg_context = NULL;
     }
 
-    hg_class_t* hg_class     = hg->hg_class;
-    uint8_t     hg_ownership = hg->hg_ownership;
-
-    if (hg->hg_class && (hg->hg_ownership & MARGO_OWNS_HG_CLASS)
-        && free_class) {
+    if (hg->hg_class && (hg->hg_ownership & MARGO_OWNS_HG_CLASS)) {
         HG_Finalize(hg->hg_class);
         hg->hg_class = NULL;
     }
 
     memset(hg, 0, sizeof(*hg));
-
-    if (!free_class) {
-        hg->hg_class     = hg_class;
-        hg->hg_ownership = hg_ownership;
-    }
 }
diff --git a/src/margo-hg-config.h b/src/margo-hg-config.h
index 77a9c48e..1b90fb05 100644
--- a/src/margo-hg-config.h
+++ b/src/margo-hg-config.h
@@ -51,6 +51,6 @@ bool                __margo_hg_init_from_json(const struct json_object*,
                                               const margo_hg_user_args_t*,
                                               margo_hg_t*);
 struct json_object* __margo_hg_to_json(const margo_hg_t*);
-void                __margo_hg_destroy(margo_hg_t*, bool free_class);
+void                __margo_hg_destroy(margo_hg_t*);
 
 #endif
diff --git a/src/margo-init.c b/src/margo-init.c
index 34926c30..379c4b35 100644
--- a/src/margo-init.c
+++ b/src/margo-init.c
@@ -385,7 +385,7 @@ margo_instance_id margo_init_ext(const char*                   address,
         if (mid->current_rpc_id_key) ABT_key_free(&(mid->current_rpc_id_key));
         free(mid);
     }
-    __margo_hg_destroy(&hg, true);
+    __margo_hg_destroy(&hg);
     __margo_abt_destroy(&abt);
     mid = MARGO_INSTANCE_NULL;
     goto finish;