/**
 * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2018-2026. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */

#ifdef HAVE_CONFIG_H
#  include "config.h"
#endif

#include "cuda_ipc_cache.h"
#include "cuda_ipc_iface.h"
#include "cuda_ipc.inl"

#include <ucs/datastruct/khash.h>
#include <ucs/debug/log.h>
#include <ucs/debug/memtrack_int.h>
#include <ucs/profile/profile.h>
#include <ucs/sys/sys.h>
#include <ucs/sys/string.h>
#include <ucs/sys/ptr_arith.h>
#include <ucs/sys/uid.h>
#include <ucs/type/rwlock.h>
#include <uct/cuda/base/cuda_ctx.inl>
typedef struct uct_cuda_ipc_cache_hash_key {
    pid_t        pid;
    ucs_sys_ns_t pid_ns;
    CUdevice     cu_device;
} uct_cuda_ipc_cache_hash_key_t;

static UCS_F_ALWAYS_INLINE int
uct_cuda_ipc_cache_hash_equal(uct_cuda_ipc_cache_hash_key_t key1,
                              uct_cuda_ipc_cache_hash_key_t key2)
{
    return (key1.pid == key2.pid) && (key1.cu_device == key2.cu_device) &&
           (key1.pid_ns == key2.pid_ns);
}

static UCS_F_ALWAYS_INLINE khint32_t
uct_cuda_ipc_cache_hash_func(uct_cuda_ipc_cache_hash_key_t key)
{
    uint64_t value = key.pid ^ ((uint64_t)key.pid_ns << 32) ^
                     ((uint32_t)key.cu_device << 24);
    return kh_int64_hash_func(value);
}

KHASH_INIT(cuda_ipc_rem_cache, uct_cuda_ipc_cache_hash_key_t,
           uct_cuda_ipc_cache_t*, 1, uct_cuda_ipc_cache_hash_func,
           uct_cuda_ipc_cache_hash_equal);

/*
 * Cache limit propagation:
 *
 * The cache limits (max_regions, max_size) are configured via MD config
 * (uct_cuda_ipc_md_config_t), which is transient and not accessible after
 * md_open returns. Per-peer caches are created on-demand in
 * uct_cuda_ipc_get_remote_cache(), which has no access to the MD. Therefore
 * the limits are stored here as globals and read directly by the eviction loop.
 * Each md_open tightens the limits to min(current, configured), so the most
 * restrictive value across all MDs is always in effect.
 */
typedef struct uct_cuda_ipc_remote_cache {
    khash_t(cuda_ipc_rem_cache) hash;
    ucs_rw_spinlock_t           lock;
    unsigned long               max_regions; /**< Global max regions limit */
    size_t                      max_size;    /**< Global max total size limit */
} uct_cuda_ipc_remote_cache_t;

uct_cuda_ipc_remote_cache_t uct_cuda_ipc_remote_cache;

#if HAVE_CUDA_FABRIC
static UCS_F_ALWAYS_INLINE int
uct_cuda_ipc_rem_mpool_cache_hash_equal(CUmemFabricHandle key1,
                                        CUmemFabricHandle key2)
{
    return !memcmp(&key1, &key2, sizeof(key1));
}

/* djb2 hash function */
static UCS_F_ALWAYS_INLINE khint_t
uct_cuda_ipc_rem_mpool_cache_hash_func(CUmemFabricHandle key)
{
    khint_t hash = 5381;
    int i;

    for (i = 0; i < sizeof(key.data); i++) {
        hash = ((hash << 5) + hash) + key.data[i];
    }

    return hash;
}

KHASH_INIT(cuda_ipc_rem_mpool_cache, CUmemFabricHandle, CUmemoryPool, 1,
           uct_cuda_ipc_rem_mpool_cache_hash_func,
           uct_cuda_ipc_rem_mpool_cache_hash_equal)

typedef struct {
    khash_t(cuda_ipc_rem_mpool_cache) hash;
    pthread_rwlock_t                  lock;
} uct_cuda_ipc_rem_mpool_cache_t;

static uct_cuda_ipc_rem_mpool_cache_t uct_cuda_ipc_rem_mpool_cache;
#endif

static ucs_pgt_dir_t *uct_cuda_ipc_cache_pgt_dir_alloc(const ucs_pgtable_t *pgtable)
{
    void *ptr;
    int ret;

    ret = ucs_posix_memalign(&ptr,
                             ucs_max(sizeof(void *), UCS_PGT_ENTRY_MIN_ALIGN),
                             sizeof(ucs_pgt_dir_t), "cuda_ipc_cache_pgdir");
    return (ret == 0) ? ptr : NULL;
}

static void uct_cuda_ipc_cache_pgt_dir_release(const ucs_pgtable_t *pgtable,
                                               ucs_pgt_dir_t *dir)
{
    ucs_free(dir);
}

static void
uct_cuda_ipc_cache_region_collect_callback(const ucs_pgtable_t *pgtable,
                                           ucs_pgt_region_t *pgt_region,
                                           void *arg)
{
    ucs_list_link_t *list = arg;
    uct_cuda_ipc_cache_region_t *region;

    region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t);
    ucs_list_add_tail(list, &region->list);
}

