opal/mca/common/ucx: Use new TSD api

Co-authored-by: Artem Y. Polyakov <artemp@mellanox.com>

Signed-off-by: Tomislav Janjusic <tomislavj@mellanox.com>

ompi/mca/osc/ucx/osc_ucx_component.c

@@ -695,8 +695,9 @@ int ompi_osc_ucx_free(struct ompi_win_t *win) {
 
     opal_common_ucx_wpctx_release(module->ctx);
 
-    if (module->disp_units)
+    if (module->disp_units) {
         free(module->disp_units);
+    }
     ompi_comm_free(&module->comm);
 
     free(module);

opal/mca/common/ucx/common_ucx_wpool.c

@@ -33,7 +33,7 @@ __thread int initialized = 0;
 static _ctx_record_t *
 _tlocal_add_ctx_rec(opal_common_ucx_ctx_t *ctx);
 static inline _ctx_record_t *
-_tlocal_get_ctx_rec(opal_tsd_key_t tls_key);
+_tlocal_get_ctx_rec(opal_tsd_tracked_key_t tls_key);
 static void _tlocal_ctx_rec_cleanup(_ctx_record_t *ctx_rec);
 static void _tlocal_mem_rec_cleanup(_mem_record_t *mem_rec);
 static void _ctx_rec_destructor(void *arg);
@@ -399,7 +399,8 @@ opal_common_ucx_wpctx_create(opal_common_ucx_wpool_t *wpool, int comm_size,
         goto error;
     }
 
-    opal_tsd_key_create(&ctx->tls_key, _ctx_rec_destructor);
+    OBJ_CONSTRUCT(&ctx->tls_key, opal_tsd_tracked_key_t);
+    opal_tsd_tracked_key_set_destructor(&ctx->tls_key, _ctx_rec_destructor);
 
     (*ctx_ptr) = ctx;
     return ret;
@@ -414,16 +415,16 @@ error:
 OPAL_DECLSPEC void
 opal_common_ucx_wpctx_release(opal_common_ucx_ctx_t *ctx)
 {
-    _ctx_record_t *ctx_rec = NULL;
+    _ctx_record_t *ctx_rec = NULL, *next;
 
     /* Application is expected to guarantee that no operation
      * is performed on the context that is being released */
 
     /* destroy key so that other threads don't invoke destructors */
-    opal_tsd_key_delete(ctx->tls_key);
+    OBJ_DESTRUCT(&ctx->tls_key);
 
     /* loop through list of records */
-    OPAL_LIST_FOREACH(ctx_rec, &ctx->ctx_records, _ctx_record_t) {
+    OPAL_LIST_FOREACH_SAFE(ctx_rec, next, &ctx->ctx_records, _ctx_record_t) {
        _tlocal_ctx_rec_cleanup(ctx_rec);
     }
 
@@ -493,7 +494,8 @@ int opal_common_ucx_wpmem_create(opal_common_ucx_ctx_t *ctx,
         goto error_rkey_pack;
     }
 
-    opal_tsd_key_create(&mem->tls_key, _mem_rec_destructor);
+    OBJ_CONSTRUCT(&mem->tls_key, opal_tsd_tracked_key_t);
+    opal_tsd_tracked_key_set_destructor(&mem->tls_key, _mem_rec_destructor);
 
     (*mem_ptr) = mem;
     (*my_mem_addr) = rkey_addr;
@@ -560,12 +562,12 @@ static int _comm_ucx_wpmem_map(opal_common_ucx_wpool_t *wpool,
 
 void opal_common_ucx_wpmem_free(opal_common_ucx_wpmem_t *mem)
 {
-    _mem_record_t *mem_rec = NULL;
-    
-    opal_tsd_key_delete(mem->tls_key);
+    _mem_record_t *mem_rec = NULL, *next;
     
+    OBJ_DESTRUCT(&mem->tls_key);
+
     /* Loop through list of records */
-    OPAL_LIST_FOREACH(mem_rec, &mem->mem_records, _mem_record_t) {
+    OPAL_LIST_FOREACH_SAFE(mem_rec, next, &mem->mem_records, _mem_record_t) {
         _tlocal_mem_rec_cleanup(mem_rec);
     }
   
@@ -581,9 +583,9 @@ void opal_common_ucx_wpmem_free(opal_common_ucx_wpmem_t *mem)
 }
 
 static inline _ctx_record_t *
-_tlocal_get_ctx_rec(opal_tsd_key_t tls_key){
+_tlocal_get_ctx_rec(opal_tsd_tracked_key_t tls_key){
     _ctx_record_t *ctx_rec = NULL;
-    int rc = opal_tsd_getspecific(tls_key, (void**)&ctx_rec);
+    int rc = opal_tsd_tracked_key_get(&tls_key, (void**)&ctx_rec);
 
     if (OPAL_SUCCESS != rc) {
         return NULL;
@@ -657,7 +659,7 @@ _tlocal_add_ctx_rec(opal_common_ucx_ctx_t *ctx)
     opal_mutex_unlock(&ctx->mutex);
 
     /* Add tls reference to record */
-    rc = opal_tsd_setspecific(ctx->tls_key, ctx_rec);
+    rc = opal_tsd_tracked_key_set(&ctx->tls_key, ctx_rec);
     if (OPAL_SUCCESS != rc) {
         OBJ_RELEASE(ctx_rec);
         return NULL;
@@ -743,7 +745,7 @@ static _mem_record_t *_tlocal_add_mem_rec(opal_common_ucx_wpmem_t *mem, _ctx_rec
     
     opal_atomic_wmb();
 
-    rc = opal_tsd_setspecific(mem->tls_key, mem_rec);
+    rc = opal_tsd_tracked_key_set(&mem->tls_key, mem_rec);
     if (OPAL_SUCCESS != rc) {
         return NULL;
     }

opal/mca/common/ucx/common_ucx_wpool.h

@@ -73,7 +73,7 @@ typedef struct {
 
     /* Thread-local key to allow each thread to have
      * local information associated with this wpctx */
-    opal_tsd_key_t tls_key;
+    opal_tsd_tracked_key_t tls_key;
 
     /* UCX addressing information */
     char *recv_worker_addrs;
@@ -107,7 +107,7 @@ typedef struct {
     /* TLS item that allows each thread to
      * store endpoints and rkey arrays
      * for faster access */
-    opal_tsd_key_t tls_key;
+    opal_tsd_tracked_key_t tls_key;
 } opal_common_ucx_wpmem_t;
 
 typedef struct __winfo_list_item_t _winfo_list_item_t;
@@ -220,7 +220,7 @@ opal_common_ucx_tlocal_fetch(opal_common_ucx_wpmem_t *mem, int target,
     int rc = OPAL_SUCCESS;
 
     /* First check the fast-path */
-    rc = opal_tsd_getspecific(mem->tls_key, (void**)&mem_rec);
+    rc = opal_tsd_tracked_key_get(&mem->tls_key, (void**)&mem_rec);
     if (OPAL_SUCCESS != rc) {
         return rc;
     }
@@ -231,7 +231,7 @@ opal_common_ucx_tlocal_fetch(opal_common_ucx_wpmem_t *mem, int target,
         if (OPAL_SUCCESS != rc) {
             return rc;
         }
-        rc = opal_tsd_getspecific(mem->tls_key, (void**)&mem_rec);
+        rc = opal_tsd_tracked_key_get(&mem->tls_key, (void**)&mem_rec);
         if (OPAL_SUCCESS != rc) {
             return rc;
         }