From c12808c8b1e648c1814444a945246f47c2ffc9b2 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Tue, 2 Jul 2024 08:48:53 -0700
Subject: [PATCH] Fix segmentation fault in `TrackedRequests`

Move `TrackedRequests` mutexes to the class itself, allowing locking the
objects even after releasing ownership. With this change, it's now
possible to lock the object within `InflightRewquests::merge` which was
a source of segmentation faults.
---
 cpp/include/ucxx/inflight_requests.h |  6 +++---
 cpp/src/inflight_requests.cpp        | 23 +++++++++++++----------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/cpp/include/ucxx/inflight_requests.h b/cpp/include/ucxx/inflight_requests.h
index 8a63dfcb..d71caf88 100644
--- a/cpp/include/ucxx/inflight_requests.h
+++ b/cpp/include/ucxx/inflight_requests.h
@@ -39,6 +39,9 @@ typedef struct TrackedRequests {
     std::make_unique<InflightRequestsMap>()};  ///< Valid requests awaiting completion.
   InflightRequestsMapPtr _canceling{
     std::make_unique<InflightRequestsMap>()};  ///< Requests scheduled for cancelation.
+  std::mutex _mutex{};  ///< Mutex to control access to inflight requests container
+  std::mutex
+    _cancelMutex{};  ///< Mutex to allow cancelation and prevent removing requests simultaneously
 } TrackedRequests;
 
 /**
@@ -61,9 +64,6 @@ class InflightRequests {
     std::make_unique<TrackedRequests>()};  ///< Container storing pointers to all inflight
                                            ///< and in cancelation process requests known to
                                            ///< the owner of this object
-  std::mutex _mutex{};  ///< Mutex to control access to inflight requests container
-  std::mutex
-    _cancelMutex{};  ///< Mutex to allow cancelation and prevent removing requests simultaneously
 
   /**
    * @brief Drop references to requests that completed cancelation.
diff --git a/cpp/src/inflight_requests.cpp b/cpp/src/inflight_requests.cpp
index f08e7a8a..4aec90a5 100644
--- a/cpp/src/inflight_requests.cpp
+++ b/cpp/src/inflight_requests.cpp
@@ -14,13 +14,13 @@ InflightRequests::~InflightRequests() { cancelAll(); }
 
 size_t InflightRequests::size()
 {
-  std::lock_guard<std::mutex> lock(_mutex);
+  std::lock_guard<std::mutex> lock(_trackedRequests->_mutex);
   return _trackedRequests->_inflight->size();
 }
 
 void InflightRequests::insert(std::shared_ptr<Request> request)
 {
-  std::lock_guard<std::mutex> lock(_mutex);
+  std::lock_guard<std::mutex> lock(_trackedRequests->_mutex);
 
   _trackedRequests->_inflight->insert({request.get(), request});
 }
@@ -28,7 +28,10 @@ void InflightRequests::insert(std::shared_ptr<Request> request)
 void InflightRequests::merge(TrackedRequestsPtr trackedRequests)
 {
   {
-    std::scoped_lock lock{_cancelMutex, _mutex};
+    std::scoped_lock lock{_trackedRequests->_cancelMutex,
+                          _trackedRequests->_mutex,
+                          trackedRequests->_cancelMutex,
+                          trackedRequests->_mutex};
     if (trackedRequests->_inflight != nullptr)
       _trackedRequests->_inflight->merge(*(trackedRequests->_inflight));
     else
@@ -43,7 +46,7 @@ void InflightRequests::merge(TrackedRequestsPtr trackedRequests)
 void InflightRequests::remove(const Request* const request)
 {
   do {
-    int result = std::try_lock(_cancelMutex, _mutex);
+    int result = std::try_lock(_trackedRequests->_cancelMutex, _trackedRequests->_mutex);
 
     /**
      * `result` can be have one of three values:
@@ -72,8 +75,8 @@ void InflightRequests::remove(const Request* const request)
         tmpRequest = search->second;
         _trackedRequests->_inflight->erase(search);
       }
-      _cancelMutex.unlock();
-      _mutex.unlock();
+      _trackedRequests->_cancelMutex.unlock();
+      _trackedRequests->_mutex.unlock();
       return;
     }
   } while (true);
@@ -84,7 +87,7 @@ size_t InflightRequests::dropCanceled()
   size_t removed = 0;
 
   {
-    std::scoped_lock lock{_cancelMutex};
+    std::scoped_lock lock{_trackedRequests->_cancelMutex};
     for (auto it = _trackedRequests->_canceling->begin();
          it != _trackedRequests->_canceling->end();) {
       auto request = it->second;
@@ -105,7 +108,7 @@ size_t InflightRequests::getCancelingSize()
   dropCanceled();
   size_t cancelingSize = 0;
   {
-    std::scoped_lock lock{_cancelMutex};
+    std::scoped_lock lock{_trackedRequests->_cancelMutex};
     cancelingSize = _trackedRequests->_canceling->size();
   }
 
@@ -117,7 +120,7 @@ size_t InflightRequests::cancelAll()
   decltype(_trackedRequests->_inflight) toCancel;
   size_t total;
   {
-    std::scoped_lock lock{_cancelMutex, _mutex};
+    std::scoped_lock lock{_trackedRequests->_cancelMutex, _trackedRequests->_mutex};
     total = _trackedRequests->_inflight->size();
 
     // Fast path when no requests have been registered or the map has been
@@ -152,7 +155,7 @@ size_t InflightRequests::cancelAll()
 
 TrackedRequestsPtr InflightRequests::release()
 {
-  std::scoped_lock lock{_cancelMutex, _mutex};
+  std::scoped_lock lock{_trackedRequests->_cancelMutex, _trackedRequests->_mutex};
 
   return std::exchange(_trackedRequests, std::make_unique<TrackedRequests>());
 }