static ucs_status_t
uct_cuda_ipc_primary_ctx_retain_and_push(CUdevice cuda_device)
{
    CUcontext cuda_ctx;
    ucs_status_t status;

    status = uct_cuda_ctx_primary_retain(cuda_device, 0, &cuda_ctx);
    if (status != UCS_OK) {
        return status;
    }

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuCtxPushCurrent(cuda_ctx));
    if (status != UCS_OK) {
        UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device));
    }

    return status;
}

static void uct_cuda_ipc_primary_ctx_pop_and_release(CUdevice cuda_device)
{
    UCT_CUDADRV_FUNC_LOG_WARN(cuCtxPopCurrent(NULL));
    UCT_CUDADRV_FUNC_LOG_WARN(cuDevicePrimaryCtxRelease(cuda_device));
}

static ucs_status_t
uct_cuda_ipc_close_memhandle_legacy(uct_cuda_ipc_cache_region_t *region)
{
    ucs_status_t status;

    status = uct_cuda_ipc_primary_ctx_retain_and_push(region->cu_dev);
    if (status != UCS_OK) {
        return status;
    }

    status = UCT_CUDADRV_FUNC_LOG_WARN(
            cuIpcCloseMemHandle((CUdeviceptr)region->mapped_addr));
    uct_cuda_ipc_primary_ctx_pop_and_release(region->cu_dev);
    return status;
}

static ucs_status_t uct_cuda_ipc_close_memhandle(uct_cuda_ipc_cache_region_t *region)
{
    ucs_status_t status;

    if ((region->key.ph.handle_type == UCT_CUDA_IPC_KEY_HANDLE_TYPE_VMM) ||
        (region->key.ph.handle_type == UCT_CUDA_IPC_KEY_HANDLE_TYPE_POSIX_FD)) {
        status = UCT_CUDADRV_FUNC_LOG_WARN(cuMemUnmap(
                    (CUdeviceptr)region->mapped_addr, region->key.b_len));
        if (status != UCS_OK) {
            return status;
        }

        return UCT_CUDADRV_FUNC_LOG_WARN(cuMemAddressFree(
                (CUdeviceptr)region->mapped_addr, region->key.b_len));
    }
#if HAVE_CUDA_FABRIC
    if (region->key.ph.handle_type == UCT_CUDA_IPC_KEY_HANDLE_TYPE_MEMPOOL) {
        return UCT_CUDADRV_FUNC_LOG_WARN(
                cuMemFree((CUdeviceptr)region->mapped_addr));
    }
#endif
    return uct_cuda_ipc_close_memhandle_legacy(region);
}

static void
uct_cuda_ipc_cache_region_remove(uct_cuda_ipc_cache_t *cache,
                                 uct_cuda_ipc_cache_region_t *region)
{
    ucs_status_t status;

    status = ucs_pgtable_remove(&cache->pgtable, &region->super);
    if (status != UCS_OK) {
        ucs_warn("failed to remove address:%p from cache (%s)",
                  (void *)region->key.d_bptr, ucs_status_string(status));
    }

    ucs_list_del(&region->lru_list);

    ucs_assert(cache->num_regions > 0);
    cache->num_regions--;
    cache->total_size -= region->key.b_len;
}

static void
uct_cuda_ipc_cache_region_destroy(uct_cuda_ipc_cache_t *cache,
                                  uct_cuda_ipc_cache_region_t *region)
{
    uct_cuda_ipc_cache_region_remove(cache, region);

    ucs_trace("%s: destroy region " UCS_PGT_REGION_FMT " size:%lu",
              cache->name, UCS_PGT_REGION_ARG(&region->super),
              region->key.b_len);

    uct_cuda_ipc_close_memhandle(region);
    ucs_free(region);
}

static void uct_cuda_ipc_cache_evict_lru(uct_cuda_ipc_cache_t *cache)
{
    unsigned long max_regions = uct_cuda_ipc_remote_cache.max_regions;
    size_t max_size           = uct_cuda_ipc_remote_cache.max_size;
    uct_cuda_ipc_cache_region_t *region, *tmp;

    ucs_list_for_each_safe(region, tmp, &cache->lru_list, lru_list) {
        if ((cache->num_regions <= max_regions) &&
            (cache->total_size  <= max_size)) {
            break;
        }

        if (region->refcount > 0) {
            /* In-use region -- keep on LRU, revisit on next eviction pass */
            continue;
        }

        uct_cuda_ipc_cache_region_destroy(cache, region);
    }
}

static void uct_cuda_ipc_cache_purge(uct_cuda_ipc_cache_t *cache)
{
    uct_cuda_ipc_cache_region_t *region, *tmp;
    ucs_list_link_t region_list;

    ucs_list_head_init(&region_list);
    ucs_pgtable_purge(&cache->pgtable, uct_cuda_ipc_cache_region_collect_callback,
                      &region_list);
    ucs_list_for_each_safe(region, tmp, &region_list, list) {
        uct_cuda_ipc_close_memhandle(region);
        ucs_free(region);
    }

    ucs_list_head_init(&cache->lru_list);
    cache->num_regions = 0;
    cache->total_size  = 0;
    ucs_trace("%s: cuda ipc cache purged", cache->name);
}

