Using ULT stack to store RPC lineage

src/margo-core.c

@@ -24,6 +24,7 @@
 #include "margo-timer-private.h"
 #include "margo-serialization.h"
 #include "margo-id.h"
+#include "margo-rpc-lineage.h"
 #include "utlist.h"
 #include "uthash.h"
 #include "abtx_prof.h"
@@ -171,7 +172,6 @@ static void margo_cleanup(margo_instance_id mid)
     ABT_mutex_free(&mid->finalize_mutex);
     ABT_cond_free(&mid->finalize_cond);
     ABT_mutex_free(&mid->pending_operations_mtx);
-    ABT_key_free(&(mid->current_rpc_id_key));
 
     /* monitoring (destroyed before Argobots since it contains mutexes) */
     __MARGO_MONITOR(mid, FN_END, finalize, monitoring_args);
@@ -2234,17 +2234,16 @@ void __margo_internal_decr_pending(margo_instance_id mid)
 hg_return_t margo_set_current_rpc_id(margo_instance_id mid, hg_id_t parent_id)
 {
     if (mid == MARGO_INSTANCE_NULL) return HG_INVALID_ARG;
-    // rely on the fact that sizeof(void*) == sizeof(hg_id_t)
     if (parent_id == 0) parent_id = mux_id(0, MARGO_DEFAULT_PROVIDER_ID);
-    int ret = ABT_key_set(mid->current_rpc_id_key, (void*)parent_id);
+    int ret = margo_lineage_set(parent_id);
     if (ret != ABT_SUCCESS) return HG_OTHER_ERROR;
     return HG_SUCCESS;
 }
 
 hg_return_t margo_get_current_rpc_id(margo_instance_id mid, hg_id_t* parent_id)
 {
     if (mid == MARGO_INSTANCE_NULL) return HG_INVALID_ARG;
-    int ret = ABT_key_get(mid->current_rpc_id_key, (void**)parent_id);
+    int ret = margo_lineage_get(parent_id);
     if (ret != ABT_SUCCESS || *parent_id == 0) {
         *parent_id = mux_id(0, MARGO_DEFAULT_PROVIDER_ID);
         return HG_OTHER_ERROR;
@@ -2298,6 +2297,8 @@ void __margo_internal_post_wrapper_hooks(
 
     __margo_internal_decr_pending(mid);
     if (__margo_internal_finalize_requested(mid)) { margo_finalize(mid); }
+
+    margo_lineage_erase();
 }
 
 static void margo_handle_data_free(void* args)

src/margo-init.c

@@ -330,10 +330,6 @@ margo_instance_id margo_init_ext(const char*                   address,
     hret                   = __margo_handle_cache_init(mid, handle_cache_size);
     if (hret != HG_SUCCESS) goto error;
 
-    // create current_rpc_id_key ABT_key
-    ret = ABT_key_create(NULL, &(mid->current_rpc_id_key));
-    if (ret != ABT_SUCCESS) goto error;
-
     // set logger
     margo_set_logger(mid, args.logger);
 
@@ -386,7 +382,6 @@ margo_instance_id margo_init_ext(const char*                   address,
         ABT_mutex_free(&mid->finalize_mutex);
         ABT_cond_free(&mid->finalize_cond);
         ABT_mutex_free(&mid->pending_operations_mtx);
-        if (mid->current_rpc_id_key) ABT_key_free(&(mid->current_rpc_id_key));
         free(mid);
     }
     __margo_hg_destroy(&hg);

src/margo-instance.h

@@ -128,9 +128,6 @@ struct margo_instance {
     _Atomic uint64_t num_progress_calls;
     _Atomic uint64_t num_trigger_calls;
 
-    /* callpath tracking */
-    ABT_key current_rpc_id_key;
-
     /* optional diagnostics data tracking */
     int abt_profiling_enabled;
 };

src/margo-rpc-lineage.h

@@ -0,0 +1,56 @@
+/*
+ * (C) 2024 The University of Chicago
+ *
+ * See COPYRIGHT in top-level directory.
+ */
+#ifndef __MARGO_ABT_KEY_H
+#define __MARGO_ABT_KEY_H
+#include <assert.h>
+#include <stdbool.h>
+#include <unistd.h>
+#include <stdint.h>
+#include <string.h>
+#include <abt.h>
+#include <mercury.h>
+
+#define __MARGO_LINEAGE_COMMON                                    \
+    static const char* magic  = "matthieu";                       \
+    ABT_thread_attr attr      = ABT_THREAD_ATTR_NULL;             \
+    ABT_thread      ult       = ABT_THREAD_NULL;                  \
+    void*           stackaddr = NULL;                             \
+    size_t          stacksize = 0;                                \
+    int             ret;                                          \
+    ret = ABT_thread_self(&ult);                                  \
+    if(ret != ABT_SUCCESS) return ret;                            \
+    ret = ABT_thread_get_attr(ult, &attr);                        \
+    if(ret != ABT_SUCCESS) return ret;                            \
+    ret = ABT_thread_attr_get_stack(attr, &stackaddr, &stacksize);\
+    if(ret != ABT_SUCCESS) return ret;                            \
+    ABT_thread_attr_free(&attr);                                  \
+    if(!stackaddr || !stacksize) return ABT_ERR_KEY;              \
+    char* stackend = (char*)stackaddr + stacksize
+
+static inline int margo_lineage_set(hg_id_t current_rpc_id) {
+    __MARGO_LINEAGE_COMMON;
+    memcpy(stackend - 8 - sizeof(hg_id_t), magic, 8);
+    memcpy(stackend - 8, &current_rpc_id, sizeof(current_rpc_id));
+    return ABT_SUCCESS;
+}
+
+static inline int margo_lineage_erase() {
+    __MARGO_LINEAGE_COMMON;
+    memset(stackend - 8 - sizeof(hg_id_t), 0, 8 + sizeof(hg_id_t));
+    return ABT_SUCCESS;
+}
+
+static inline int margo_lineage_get(hg_id_t* current_rpc_id) {
+    __MARGO_LINEAGE_COMMON;
+    if(memcmp(stackend - 8 - sizeof(hg_id_t), magic, 8) != 0)
+        return ABT_ERR_KEY;
+    memcpy(current_rpc_id, stackend - 8, sizeof(*current_rpc_id));
+    return ABT_SUCCESS;
+}
+
+#undef __MARGO_LINEAGE_COMMON
+
+#endif