mpool and thread safety

[amirgon]

Hi,
I would like to bring up an issue we are seeing with UCX when running Apache ArrowFlight over UCX.
I'm not sure if this is a bug in UCX, in ArrowFlight or in the way we use ArrowFlight with UCX, so I'm opening this for discussion.

The problem is a bit hard to reproduce, it only happens occasionally with very large amounts of data transferred over ArrowFlight with UCX/Posix-shmem transport.

When the problem happens, communication freezes and UCX shows these errors:

mpool.c:54   UCX  WARN  object 0x7efdabbb6f10 was not returned to mpool mm_recv_desc

and also:

mm_posix.c:218  UCX  ERROR open(file_name=(null)/ucx_shm_posix_7f46580058c0 flags=0x0) failed: No such file or directory

The last one is unusual since we use UCX_POSIX_USE_PROC_LINK=n and in such case uct_posix_shm_open should be called and not uct_posix_file_open.
(null) appears in the file name since posix_config->dir was null, as expected when using uct_posix_shm_open, however UCT_POSIX_SEG_FLAG_SHM_OPEN was unexpectedly read as 0 on that specific segment although it should have been 1 for all segments.

Digging into this further, we found out that the issue was related to memory pool corruption, specifically mm_recv_desc pool.
The corruption caused elements from the allocated list to be linked to the free list, which caused eventually the errors above.

The reason mm_recv_desc got corrupted was that it was used from multiple threads.
To support zero-copy, the received buffers were released only when they reached their final destination on a different thread, so the allocating thread and the releasing thread are different, while UCX mpool is not thread safe.

To fix that we added a spinlock on ucx mpool functions ucs_mpool_get_inline and ucs_mpool_add_to_freelist, and this seems to resolve the issue without impacting performance. (If this is a valid solution, I can create a PR)

After this fix the issue doesn't block us any more, however, there are still open questions:

• Has anyone else ever seen this issue? I couldn't find any references for that.
• How is thread safety handled in the context of memory allocation and deallocation in UCX?
• Are there assumptions on how the user is supposed to manage ucx buffers? Are there limitations as to which thread is allowed to release a buffer?

[evgeny-leksikov]

Hi @amirgon ,

does ArrowFlight uses UCT or UCP API?
UCT API is not thread safe, thread safety is guaranteed by only by UCP API and only if it's enabled in compile and run time.
If UCP API is used then could you please check:

1. ucx_info -b | grep ENABLE_MT
2. parameters of ucp_init and ucp_worker_create

[amirgon]

Hi @evgeny-leksikov

Thank you for your answer!
Indeed I can see that ArrowFlight is setting UCS_THREAD_MODE_SERIALIZED and not UCS_THREAD_MODE_MULTI when calling ucp_worker_create.

arrow_src/cpp/src/arrow/flight/transport/ucx/ucx_client.cc:

      ucs_status_t status;
      ucp_worker_params_t worker_params;
      std::memset(&worker_params, 0, sizeof(worker_params));
      worker_params.field_mask = UCP_WORKER_PARAM_FIELD_THREAD_MODE;
      worker_params.thread_mode = UCS_THREAD_MODE_SERIALIZED;

      ucp_worker_h ucp_worker;
      status = ucp_worker_create(ucp_context->get(), &worker_params, &ucp_worker);

I've opened an Arrow PR to fix that: apache/arrow#43120

[evgeny-leksikov]

does it fix original issue?

[amirgon]

I think it does.
I removed the spinlocks from UCX and replaced UCS_THREAD_MODE_SERIALIZED with UCS_THREAD_MODE_MULTI on ArrowFlight.
So far, after these changes, I wasn't able to reproduce the original issue.