static ucs_status_t
uct_cuda_ipc_open_memhandle_legacy(CUipcMemHandle memh, CUdevice cu_dev,
                                   CUdeviceptr *mapped_addr,
                                   ucs_log_level_t log_level)
{
    CUresult cuerr;
    ucs_status_t status;

    status = uct_cuda_ipc_primary_ctx_retain_and_push(cu_dev);
    if (status != UCS_OK) {
        return status;
    }

    cuerr = cuIpcOpenMemHandle(mapped_addr, memh,
                               CU_IPC_MEM_LAZY_ENABLE_PEER_ACCESS);
    if (cuerr != CUDA_SUCCESS) {
        ucs_log(log_level, "cuIpcOpenMemHandle() failed: %s",
                uct_cuda_cu_get_error_string(cuerr));
        status = (cuerr == CUDA_ERROR_ALREADY_MAPPED) ?
            UCS_ERR_ALREADY_EXISTS : UCS_ERR_INVALID_PARAM;
    }

    uct_cuda_ipc_primary_ctx_pop_and_release(cu_dev);
    return status;
}

#if HAVE_CUDA_FABRIC || HAVE_DECL_SYS_PIDFD_GETFD
static void
uct_cuda_ipc_init_access_desc(CUmemAccessDesc *access_desc, CUdevice cu_dev)
{
    access_desc->location.type = CU_MEM_LOCATION_TYPE_DEVICE;
    access_desc->flags         = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
    access_desc->location.id   = cu_dev;
}

static ucs_status_t
uct_cuda_ipc_open_memhandle_vmm(const uct_cuda_ipc_rkey_t *key, CUdevice cu_dev,
                                CUdeviceptr *mapped_addr,
                                void *shareable_handle,
                                CUmemAllocationHandleType handle_type,
                                ucs_log_level_t log_level)
{
    CUmemAccessDesc access_desc = {};
    ucs_status_t status;
    CUdeviceptr dptr;
    CUmemGenericAllocationHandle handle;

    status = UCT_CUDADRV_FUNC(cuMemImportFromShareableHandle(&handle,
                                                             shareable_handle,
                                                             handle_type),
                              log_level);
    if (status != UCS_OK) {
        goto out;
    }

    status = UCT_CUDADRV_FUNC(cuMemAddressReserve(&dptr, key->b_len, 0, 0, 0),
                              log_level);
    if (status != UCS_OK) {
        goto release_handle;
    }

    status = UCT_CUDADRV_FUNC(cuMemMap(dptr, key->b_len, 0, handle, 0),
                              log_level);
    if (status != UCS_OK) {
        goto release_va_range;
    }

    uct_cuda_ipc_init_access_desc(&access_desc, cu_dev);

    status = UCT_CUDADRV_FUNC(cuMemSetAccess(dptr, key->b_len, &access_desc, 1),
                              log_level);
    if (status != UCS_OK) {
        goto unmap_range;
    }

    *mapped_addr = dptr;
    goto release_handle;

unmap_range:
    UCT_CUDADRV_FUNC_LOG_WARN(cuMemUnmap(dptr, key->b_len));
release_va_range:
    UCT_CUDADRV_FUNC_LOG_WARN(cuMemAddressFree(dptr, key->b_len));
release_handle:
    UCT_CUDADRV_FUNC_LOG_WARN(cuMemRelease(handle));
out:
    return status;
}

#endif /* HAVE_CUDA_FABRIC || HAVE_DECL_SYS_PIDFD_GETFD */
#if HAVE_CUDA_FABRIC
static ucs_status_t cuda_ipc_rem_mpool_cache_create(uct_cuda_ipc_rkey_t *key,
                                                    CUdevice cu_dev,
                                                    CUmemoryPool *mpool,
                                                    CUdeviceptr *mapped_addr)
{
    CUmemAccessDesc access_desc = {};
    CUdeviceptr dptr;
    ucs_status_t status;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuMemPoolImportFromShareableHandle(
                mpool, (void *)&key->ph.handle.fabric_handle,
                CU_MEM_HANDLE_TYPE_FABRIC, 0));
    if (status != UCS_OK) {
        goto err;
    }

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuMemPoolImportPointer(&dptr, *mpool,
                (CUmemPoolPtrExportData*)&key->ph.ptr));
    if (status != UCS_OK) {
        goto err_free_mpool;
    }

    uct_cuda_ipc_init_access_desc(&access_desc, cu_dev);

    status = UCT_CUDADRV_FUNC_LOG_ERR(
                cuMemPoolSetAccess(*mpool, &access_desc, 1));
    if (status != UCS_OK) {
        goto err_free_ptr;
    }

    *mapped_addr = dptr;
    return UCS_OK;

err_free_ptr:
    cuMemFree(dptr);
err_free_mpool:
    cuMemPoolDestroy(*mpool);
err:
    return status;
}

