opal/common/ucx: add refcnt in tlocal_ctx_tbl entry to keep track of usage

Signed-off-by: Artem Polyakov <artpol84@gmail.com>

opal/mca/common/ucx/common_ucx_wpool.c

@@ -840,6 +840,7 @@ static void _common_ucx_tls_cleanup(_tlocal_table_t *tls)
     size = tls->ctx_tbl_size;
     for (i = 0; i < size; i++) {
         if (NULL != tls->ctx_tbl[i]->gctx){
+            assert(tls->ctx_tbl[i]->refcnt == 0);
             _tlocal_ctx_record_cleanup(tls->ctx_tbl[i]);
         }
         free(tls->ctx_tbl[i]);
@@ -909,6 +910,11 @@ _tlocal_ctx_record_cleanup(_tlocal_ctx_t *ctx_rec)
     if (NULL == ctx_rec->gctx) {
         return OPAL_SUCCESS;
     }
+
+    if (ctx_rec->refcnt > 0) {
+        return OPAL_SUCCESS;
+    }
+
     /* Remove myself from the communication context structure
      * This may result in context release as we are using
      * delayed cleanup */
@@ -934,7 +940,7 @@ _tlocal_add_ctx(_tlocal_table_t *tls, opal_common_ucx_ctx_t *ctx)
     /* Try to find available record in the TLS table
      * In parallel perform deferred cleanups */
     for (i=0; i<tls->ctx_tbl_size; i++) {
-        if (NULL != tls->ctx_tbl[i]->gctx) {
+        if (NULL != tls->ctx_tbl[i]->gctx && tls->ctx_tbl[i]->refcnt == 0) {
             if (tls->ctx_tbl[i]->gctx->released ) {
                 /* Found dirty record, need to clean first */
                 _tlocal_ctx_record_cleanup(tls->ctx_tbl[i]);
@@ -1059,6 +1065,10 @@ _tlocal_mem_record_cleanup(_tlocal_mem_t *mem_rec)
         free(mem_rec->mem_tls_ptr);
     }
 
+    assert(mem_rec->ctx_rec != NULL);
+    OPAL_ATOMIC_ADD_FETCH32(&mem_rec->ctx_rec->refcnt, -1);
+    assert(mem_rec->ctx_rec->refcnt >= 0);
+
     free(mem_rec->mem);
 
     memset(mem_rec, 0, sizeof(*mem_rec));
@@ -1107,6 +1117,9 @@ static _tlocal_mem_t *_tlocal_add_mem(_tlocal_table_t *tls,
     WPOOL_DBG_OUT("tls = %p, ctx = %p\n",
                   (void *)tls, (void*)mem->ctx);
 
+    tls->mem_tbl[free_idx]->ctx_rec = ctx_rec;
+    OPAL_ATOMIC_ADD_FETCH32(&ctx_rec->refcnt, 1);
+
     tls->mem_tbl[free_idx]->mem->worker = ctx_rec->winfo;
     tls->mem_tbl[free_idx]->mem->rkeys = calloc(mem->ctx->comm_size,
                                          sizeof(*tls->mem_tbl[free_idx]->mem->rkeys));

opal/mca/common/ucx/common_ucx_wpool_int.h

@@ -8,6 +8,7 @@
 typedef struct {
     opal_common_ucx_ctx_t *gctx;
     opal_common_ucx_winfo_t *winfo;
+    opal_atomic_int32_t refcnt;
 } _tlocal_ctx_t;
 
 typedef struct {
@@ -19,6 +20,7 @@ typedef struct {
     opal_common_ucx_wpmem_t *gmem;
     _mem_info_t *mem;
     opal_common_ucx_tlocal_fast_ptrs_t *mem_tls_ptr;
+    _tlocal_ctx_t *ctx_rec;
 } _tlocal_mem_t;
 
 typedef struct {