[Bugfix][Store] Recover Master NoF heartbeat connections after failures

mooncake-store/benchmarks/nof_worker_pool_bench.cpp

@@ -731,12 +731,13 @@ int main(int argc, char **argv) {
             return 1;
         }
 
+        mooncake::SpdkNofWorkerPool pool;
         std::vector<EndpointContext> endpoints;
         endpoints.reserve(endpoint_strings.size());
         std::unordered_set<mooncake::nof_seg_handle *> unique_handles;
         for (size_t i = 0; i < endpoint_strings.size(); ++i) {
             EndpointContext endpoint(endpoint_strings[i], FLAGS_seed + i);
-            endpoint.seg_handle = wrapper.OpenNofSegment(endpoint.endpoint);
+            endpoint.seg_handle = pool.OpenNofSegment(endpoint.endpoint);
             if (!endpoint.seg_handle) {
                 LOG(ERROR) << "Failed to open NoF endpoint: "
                            << endpoint.endpoint;
@@ -837,7 +838,6 @@ int main(int argc, char **argv) {
         }
 
         std::vector<EndpointStats> endpoint_stats(endpoints.size());
-        mooncake::SpdkNofWorkerPool pool;
 
         TimePoint start = Clock::now();
         TimePoint warmup_end = start + std::chrono::seconds(FLAGS_warmup_sec);

mooncake-store/include/spdk/spdk_wrapper.h

@@ -2,11 +2,15 @@
 #include <atomic>
 #include <cstddef>
 #include <cstdint>
+#include <condition_variable>
+#include <functional>
 #include <map>
 #include <memory>
 #include <mutex>
 #include <stack>
 #include <string>
+#include <thread>
+#include <utility>
 #include <vector>
 #include <spdk/env.h>
 #include <spdk/nvme.h>
@@ -22,6 +26,7 @@ constexpr int kSpdkNofOpNum = 2;
 struct nof_seg_handle;
 struct tr_info;
 struct ctrlr_info;
+class SpdkNofWorkerPool;
 
 class SpdkWrapper {
    public:
@@ -41,8 +46,10 @@ class SpdkWrapper {
     int64_t NvmePollProcessCompletion(nof_seg_handle *seg,
                                       uint32_t complete_per_seg);
 
-    /** @brief Open a NoF segment. */
-    nof_seg_handle *OpenNofSegment(const std::string &tr_str);
+    // When client_managed is true, only prepare the segment metadata;
+    // the owner worker creates the qpair.
+    nof_seg_handle *OpenNofSegment(const std::string &tr_str,
+                                   bool client_managed = false);
 
     uint32_t GetBlockSize(const nof_seg_handle *seg_handle);
 
@@ -53,6 +60,9 @@ class SpdkWrapper {
     bool ProbeNofSegment(const std::string &tr_str, uint32_t timeout_ms,
                          std::string *error_reason = nullptr);
 
+    void PollProbeControllers();
+    void ShutdownProbeConnections();
+
     /** @brief Register external memory with SPDK for NoF zero-copy transfers.
      *
      * Memory returned by Alloc() is already registered with SPDK; use this to
@@ -99,6 +109,30 @@ class SpdkWrapper {
      */
     static bool IsRegistrableRange(void *addr, size_t size);
 
+    // Non-owning identity; the pool unregisters its workers before destruction.
+    using WorkerId = std::pair<SpdkNofWorkerPool *, int>;
+    // Dispatch through a registered callback to keep SpdkWrapper usable by
+    // Master without linking the client-side SpdkNofWorkerPool implementation.
+    using RecoveryTaskPoster = void (*)(WorkerId, ctrlr_info *, bool resume);
+
+    // Register the client-side dispatcher before any worker is admitted.
+    // Repeated registrations must use the same function.
+    void RegisterRecoveryTaskPoster(RecoveryTaskPoster poster);
+    void StartControllerThread();
+
+    ctrlr_info *GetController(const nof_seg_handle *seg) const;
+    bool BindSegmentToWorker(nof_seg_handle *seg, WorkerId owner,
+                             const std::function<void()> &enqueue);
+
+    void RequestRecovery(ctrlr_info *ctrlr);
+    std::vector<nof_seg_handle *> OwnedSegments(ctrlr_info *ctrlr,
+                                                WorkerId owner);
+    void OwnerDrained(ctrlr_info *ctrlr, WorkerId owner);
+    void UnregisterOwner(ctrlr_info *ctrlr, WorkerId owner);
+
+    int SetupQpairIfNeeded(nof_seg_handle *seg);
+    bool DrainQpair(nof_seg_handle *seg);
+
    private:
     struct ProbeBuffer {
         void *ptr{nullptr};
@@ -141,10 +175,21 @@ class SpdkWrapper {
     void ReplenishProbeRequestContextPoolLocked(size_t count);
     static void ProbeReadComplete(void *ctx, const struct spdk_nvme_cpl *cpl);
 
+    void MarkProbeControllerDraining(ctrlr_info *info);
+
+    void ControllerLoop();
+
     std::atomic<bool> initialized{false};
     std::mutex init_mutex;
     std::map<std::string, std::unique_ptr<ctrlr_info>> connected_ctrlrs;
     std::mutex ctrlrs_mutex;
+
+    std::mutex controller_mutex_;
+    std::condition_variable controller_cv_;
+    bool controller_stop_{false};
+    std::thread controller_thread_;
+    RecoveryTaskPoster post_recovery_task_{nullptr};
+
     std::map<std::string, std::unique_ptr<ProbeBuffer>> probe_buffers_;
     std::mutex probe_buffers_mutex_;
     std::vector<std::unique_ptr<ProbeRequestContext>> probe_request_contexts_;

mooncake-store/include/transfer_task.h

@@ -4,6 +4,7 @@
 #include <condition_variable>
 #include <cstdint>
 #include <cstring>
+#include <functional>
 #include <map>
 #include <memory>
 #include <mutex>
@@ -406,6 +407,7 @@ class MemcpyWorkerPool {
 #ifdef USE_NOF
 // struct SpdkNofSubTask;
 struct SpdkNofQos;
+struct SpdkNofControllerState;
 
 /**
  * @brief Spdk nvmf operation descriptor
@@ -455,6 +457,8 @@ struct SpdkNofSubTask {
 constexpr int kDefaultSpdkNofSubmitChunkBytes = (1 << 17);    // 128k
 constexpr int kDefaultSpdkNofInflightBytesLimit = (1 << 25);  // 32M
 struct SpdkNofQos {
+    SpdkNofControllerState* controller_state{nullptr};
+
     int inflight_blocks[kSpdkNofOpNum];
     int blocks_per_chunk;
     int inflight_blocks_limit;
@@ -486,6 +490,13 @@ struct SpdkNofQos {
     }
 };
 
+struct SpdkNofControllerState {
+    enum class Phase { Running, Cleaning, Paused };
+
+    Phase phase{Phase::Running};
+    std::vector<nof_seg_handle*> draining_segments;
+};
+
 /**
  * @brief Thread pool for asynchronous spdk nvmf operations
  *
@@ -495,6 +506,13 @@ struct SpdkNofQos {
 constexpr int kDefaultSpdkNofWorkers = 4;
 class SpdkNofWorkerPool {
    public:
+    struct WorkerContext;
+    using ControlTask = std::function<void(WorkerContext&)>;
+
+    // Callbacks run on the target worker and must not wait for other tasks on
+    // that worker. Callers must stop posting before destroying the pool.
+    void PostControlTask(int worker_idx, ControlTask task);
+
     explicit SpdkNofWorkerPool(int numa_socket_id = 0);
     ~SpdkNofWorkerPool();
 
@@ -504,6 +522,10 @@ class SpdkNofWorkerPool {
     SpdkNofWorkerPool(SpdkNofWorkerPool&&) = delete;
     SpdkNofWorkerPool& operator=(SpdkNofWorkerPool&&) = delete;
 
+    // Bind the segment to a worker before returning. That worker creates the
+    // qpair and starts the asynchronous connection.
+    nof_seg_handle* OpenNofSegment(const std::string& tr_str);
+
     /**
      * @brief Submit a spdk nvmf task for async execution
      * @param task The spdk nvmf task to execute
@@ -512,11 +534,13 @@ class SpdkNofWorkerPool {
 
    private:
     void workerThread(int work_idx);
+    void PostRecoveryTask(int worker_idx, ctrlr_info* ctrlr, bool resume);
 
     int worker_count_;
     int numa_socket_id_;
     std::vector<std::thread> workers_;
     std::unique_ptr<std::queue<SpdkNofTask>[]> task_queue_;
+    std::unique_ptr<std::queue<ControlTask>[]> control_queue_;
     std::unique_ptr<std::mutex[]> queue_mutex_;
     std::unique_ptr<std::condition_variable[]> queue_cv_;
     std::atomic<bool> shutdown_;

mooncake-store/src/master_service.cpp

@@ -12411,9 +12411,11 @@ bool MasterService::TryUnmountNoFSegmentByHeartbeat(
     return true;
 }
 
+#ifdef USE_NOF
 void MasterService::NofHeartbeatThreadFunc() {
     size_t next_probe_index = 0;
     while (nof_heartbeat_running_) {
+        SpdkWrapper::GetInstance().PollProbeControllers();
         auto now = std::chrono::steady_clock::now();
         std::vector<MountedNoFSegmentSnapshot> mounted_segments;
         nof_segment_manager_.GetMountedSegmentsSnapshot(mounted_segments);
@@ -12561,7 +12563,9 @@ void MasterService::NofHeartbeatThreadFunc() {
             TryUnmountNoFSegmentByHeartbeat(*probe_target, error_reason);
         }
     }
+    SpdkWrapper::GetInstance().ShutdownProbeConnections();
 }
+#endif
 
 tl::expected<void, ErrorCode>
 MasterService::ValidateDynamicReplicaPendingForCopyStart(