static ucs_status_t
uct_cuda_ipc_open_memhandle_mempool(uct_cuda_ipc_rkey_t *key, CUdevice cu_dev,
                                    CUdeviceptr *mapped_addr,
                                    ucs_log_level_t log_level)
{
    khash_t(cuda_ipc_rem_mpool_cache) *hash = &uct_cuda_ipc_rem_mpool_cache.hash;
    const CUmemFabricHandle *hkey           = &key->ph.handle.fabric_handle;
    ucs_status_t status                     = UCS_OK;
    CUmemoryPool mpool;
    khiter_t khiter;
    int khret;

    pthread_rwlock_rdlock(&uct_cuda_ipc_rem_mpool_cache.lock);
    khiter = kh_get(cuda_ipc_rem_mpool_cache, hash, *hkey);
    if (ucs_likely(khiter != kh_end(hash))) {
        key->ph.pool = kh_val(hash, khiter);
        ucs_assert(key->ph.pool != NULL);
        goto out_import_pointer;
    }

    /* acquire writer lock as we did not find cached mpool in hashmap */
    pthread_rwlock_unlock(&uct_cuda_ipc_rem_mpool_cache.lock);
    pthread_rwlock_wrlock(&uct_cuda_ipc_rem_mpool_cache.lock);

    khiter = kh_put(cuda_ipc_rem_mpool_cache, hash, *hkey, &khret);
    if ((khret == UCS_KH_PUT_BUCKET_EMPTY) ||
        (khret == UCS_KH_PUT_BUCKET_CLEAR)) {
        status = cuda_ipc_rem_mpool_cache_create(key, cu_dev, &mpool,
                                                 mapped_addr);
        if (status != UCS_OK) {
            /* Remove stale key from the hash if failed to construct new value.
             * Otherwise next lookup will return NULL value. */
            kh_del(cuda_ipc_rem_mpool_cache, hash, khiter);
            goto err;
        }

        kh_val(hash, khiter) = mpool;
        key->ph.pool         = mpool;
    } else if (khret == UCS_KH_PUT_KEY_PRESENT) {
        key->ph.pool = kh_val(hash, khiter);
    } else {
        ucs_error("unable to use cuda_ipc remote_cache hash: %d", khret);
        status = UCS_ERR_NO_RESOURCE;
        goto err;
    }

out_import_pointer:
    status = UCT_CUDADRV_FUNC(cuMemPoolImportPointer(mapped_addr, key->ph.pool,
            (CUmemPoolPtrExportData*)&key->ph.ptr), log_level);

err:
    pthread_rwlock_unlock(&uct_cuda_ipc_rem_mpool_cache.lock);
    return status;
}
#endif /* HAVE_CUDA_FABRIC */

#if HAVE_DECL_SYS_PIDFD_GETFD
static ucs_status_t
uct_cuda_ipc_open_memhandle_posix_fd(uct_cuda_ipc_extended_rkey_t *ext_key,
                                     CUdevice cu_dev, CUdeviceptr *mapped_addr,
                                     ucs_log_level_t log_level)
{
    uct_cuda_ipc_rkey_t *key = &ext_key->super;
    int pidfd, local_fd;
    ucs_status_t status;

    if (key->ph.handle.posix_fd.system_id != ucs_get_system_id()) {
        ucs_debug("posix_fd import: different machine "
                  "(remote system_id=0x%" PRIx64
                  " local system_id=0x%" PRIx64 ")",
                  key->ph.handle.posix_fd.system_id,
                  ucs_get_system_id());
        return UCS_ERR_UNREACHABLE;
    }

    if (ext_key->pid_ns != ucs_sys_get_ns(UCS_SYS_NS_TYPE_PID)) {
        ucs_debug("posix_fd import: peer pid=%d is in a different pid namespace "
                  "(remote=%u local=%u)", (int)key->pid, ext_key->pid_ns,
                  ucs_sys_get_ns(UCS_SYS_NS_TYPE_PID));
        return UCS_ERR_UNREACHABLE;
    }

    pidfd = syscall(SYS_pidfd_open, (int)key->pid, 0);
    if (pidfd < 0) {
        ucs_log(log_level, "pidfd_open(%d) failed: %m", (int)key->pid);
        return UCS_ERR_IO_ERROR;
    }

    local_fd = syscall(SYS_pidfd_getfd, pidfd, key->ph.handle.posix_fd.fd, 0);
    if (local_fd < 0) {
        ucs_log(log_level, "pidfd_getfd(pidfd=%d, pid=%d, fd=%d) failed: %m",
                pidfd, (int)key->pid, key->ph.handle.posix_fd.fd);
        status = UCS_ERR_IO_ERROR;
        goto close_pidfd;
    }

    status = uct_cuda_ipc_open_memhandle_vmm(
            key, cu_dev, mapped_addr, (void*)(uintptr_t)local_fd,
            CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, log_level);
    close(local_fd);
    ucs_trace("posix_fd import from pid=%d: %s", (int)key->pid,
              ucs_status_string(status));
close_pidfd:
    close(pidfd);
    return status;
}
#endif /* HAVE_DECL_SYS_PIDFD_GETFD */

