ucp_init() can deadlock on the GLIBC loader lock in cuInit() #12078

/*
 * ucx_cuinit_dlopen_deadlock.c
 *
 * Shows that ucp_init() and a concurrent first cuInit() on another thread
 * deadlock. Two locks are involved:
 *
 *   - ucp_init() loads the "cuda" uct module with dlopen(). That module's
 *     UCS_STATIC_INIT (src/uct/cuda/base/cuda_md.c) calls cuInit(0) from the
 *     library constructor, so cuInit runs while glibc's loader lock
 *     (_dl_load_lock) is held.
 *   - The process's first cuInit() takes libcuda's internal init lock and,
 *     while holding it, calls dlopen("libcuda.so.1"), which needs
 *     _dl_load_lock.
 *
 * A barrier releases both threads together. Each thread records when it
 * enters and leaves its call. At the deadline the main thread takes one
 * snapshot of both and reports one verdict:
 *
 *   OK        both calls returned and succeeded               exit 0
 *   FAIL      both calls returned, at least one failed        exit 1
 *   DEADLOCK  both threads still inside their calls           exit 2
 *   STUCK     exactly one thread still inside its call        exit 3
 *
 * For DEADLOCK and STUCK it prints, for each thread still inside its call,
 * which call, how long it has been inside, its wchan and its stack.
 *
 * Each run is one process, because only the first ucp_init() and the first
 * cuInit() in a process take these paths.
 *
 * Build:
 *   cc -std=gnu11 -O1 -g -o ucx_cuinit_dlopen_deadlock ucx_cuinit_dlopen_deadlock.c \
 *      -I$UCX/include -L$UCX/lib -Wl,-rpath,$UCX/lib -lucp -lucs -lpthread -ldl
 *
 * Usage: ucx_cuinit_dlopen_deadlock [timeout_s=10] [cuinit_delay_us=0]
 *
 * libcuda is not linked: it is dlopen()ed and cuInit is found with dlsym(),
 * as a CUDA runtime or a binding such as cudarc does, before the threads
 * start. Nothing calls cuInit before the barrier.
 */
#define _GNU_SOURCE
#include <dlfcn.h>
#include <execinfo.h>
#include <pthread.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/syscall.h>
#include <time.h>
#include <unistd.h>

#include <ucp/api/ucp.h>

typedef int (*cuInit_t)(unsigned int);

enum { NOT_STARTED, INSIDE, RETURNED };

#define MAX_FRAMES 48

struct worker {
    const char *call;        /* the call this thread makes */
    pthread_t   thread;
    atomic_int  tid;
    atomic_int  state;       /* NOT_STARTED, INSIDE or RETURNED */
    double      entered_ms;  /* written before state becomes INSIDE */
    double      returned_ms; /* written before state becomes RETURNED */
    int         result;      /* the call's return value, likewise */
    void       *frames[MAX_FRAMES];
    atomic_int  nframes;     /* -1 until the stack has been captured */
};

static struct worker     ucp_thread = { .call = "ucp_init" };
static struct worker     cu_thread  = { .call = "cuInit" };
static pthread_barrier_t barrier;
static cuInit_t          cu_init;
static ucp_config_t     *config;
static ucp_context_h     context;
static long              cuinit_delay_us;

static double now_ms(void)
{
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec * 1e3 + ts.tv_nsec * 1e-6;
}

/* Captures the interrupted thread's stack. backtrace() was called once in
 * main, so it loads nothing here, and its dl_iterate_phdr() takes
 * _dl_load_write_lock, not the _dl_load_lock this deadlock holds. */
static void capture_stack(int sig)
{
    pid_t          tid = (pid_t)syscall(SYS_gettid);
    struct worker *w   = (tid == atomic_load(&ucp_thread.tid)) ? &ucp_thread : &cu_thread;
    (void)sig;
    atomic_store(&w->nframes, backtrace(w->frames, MAX_FRAMES));
}

static void enter(struct worker *w)
{
    w->entered_ms = now_ms();
    atomic_store(&w->state, INSIDE);
}

static void leave(struct worker *w, int result)
{
    w->returned_ms = now_ms();
    w->result      = result;
    atomic_store(&w->state, RETURNED);
}

static void *ucp_main(void *arg)
{
    ucp_params_t params;
    ucs_status_t status;
    (void)arg;

    atomic_store(&ucp_thread.tid, (int)syscall(SYS_gettid));
    pthread_setname_np(pthread_self(), "ucp_init");
    memset(&params, 0, sizeof(params));
    params.field_mask = UCP_PARAM_FIELD_FEATURES;
    params.features   = UCP_FEATURE_TAG;

    pthread_barrier_wait(&barrier);
    enter(&ucp_thread);
    status = ucp_init(&params, config, &context);
    leave(&ucp_thread, (int)status);
    return NULL;
}

static void *cu_main(void *arg)
{
    int rc;
    (void)arg;

    atomic_store(&cu_thread.tid, (int)syscall(SYS_gettid));
    pthread_setname_np(pthread_self(), "cuInit");

    pthread_barrier_wait(&barrier);
    if (cuinit_delay_us > 0) {
        usleep(cuinit_delay_us);
    }
    enter(&cu_thread);
    rc = cu_init(0);
    leave(&cu_thread, rc);
    return NULL;
}

/* Prints addr as "<file>+0x<offset>" from /proc/self/maps, which needs no
 * loader lock (dladdr() takes the one that is held). */
