diff --git a/csrc/inc/page_allocator.hpp b/csrc/inc/page_allocator.hpp index 6538bde5..43403d1b 100644 --- a/csrc/inc/page_allocator.hpp +++ b/csrc/inc/page_allocator.hpp @@ -155,6 +155,7 @@ class PageAllocator { BroadcastMapCallback broadcast_map_callback_; BroadcastUnmapCallback broadcast_unmap_callback_; ShouldUseWorkerIpcCallback should_use_worker_ipc_callback_; + mutable std::atomic should_use_worker_ipc_cached_{false}; }; } // namespace kvcached diff --git a/csrc/page_allocator.cpp b/csrc/page_allocator.cpp index f387abce..9c77ffe7 100644 --- a/csrc/page_allocator.cpp +++ b/csrc/page_allocator.cpp @@ -507,8 +507,11 @@ void PageAllocator::set_broadcast_unmap_callback( void PageAllocator::set_should_use_worker_ipc_callback( ShouldUseWorkerIpcCallback callback) { + bool use_worker_ipc = callback ? callback() : false; std::lock_guard lock(lock_); - should_use_worker_ipc_callback_ = callback; + should_use_worker_ipc_callback_ = std::move(callback); + should_use_worker_ipc_cached_.store(use_worker_ipc, + std::memory_order_release); LOGGER(INFO, "Should-use-worker-ipc callback set for PageAllocator"); } @@ -731,10 +734,26 @@ void PageAllocator::stop_prealloc_thread_internal() { } bool PageAllocator::should_use_worker_ipc() const { - if (should_use_worker_ipc_callback_) { - return should_use_worker_ipc_callback_(); + ShouldUseWorkerIpcCallback callback; + { + std::lock_guard lock(lock_); + if (prealloc_thread_ && + prealloc_thread_->get_id() == std::this_thread::get_id()) { + // The background prealloc thread must not re-enter Python here. It only + // consumes the cached decision, which non-prealloc callers refresh. + return should_use_worker_ipc_cached_.load(std::memory_order_acquire); + } + callback = should_use_worker_ipc_callback_; } - return false; + + if (callback) { + bool use_worker_ipc = callback(); + should_use_worker_ipc_cached_.store(use_worker_ipc, + std::memory_order_release); + return use_worker_ipc; + } + + return should_use_worker_ipc_cached_.load(std::memory_order_acquire); } } // namespace kvcached \ No newline at end of file