static ucs_status_t
uct_cuda_ipc_open_memhandle(uct_cuda_ipc_extended_rkey_t *ext_key,
                            CUdevice cu_dev, CUdeviceptr *mapped_addr,
                            ucs_log_level_t log_level)
{
    uct_cuda_ipc_rkey_t *key = &ext_key->super;
    ucs_log_level_t level;

    ucs_trace("key handle type %u", key->ph.handle_type);

    switch(key->ph.handle_type) {
    case UCT_CUDA_IPC_KEY_HANDLE_TYPE_LEGACY:
        return uct_cuda_ipc_open_memhandle_legacy(key->ph.handle.legacy, cu_dev,
                                                  mapped_addr, log_level);
#if HAVE_CUDA_FABRIC
    case UCT_CUDA_IPC_KEY_HANDLE_TYPE_VMM:
        return uct_cuda_ipc_open_memhandle_vmm(
                key, cu_dev, mapped_addr, (void*)&key->ph.handle.fabric_handle,
                CU_MEM_HANDLE_TYPE_FABRIC, log_level);
    case UCT_CUDA_IPC_KEY_HANDLE_TYPE_MEMPOOL:
        return uct_cuda_ipc_open_memhandle_mempool(key, cu_dev, mapped_addr,
                                                   log_level);
#endif
#if HAVE_DECL_SYS_PIDFD_GETFD
    case UCT_CUDA_IPC_KEY_HANDLE_TYPE_POSIX_FD:
        return uct_cuda_ipc_open_memhandle_posix_fd(ext_key, cu_dev, mapped_addr,
                                                    log_level);
#endif
    case UCT_CUDA_IPC_KEY_HANDLE_TYPE_NO_IPC:
        level = UCS_LOG_LEVEL_DEBUG;
        break;
    default:
        level = UCS_LOG_LEVEL_ERROR;
        break;
    }

    ucs_log(level, "unsupported key handle type: %u", key->ph.handle_type);
    return UCS_ERR_INVALID_PARAM;
}

static void uct_cuda_ipc_cache_invalidate_regions(uct_cuda_ipc_cache_t *cache,
                                                  void *from, void *to)
{
    ucs_list_link_t region_list;
    ucs_status_t status;
    uct_cuda_ipc_cache_region_t *region, *tmp;

    ucs_list_head_init(&region_list);
    ucs_pgtable_search_range(&cache->pgtable, (ucs_pgt_addr_t)from,
                             (ucs_pgt_addr_t)to - 1,
                             uct_cuda_ipc_cache_region_collect_callback,
                             &region_list);
    ucs_list_for_each_safe(region, tmp, &region_list, list) {
        uct_cuda_ipc_cache_region_remove(cache, region);

        status = uct_cuda_ipc_close_memhandle(region);
        if (status != UCS_OK) {
            ucs_error("failed to close memhandle for base addr:%p type:%d (%s)",
                      (void*)region->key.d_bptr, region->key.ph.handle_type,
                      ucs_status_string(status));
        }

        ucs_free(region);
    }
    ucs_trace("%s: closed memhandles in the range [%p..%p]",
              cache->name, from, to);
}

static uct_cuda_ipc_cache_t *
uct_cuda_ipc_remote_cache_get(uct_cuda_ipc_cache_hash_key_t key)
{
    const khint_t it = kh_get(cuda_ipc_rem_cache,
                              &uct_cuda_ipc_remote_cache.hash, key);

    if (it != kh_end(&uct_cuda_ipc_remote_cache.hash)) {
        return kh_value(&uct_cuda_ipc_remote_cache.hash, it);
    }

    return NULL;
}

static ucs_status_t
uct_cuda_ipc_remote_cache_put(uct_cuda_ipc_cache_hash_key_t key,
                              uct_cuda_ipc_cache_t **cache_p)
{
    uct_cuda_ipc_cache_t *cache;
    int ret;
    khint_t it;
    char target_name[64];
    ucs_status_t status;

    cache = uct_cuda_ipc_remote_cache_get(key);
    if (cache != NULL) {
        goto out;
    }

    it = kh_put(cuda_ipc_rem_cache, &uct_cuda_ipc_remote_cache.hash, key, &ret);
    if (ret == UCS_KH_PUT_FAILED) {
        ucs_error("failed to allocate cuda_ipc remote_cache hash entry");
        return UCS_ERR_NO_MEMORY;
    }

    ucs_assertv_always(ret != UCS_KH_PUT_KEY_PRESENT, "key %d:%u:%d is present",
                       key.pid, key.pid_ns, key.cu_device);
    ucs_assertv_always((ret == UCS_KH_PUT_BUCKET_EMPTY) ||
                               (ret == UCS_KH_PUT_BUCKET_CLEAR),
                       "invalid return value: %d", ret);

    ucs_snprintf_safe(target_name, sizeof(target_name), "dest:%d:%u:%d",
                      key.pid, key.pid_ns, key.cu_device);
    status = uct_cuda_ipc_create_cache(&cache, target_name);
    if (status != UCS_OK) {
        kh_del(cuda_ipc_rem_cache, &uct_cuda_ipc_remote_cache.hash, it);
        ucs_error("failed to create cuda ipc cache: %s",
                  ucs_status_string(status));
        return status;
    }

    kh_val(&uct_cuda_ipc_remote_cache.hash, it) = cache;

out:
    *cache_p = cache;
    return UCS_OK;
}

static void
uct_cuda_ipc_cache_destroy_region(uct_cuda_ipc_cache_t *cache, uintptr_t d_bptr,
                                  const void *mapped_addr, int cache_enabled)
{
    ucs_pgt_region_t *pgt_region;
    uct_cuda_ipc_cache_region_t *region;

    /* use write lock because cache maybe modified */
    pthread_rwlock_wrlock(&cache->lock);
    pgt_region = UCS_PROFILE_CALL(ucs_pgtable_lookup, &cache->pgtable, d_bptr);
    ucs_assert(pgt_region != NULL);
    region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t);

    ucs_assert(region->refcount >= 1);
    region->refcount--;

    if (!region->refcount && !cache_enabled) {
        ucs_assert(region->mapped_addr == mapped_addr);
        uct_cuda_ipc_cache_region_destroy(cache, region);
    }

    pthread_rwlock_unlock(&cache->lock);
}