static void print_frame(int i, void *addr)
{
    char      line[1024], path[768];
    uintptr_t a = (uintptr_t)addr, lo, hi, off;
    FILE     *maps = fopen("/proc/self/maps", "r");

    while (maps != NULL && fgets(line, sizeof(line), maps) != NULL) {
        path[0] = '\0';
        if (sscanf(line, "%lx-%lx %*s %lx %*s %*s %767s", &lo, &hi, &off, path) >= 3 &&
            a >= lo && a < hi) {
            printf("    #%-2d %p %s+0x%lx\n", i, addr, path[0] ? path : "[anon]",
                   (unsigned long)(a - lo + off));
            fclose(maps);
            return;
        }
    }
    if (maps != NULL) {
        fclose(maps);
    }
    printf("    #%-2d %p\n", i, addr);
}

/* Reports one thread as it was in the deadline snapshot. */
static void report_thread(struct worker *w, int state, double deadline_ms)
{
    char  path[64], wchan[64] = "?";
    FILE *f;
    int   tid = atomic_load(&w->tid), i, n;

    if (state == RETURNED) {
        printf("  %-8s (tid %d): returned %d after %.1f ms\n", w->call, tid, w->result,
               w->returned_ms - w->entered_ms);
        return;
    }
    if (state == NOT_STARTED) {
        printf("  %-8s (tid %d): never entered\n", w->call, tid);
        return;
    }
    snprintf(path, sizeof(path), "/proc/self/task/%d/wchan", tid);
    if ((f = fopen(path, "r")) != NULL) {
        if (fgets(wchan, sizeof(wchan), f) == NULL) {
            strcpy(wchan, "?");
        }
        fclose(f);
    }
    printf("  %-8s (tid %d): inside %s() for %.1f s at the deadline, wchan=%s\n", w->call,
           tid, w->call, (deadline_ms - w->entered_ms) / 1e3, wchan);
    atomic_store(&w->nframes, -1);
    pthread_kill(w->thread, SIGUSR1);
    for (i = 0; i < 100 && atomic_load(&w->nframes) < 0; ++i) {
        usleep(10000);
    }
    n = atomic_load(&w->nframes);
    if (n < 0) {
        printf("    (stack not captured)\n");
    }
    for (i = 0; i < n; ++i) {
        print_frame(i, w->frames[i]);
    }
}

int main(int argc, char **argv)
{
    double           timeout_s = (argc > 1) ? atof(argv[1]) : 10.0;
    struct sigaction sa;
    void            *libcuda, *warm[1];
    double           deadline_ms;
    int              ucp_state, cu_state, inside;
    ucs_status_t     status;

    cuinit_delay_us = (argc > 2) ? atol(argv[2]) : 0;
    setvbuf(stdout, NULL, _IOLBF, 0);

    libcuda = dlopen("libcuda.so.1", RTLD_NOW | RTLD_GLOBAL);
    if (libcuda == NULL || (cu_init = (cuInit_t)dlsym(libcuda, "cuInit")) == NULL) {
        fprintf(stderr, "cannot load cuInit from libcuda.so.1: %s\n", dlerror());
        return 1;
    }
    status = ucp_config_read(NULL, NULL, &config);
    if (status != UCS_OK) {
        fprintf(stderr, "ucp_config_read failed: %s\n", ucs_status_string(status));
        return 1;
    }
    printf("UCX %s, UCX_MODULES=%s, timeout %.0f s, cuInit delay %ld us\n",
           ucp_get_version_string(),
           getenv("UCX_MODULES") ? getenv("UCX_MODULES") : "<unset>", timeout_s,
           cuinit_delay_us);

    backtrace(warm, 1); /* loads libgcc_s now, not inside the signal handler */
    memset(&sa, 0, sizeof(sa));
    sa.sa_handler = capture_stack;
    sigaction(SIGUSR1, &sa, NULL);

    atomic_store(&ucp_thread.nframes, -1);
    atomic_store(&cu_thread.nframes, -1);
    pthread_barrier_init(&barrier, NULL, 3);
    pthread_create(&ucp_thread.thread, NULL, ucp_main, NULL);
    pthread_create(&cu_thread.thread, NULL, cu_main, NULL);
    pthread_barrier_wait(&barrier);

    deadline_ms = now_ms() + timeout_s * 1e3;
    while (now_ms() < deadline_ms) {
        if (atomic_load(&ucp_thread.state) == RETURNED &&
            atomic_load(&cu_thread.state) == RETURNED) {
            break;
        }
        usleep(10000);
    }

    /* One snapshot decides the verdict, and the report uses the same one. */
    ucp_state = atomic_load(&ucp_thread.state);
    cu_state  = atomic_load(&cu_thread.state);

    if (ucp_state == RETURNED && cu_state == RETURNED) {
        pthread_join(ucp_thread.thread, NULL);
        pthread_join(cu_thread.thread, NULL);
        if (ucp_thread.result != UCS_OK || cu_thread.result != 0) {
            printf("FAIL: ucp_init returned %s, cuInit returned CUresult %d\n",
                   ucs_status_string((ucs_status_t)ucp_thread.result), cu_thread.result);
            return 1;
        }
        printf("OK: ucp_init %.1f ms, cuInit %.1f ms\n",
               ucp_thread.returned_ms - ucp_thread.entered_ms,
               cu_thread.returned_ms - cu_thread.entered_ms);
        ucp_cleanup(context);
        ucp_config_release(config);
        return 0;
    }

    inside = (ucp_state == INSIDE) + (cu_state == INSIDE);
    if (inside == 2) {
        printf("DEADLOCK: ucp_init and cuInit both still inside their calls after %.0f s\n",
               timeout_s);
    } else {
        printf("STUCK: %d of the two calls still inside after %.0f s\n", inside, timeout_s);
    }
    report_thread(&ucp_thread, ucp_state, deadline_ms);
    report_thread(&cu_thread, cu_state, deadline_ms);
    _exit(inside == 2 ? 2 : 3); /* exit() would run destructors that may need the held locks */
}