diff --git a/include/margo.h b/include/margo.h index b4d51f54e28da2de3d19ed46f79191ec6e7919e5..f0229747924f6f358db91f2e985b17eca351567b 100644 --- a/include/margo.h +++ b/include/margo.h @@ -100,6 +100,23 @@ void margo_finalize( void margo_wait_for_finalize( margo_instance_id mid); +/** + * Installs a callback to be called before the margo instance is finalize. + * Callbacks installed will be called in reverse ordered than they have been + * pushed, and with the user-provider pointer as argument. + * + * Note that callbacks may not be called within margo_finalize. They are called + * when the margo instance is cleaned up, which may happen in margo_wait_for_finalize. + * + * @param mid The margo instance + * @param cb Callback to install + * @param uargs User-provider argument to pass to the callback when called + */ +void margo_push_finalize_callback( + margo_instance_id mid, + void(*cb)(void*), + void* uargs); + /** * Registers an RPC with margo * diff --git a/src/margo.c b/src/margo.c index 0ce1b3c490d9e103bb9721f7b4b909ba36024560..36fe8cd54179c0139643cdbe9aa996ef43ad9721 100644 --- a/src/margo.c +++ b/src/margo.c @@ -63,6 +63,13 @@ struct margo_handle_cache_el struct margo_handle_cache_el *next; /* free list link */ }; +struct margo_finalize_cb +{ + void(*callback)(void*); + void* uargs; + struct margo_finalize_cb* next; +}; + struct margo_instance { /* mercury/argobots state */ @@ -87,6 +94,7 @@ struct margo_instance int refcount; ABT_mutex finalize_mutex; ABT_cond finalize_cond; + struct margo_finalize_cb* finalize_cb; /* hash table to track multiplexed rpcs registered with margo */ struct mplex_element *mplex_table; @@ -279,6 +287,7 @@ margo_instance_id margo_init_pool(ABT_pool progress_pool, ABT_pool handler_pool, mid->hg_progress_timeout_ub = DEFAULT_MERCURY_PROGRESS_TIMEOUT_UB; mid->mplex_table = NULL; mid->refcount = 1; + mid->finalize_cb = NULL; ret = margo_timer_instance_init(mid); if(ret != 0) goto err; @@ -309,6 +318,15 @@ static void margo_cleanup(margo_instance_id mid) { int i; + /* call finalize callbacks */ + struct margo_finalize_cb* fcb = mid->finalize_cb; + while(fcb) { + (fcb->callback)(fcb->uargs); + struct margo_finalize_cb* tmp = fcb; + fcb = fcb->next; + free(tmp); + } + margo_timer_instance_finalize(mid); /* delete the hash used for multiplexing */ @@ -399,6 +417,23 @@ void margo_wait_for_finalize(margo_instance_id mid) return; } +void margo_push_finalize_callback( + margo_instance_id mid, + void(*cb)(void*), + void* uargs) +{ + if(cb == NULL) return; + + struct margo_finalize_cb* fcb = + (struct margo_finalize_cb*)malloc(sizeof(*fcb)); + fcb->callback = cb; + fcb->uargs = uargs; + + struct margo_finalize_cb* next = mid->finalize_cb; + fcb->next = next; + mid->finalize_cb = fcb; +} + hg_id_t margo_register_name(margo_instance_id mid, const char *func_name, hg_proc_cb_t in_proc_cb, hg_proc_cb_t out_proc_cb, hg_rpc_cb_t rpc_cb) {