void uct_cuda_ipc_unmap_memhandle(pid_t pid, ucs_sys_ns_t pid_ns,
                                  uintptr_t d_bptr, const void *mapped_addr,
                                  CUdevice cu_dev, int cache_enabled)
{
    const uct_cuda_ipc_cache_hash_key_t key = {pid, pid_ns, cu_dev};
    uct_cuda_ipc_cache_t *cache;

    /* checking if the mapped address is the same as the d_bptr
     * this is true for the case of single process memory mapping
     * see uct_cuda_ipc_map_memhandle for more details */
    if ((d_bptr == (uintptr_t)mapped_addr) &&
        uct_cuda_ipc_is_rkey_local(pid, pid_ns)) {
        return;
    }

    ucs_rw_spinlock_read_lock(&uct_cuda_ipc_remote_cache.lock);
    cache = uct_cuda_ipc_remote_cache_get(key);
    if (cache != NULL) {
        uct_cuda_ipc_cache_destroy_region(cache, d_bptr, mapped_addr,
                                          cache_enabled);
    } else {
        ucs_debug("no remote cache found for key: %d:%u:%d", pid, pid_ns,
                  cu_dev);
    }

    ucs_rw_spinlock_read_unlock(&uct_cuda_ipc_remote_cache.lock);
}

static ucs_status_t
uct_cuda_ipc_cache_put_region(uct_cuda_ipc_cache_t *cache,
                              uct_cuda_ipc_extended_rkey_t *ext_key,
                              CUdevice cu_dev, void **mapped_addr,
                              ucs_log_level_t log_level)
{
    uct_cuda_ipc_rkey_t *key = &ext_key->super;
    ucs_pgt_region_t *pgt_region;
    uct_cuda_ipc_cache_region_t *region;
    ucs_status_t status;
    int ret;

    pthread_rwlock_wrlock(&cache->lock);
    pgt_region = UCS_PROFILE_CALL(ucs_pgtable_lookup,
                                  &cache->pgtable, key->d_bptr);
    if (ucs_likely(pgt_region != NULL)) {
        region = ucs_derived_of(pgt_region, uct_cuda_ipc_cache_region_t);

        if (key->ph.buffer_id == region->key.ph.buffer_id) {
            /*cache hit */
            ucs_trace("%s: cuda_ipc cache hit addr:%p size:%lu region:"
                      UCS_PGT_REGION_FMT, cache->name, (void *)key->d_bptr,
                      key->b_len, UCS_PGT_REGION_ARG(&region->super));

            /* Move to LRU tail (most recently used). Region is always on LRU
             * while alive, so unconditional del/add_tail is safe. */
            ucs_list_del(&region->lru_list);
            ucs_list_add_tail(&cache->lru_list, &region->lru_list);

            *mapped_addr = region->mapped_addr;
            ucs_assert(region->refcount < UINT64_MAX);
            region->refcount++;
            pthread_rwlock_unlock(&cache->lock);
            return UCS_OK;
        } else {
            ucs_trace("%s: cuda_ipc cache remove stale region:"
                      UCS_PGT_REGION_FMT " new_addr:%p new_size:%lu",
                      cache->name, UCS_PGT_REGION_ARG(&region->super),
                      (void *)key->d_bptr, key->b_len);

            uct_cuda_ipc_cache_region_destroy(cache, region);
        }
    }

    status = uct_cuda_ipc_open_memhandle(ext_key, cu_dev,
                                         (CUdeviceptr*)mapped_addr, log_level);
    if (ucs_unlikely(status != UCS_OK)) {
        if (ucs_likely(status == UCS_ERR_ALREADY_EXISTS)) {
            /* unmap all overlapping regions and retry*/
            uct_cuda_ipc_cache_invalidate_regions(cache, (void *)key->d_bptr,
                                                  UCS_PTR_BYTE_OFFSET(key->d_bptr,
                                                                      key->b_len));
            status = uct_cuda_ipc_open_memhandle(ext_key, cu_dev,
                                                 (CUdeviceptr*)mapped_addr,
                                                 log_level);
            if (ucs_unlikely(status != UCS_OK)) {
                if (ucs_likely(status == UCS_ERR_ALREADY_EXISTS)) {
                    /* unmap all cache entries and retry */
                    uct_cuda_ipc_cache_purge(cache);
                    status = uct_cuda_ipc_open_memhandle(
                            ext_key, cu_dev, (CUdeviceptr*)mapped_addr,
                            log_level);
                    if (status != UCS_OK) {
                        ucs_fatal("%s: failed to open ipc mem handle. addr:%p "
                                  "len:%lu (%s)", cache->name,
                                  (void *)key->d_bptr, key->b_len,
                                  ucs_status_string(status));
                    }
                } else {
                    ucs_fatal("%s: failed to open ipc mem handle. addr:%p len:%lu",
                              cache->name, (void *)key->d_bptr, key->b_len);
                }
            }
        } else {
            ucs_debug("%s: failed to open ipc mem handle. addr:%p len:%lu",
                      cache->name, (void *)key->d_bptr, key->b_len);
            goto err;
        }
    }

    /*create new cache entry */
    ret = ucs_posix_memalign((void **)&region,
                             ucs_max(sizeof(void *), UCS_PGT_ENTRY_MIN_ALIGN),
                             sizeof(uct_cuda_ipc_cache_region_t),
                             "uct_cuda_ipc_cache_region");
    if (ret != 0) {
        ucs_warn("failed to allocate uct_cuda_ipc_cache region");
        status = UCS_ERR_NO_MEMORY;
        goto err;
    }

    region->super.start = ucs_align_down_pow2((uintptr_t)key->d_bptr,
                                               UCS_PGT_ADDR_ALIGN);
    region->super.end   = ucs_align_up_pow2  ((uintptr_t)key->d_bptr + key->b_len,
                                               UCS_PGT_ADDR_ALIGN);
    region->key         = *key;
    region->mapped_addr = *mapped_addr;
    region->refcount    = 1;
    region->cu_dev      = cu_dev;

    status = UCS_PROFILE_CALL(ucs_pgtable_insert,
                              &cache->pgtable, &region->super);
    if (status == UCS_ERR_ALREADY_EXISTS) {
        /* overlapped region means memory freed at source. remove and try insert */
        uct_cuda_ipc_cache_invalidate_regions(cache,
                                              (void *)region->super.start,
                                              (void *)region->super.end);
        status = UCS_PROFILE_CALL(ucs_pgtable_insert,
                                  &cache->pgtable, &region->super);
    }
    if (status != UCS_OK) {

        ucs_error("%s: failed to insert region:"UCS_PGT_REGION_FMT" size:%lu :%s",
                  cache->name, UCS_PGT_REGION_ARG(&region->super), key->b_len,
                  ucs_status_string(status));
        ucs_free(region);
        goto err;
    }

    cache->num_regions++;
    cache->total_size += key->b_len;
    ucs_list_add_tail(&cache->lru_list, &region->lru_list);

    uct_cuda_ipc_cache_evict_lru(cache);

    ucs_trace("%s: cuda_ipc cache new region:"UCS_PGT_REGION_FMT" size:%lu"
              " (regions:%lu/%lu size:%zu/%zu)",
              cache->name, UCS_PGT_REGION_ARG(&region->super), key->b_len,
              cache->num_regions, uct_cuda_ipc_remote_cache.max_regions,
              cache->total_size, uct_cuda_ipc_remote_cache.max_size);

    status = UCS_OK;

err:
    pthread_rwlock_unlock(&cache->lock);
    return status;
}

UCS_PROFILE_FUNC(ucs_status_t, uct_cuda_ipc_map_memhandle,
                 (ext_key, cu_dev, mapped_addr, log_level),
                 uct_cuda_ipc_extended_rkey_t *ext_key, CUdevice cu_dev,
                 void **mapped_addr, ucs_log_level_t log_level)
{
    uct_cuda_ipc_rkey_t *key                     = &ext_key->super;
    const uct_cuda_ipc_cache_hash_key_t hash_key = {key->pid, ext_key->pid_ns,
                                                    cu_dev};
    uct_cuda_ipc_cache_t *cache;
    ucs_status_t status;
    CUuuid uuid;

    status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGetUuid(&uuid, cu_dev));
    if (status != UCS_OK) {
        return status;
    }

    if (uct_cuda_ipc_is_rkey_local(key->pid, ext_key->pid_ns) &&
        (memcmp(uuid.bytes, key->uuid.bytes, sizeof(uuid.bytes)) == 0)) {
        /* TODO: added for test purpose to enable cuda_ipc tests in gtest
         * mapped addrr is set to be same as d_bptr avoiding any calls to
         * uct_cuda_ipc_open_memhandle which would fail with invalid argument
         * error
         */
        *mapped_addr = (CUdeviceptr*)key->d_bptr;
        return UCS_OK;
    }

    ucs_rw_spinlock_write_lock(&uct_cuda_ipc_remote_cache.lock);
    status = uct_cuda_ipc_remote_cache_put(hash_key, &cache);
    if (status == UCS_OK) {
        status = uct_cuda_ipc_cache_put_region(cache, ext_key, cu_dev,
                                               mapped_addr, log_level);
    }

    ucs_rw_spinlock_write_unlock(&uct_cuda_ipc_remote_cache.lock);
    return status;
}

ucs_status_t uct_cuda_ipc_create_cache(uct_cuda_ipc_cache_t **cache,
                                       const char *name)
{
    ucs_status_t status;
    uct_cuda_ipc_cache_t *cache_desc;
    int ret;

    cache_desc = ucs_malloc(sizeof(uct_cuda_ipc_cache_t), "uct_cuda_ipc_cache_t");
    if (cache_desc == NULL) {
        ucs_error("failed to allocate memory for cuda_ipc cache");
        return UCS_ERR_NO_MEMORY;
    }

    ret = pthread_rwlock_init(&cache_desc->lock, NULL);
    if (ret) {
        ucs_error("pthread_rwlock_init() failed: %m");
        status = UCS_ERR_INVALID_PARAM;
        goto err;
    }

    status = ucs_pgtable_init(&cache_desc->pgtable,
                              uct_cuda_ipc_cache_pgt_dir_alloc,
                              uct_cuda_ipc_cache_pgt_dir_release);
    if (status != UCS_OK) {
        goto err_destroy_rwlock;
    }

    cache_desc->name = strdup(name);
    if (cache_desc->name == NULL) {
        status = UCS_ERR_NO_MEMORY;
        goto err_destroy_rwlock;
    }

    ucs_list_head_init(&cache_desc->lru_list);
    cache_desc->num_regions = 0;
    cache_desc->total_size  = 0;

    *cache = cache_desc;
    return UCS_OK;

err_destroy_rwlock:
    pthread_rwlock_destroy(&cache_desc->lock);
err:
    free(cache_desc);
    return status;
}

void uct_cuda_ipc_destroy_cache(uct_cuda_ipc_cache_t *cache)
{
    uct_cuda_ipc_cache_purge(cache);
    ucs_pgtable_cleanup(&cache->pgtable);
    pthread_rwlock_destroy(&cache->lock);
    free(cache->name);
    ucs_free(cache);
}

void uct_cuda_ipc_cache_set_global_limits(unsigned long max_regions,
                                          size_t max_size)
{
    uct_cuda_ipc_remote_cache.max_regions = ucs_min(uct_cuda_ipc_remote_cache.max_regions, 
                                                    max_regions);
    uct_cuda_ipc_remote_cache.max_size    = ucs_min(uct_cuda_ipc_remote_cache.max_size,
                                                    max_size);
}

void uct_cuda_ipc_destroy_cache_by_iface_address(
        const uct_cuda_ipc_iface_address_t *iface_address)
{
    uct_cuda_ipc_cache_hash_key_t cache_hash_key;
    int num_devices;
    ucs_status_t status;
    int device_index;
    CUdevice cuda_device;
    khint_t khiter;
    uct_cuda_ipc_cache_t *cache;

    cache_hash_key.pid    = iface_address->pid;
    cache_hash_key.pid_ns = iface_address->pid_ns;

    status = UCT_CUDADRV_FUNC_LOG_WARN(cuDeviceGetCount(&num_devices));
    if (status != UCS_OK) {
        return;
    }

    ucs_rw_spinlock_write_lock(&uct_cuda_ipc_remote_cache.lock);

    for (device_index = 0; device_index < num_devices; ++device_index) {
        status = UCT_CUDADRV_FUNC_LOG_WARN(
                cuDeviceGet(&cuda_device, device_index));
        if (status != UCS_OK) {
            continue;
        }

        cache_hash_key.cu_device = cuda_device;

        khiter = kh_get(cuda_ipc_rem_cache, &uct_cuda_ipc_remote_cache.hash,
                        cache_hash_key);
        if (khiter == kh_end(&uct_cuda_ipc_remote_cache.hash)) {
            continue;
        }

        cache = kh_val(&uct_cuda_ipc_remote_cache.hash, khiter);
        uct_cuda_ipc_destroy_cache(cache);
        kh_del(cuda_ipc_rem_cache, &uct_cuda_ipc_remote_cache.hash, khiter);
    }

    ucs_rw_spinlock_write_unlock(&uct_cuda_ipc_remote_cache.lock);
}

UCS_STATIC_INIT {
    ucs_rw_spinlock_init(&uct_cuda_ipc_remote_cache.lock);
    kh_init_inplace(cuda_ipc_rem_cache, &uct_cuda_ipc_remote_cache.hash);
    uct_cuda_ipc_remote_cache.max_regions = ULONG_MAX;
    uct_cuda_ipc_remote_cache.max_size    = SIZE_MAX;

#if HAVE_CUDA_FABRIC
    pthread_rwlock_init(&uct_cuda_ipc_rem_mpool_cache.lock, NULL);
    /* Assumption: If import process succeeds, then the two nodes are in the
     * same domain. Within a domain, fabric handles are expected to be unique.
     * For this reason, there is no need to maintain a hashmap per peer OS as
     * key collisions are not expected to occur. */
    kh_init_inplace(cuda_ipc_rem_mpool_cache, &uct_cuda_ipc_rem_mpool_cache.hash);
#endif
}

UCS_STATIC_CLEANUP {
    uct_cuda_ipc_cache_t *rem_cache;

#if HAVE_CUDA_FABRIC
    CUmemoryPool mpool;

    kh_foreach_value(&uct_cuda_ipc_rem_mpool_cache.hash, mpool, {
        cuMemPoolDestroy(mpool);
    });
    kh_destroy_inplace(cuda_ipc_rem_mpool_cache,
                       &uct_cuda_ipc_rem_mpool_cache.hash);
    pthread_rwlock_destroy(&uct_cuda_ipc_rem_mpool_cache.lock);
#endif

    kh_foreach_value(&uct_cuda_ipc_remote_cache.hash, rem_cache, {
        uct_cuda_ipc_destroy_cache(rem_cache);
    })
    kh_destroy_inplace(cuda_ipc_rem_cache, &uct_cuda_ipc_remote_cache.hash);
    ucs_rw_spinlock_cleanup(&uct_cuda_ipc_remote_cache.lock);
}
