fix(p2p): use pthread_mutex for MR lifecycle serialization

Author: MC952-arch

flagcx/core/flagcx_mr_registry.cc

@@ -174,6 +174,12 @@ static void freeEntryExtensions(struct flagcxMrEntry *entry) {
  */
 static flagcxResult_t idIndexAppend(struct flagcxMrRegistry *reg, uint64_t mrId,
                                     uintptr_t baseAddr) {
+  if (reg->idCount > 0 && mrId <= reg->idIndex[reg->idCount - 1].mrId) {
+    WARN("flagcxMrRegistry: non-monotonic mrId %lu (last %lu)",
+         (unsigned long)mrId,
+         (unsigned long)reg->idIndex[reg->idCount - 1].mrId);
+    return flagcxInternalError;
+  }
   if (reg->idCount >= reg->idCapacity) {
     int newCap =
         reg->idCapacity == 0 ? ID_INDEX_INITIAL_CAPACITY : reg->idCapacity * 2;
@@ -373,6 +379,8 @@ flagcxResult_t flagcxMrRegistryRegister(struct flagcxMrRegistry *reg,
         if (existing->p2p) {
           if (existing->p2p->mrId == 0)
             existing->p2p->mrId = reg->nextId++;
+          if (existing->p2p->mrId >= reg->nextId)
+            reg->nextId = existing->p2p->mrId + 1;
           if (idIndexAppend(reg, existing->p2p->mrId, addr) != flagcxSuccess) {
             /* Roll back: remove P2P ownership, caller retains ext */
             existing->p2p = NULL;
@@ -469,6 +477,8 @@ flagcxResult_t flagcxMrRegistryRegister(struct flagcxMrRegistry *reg,
         /* Assign mrId from registry's monotonic counter if not pre-set */
         if (entry->p2p->mrId == 0)
           entry->p2p->mrId = reg->nextId++;
+        if (entry->p2p->mrId >= reg->nextId)
+          reg->nextId = entry->p2p->mrId + 1;
         if (idIndexAppend(reg, entry->p2p->mrId, addr) != flagcxSuccess) {
           /* Roll back: remove the entry we just inserted */
           entry->p2p = NULL; /* caller retains ext ownership */

flagcx/core/flagcx_p2p.cc

@@ -372,6 +372,8 @@ static uint64_t &nextXferId() {
 #define gXferMutex xferMutex()
 #define gNextXferId nextXferId()
 
+static pthread_mutex_t gMrLifecycleMutex = PTHREAD_MUTEX_INITIALIZER;
+
 struct FlagcxSliceCache {
   static constexpr size_t kCap = 4096;
   std::vector<FlagcxSlice *> ring;
@@ -2098,15 +2100,17 @@ void flagcxP2pEngineDestroy(FlagcxP2pEngine *engine) {
   }
 
   {
-    /* Phase 1: collect P2P mhandle info under write lock */
+    pthread_mutex_lock(&gMrLifecycleMutex);
+
+    /* Phase 1: collect P2P mhandle info under read lock */
     struct P2pDeregInfo {
       int ibDevN;
       void *mhandle;
       uintptr_t baseAddr;
     };
     std::vector<P2pDeregInfo> deregList;
 
-    flagcxMrRegistryWrLock(flagcxGlobalMrRegistry);
+    flagcxMrRegistryRdLock(flagcxGlobalMrRegistry);
     int count = flagcxMrRegistryCount(flagcxGlobalMrRegistry);
     struct flagcxMrEntry *entries =
         flagcxMrRegistryEntries(flagcxGlobalMrRegistry);
@@ -2119,7 +2123,7 @@ void flagcxP2pEngineDestroy(FlagcxP2pEngine *engine) {
       info.baseAddr = entries[i].baseAddr;
       deregList.push_back(info);
     }
-    flagcxMrRegistryWrUnlock(flagcxGlobalMrRegistry);
+    flagcxMrRegistryRdUnlock(flagcxGlobalMrRegistry);
 
     /* Phase 2: deregister from registry first (atomically removes from lookup)
      */
@@ -2135,6 +2139,8 @@ void flagcxP2pEngineDestroy(FlagcxP2pEngine *engine) {
       } devCtx = {info.ibDevN};
       engine->adaptor->deregMr(&devCtx, info.mhandle);
     }
+
+    pthread_mutex_unlock(&gMrLifecycleMutex);
   }
 
   /* Release P2P engine's refcount on the global MR registry */
@@ -2477,6 +2483,8 @@ int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
   if (engine == NULL || data == 0)
     return -1;
 
+  pthread_mutex_lock(&gMrLifecycleMutex);
+
   /* Check for existing exact-match registration (dedup) */
   {
     struct flagcxMrEntry existing;
@@ -2490,9 +2498,11 @@ int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
           WARN("P2P Reg: addr 0x%lx size mismatch: existing %zu vs requested "
                "%zu",
                (unsigned long)data, existing.size, size);
+          pthread_mutex_unlock(&gMrLifecycleMutex);
           return -1;
         }
         mrId = p2pExt.p2p.mrId;
+        pthread_mutex_unlock(&gMrLifecycleMutex);
         return 0;
       }
     }
@@ -2520,6 +2530,7 @@ int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
                              ptrType, FLAGCX_NET_MR_FLAG_NONE,
                              &mhandle) != flagcxSuccess ||
       mhandle == NULL) {
+    pthread_mutex_unlock(&gMrLifecycleMutex);
     return -1;
   }
 
@@ -2528,6 +2539,7 @@ int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
       (struct flagcxMrP2pExt *)calloc(1, sizeof(struct flagcxMrP2pExt));
   if (p2pExt == NULL) {
     engine->adaptor->deregMr(&devCtx, mhandle);
+    pthread_mutex_unlock(&gMrLifecycleMutex);
     return -1;
   }
   /* mrId=0 signals registry to assign from its monotonic counter */
@@ -2544,34 +2556,48 @@ int flagcxP2pEngineReg(FlagcxP2pEngine *engine, uintptr_t data, size_t size,
   if (res != flagcxSuccess) {
     engine->adaptor->deregMr(&devCtx, mhandle);
     free(p2pExt);
+    pthread_mutex_unlock(&gMrLifecycleMutex);
     return -1;
   }
 
   mrId = assignedId;
+  pthread_mutex_unlock(&gMrLifecycleMutex);
   return 0;
 }
 
 void flagcxP2pEngineMrDestroy(FlagcxP2pEngine *engine, FlagcxP2pMr mr) {
   if (engine == NULL)
     return;
 
+  pthread_mutex_lock(&gMrLifecycleMutex);
+
   /* Find entry by mrId to get baseAddr for deregister */
   struct flagcxMrEntry mrEntry;
   if (flagcxMrRegistryLookupById(flagcxGlobalMrRegistry, mr, &mrEntry, NULL) !=
-      flagcxSuccess)
+      flagcxSuccess) {
+    pthread_mutex_unlock(&gMrLifecycleMutex);
     return;
+  }
 
   /* Remove from registry first — prevents concurrent readers from finding it */
   void *removedExt = NULL;
-  flagcxMrRegistryDeregister(flagcxGlobalMrRegistry, mrEntry.baseAddr,
-                             FLAGCX_MR_OWNER_P2P, NULL, &removedExt);
-  free(removedExt);
-
-  /* Now safe to deregister the adaptor handle */
+  flagcxResult_t res;
   struct {
     int ibDevN;
   } devCtx = {mrEntry.ibDevN};
+  FLAGCXCHECKGOTO(
+      flagcxMrRegistryDeregister(flagcxGlobalMrRegistry, mrEntry.baseAddr,
+                                 FLAGCX_MR_OWNER_P2P, NULL, &removedExt),
+      res, fail);
+  free(removedExt);
+
+  /* Now safe to deregister the adaptor handle */
   engine->adaptor->deregMr(&devCtx, mrEntry.mhandles[FLAGCX_MR_OWNER_IDX_P2P]);
+  pthread_mutex_unlock(&gMrLifecycleMutex);
+  return;
+
+fail:
+  pthread_mutex_unlock(&gMrLifecycleMutex);
 }
 
 int flagcxP2pEnginePrepareDesc(FlagcxP2pEngine *engine, FlagcxP2pMr mr,
@@ -2616,6 +2642,23 @@ int flagcxP2pEnginePrepareDesc(FlagcxP2pEngine *engine, FlagcxP2pMr mr,
     return -1;
   }
 
+  /* Verify (data, size) falls within registered MR region (overflow-safe) */
+  uintptr_t dataAddr = (uintptr_t)data;
+  if (dataAddr < entries[idx].baseAddr) {
+    flagcxMrRegistryWrUnlock(flagcxGlobalMrRegistry);
+    return -1;
+  }
+  size_t offset = (size_t)(dataAddr - entries[idx].baseAddr);
+  if (offset > entries[idx].size || size > entries[idx].size - offset) {
+    flagcxMrRegistryWrUnlock(flagcxGlobalMrRegistry);
+    return -1;
+  }
+
+  if (size > UINT32_MAX) {
+    flagcxMrRegistryWrUnlock(flagcxGlobalMrRegistry);
+    return -1;
+  }
+
   FlagcxP2pRdmaDesc desc;
   memset(&desc, 0, sizeof(desc));
   desc.addr = (uint64_t)(uintptr_t)data;

flagcx/core/include/flagcx_mr_registry.h

@@ -105,7 +105,7 @@ struct flagcxMrRegistry {
   int capacity;
   uint64_t nextId; /* monotonic ID generator for all subsystems */
 
-  struct flagcxMrIdEntry *idIndex; /* sorted by mrId (ascending, append-only) */
+  struct flagcxMrIdEntry *idIndex; /* sorted by mrId (ascending) */
   int idCount;
   int idCapacity;
 
@@ -130,9 +130,10 @@ flagcxResult_t flagcxMrRegistryDestroy(struct flagcxMrRegistry *reg);
  * ownerBit: one of FLAGCX_MR_OWNER_{P2P,COLL,RMA}
  * mhandle:  adaptor handle (stored in mhandles[ownerIdx])
  * ext:      subsystem extension struct pointer (ownership transferred to entry)
- * outId:    if non-NULL, returns a monotonic ID. For P2P, this is persisted
- *           in p2p->mrId and stable across repeated calls. For COLL/RMA,
- *           this is a one-shot assignment (not stored on the entry).
+ * outId:    if non-NULL, receives a monotonic ID on first registration.
+ *           For P2P: persisted in p2p->mrId, stable across repeated calls.
+ *           For COLL/RMA: one-shot assignment (not stored on entry);
+ *           re-registration of an already-owned entry sets *outId = 0.
  */
 flagcxResult_t flagcxMrRegistryRegister(struct flagcxMrRegistry *reg,
                                         uintptr_t addr, size_t size, int ibDevN,
@@ -218,6 +219,16 @@ struct flagcxMrEntry *flagcxMrRegistryEntries(struct flagcxMrRegistry *reg);
 extern struct flagcxMrRegistry *flagcxGlobalMrRegistry;
 
 flagcxResult_t flagcxMrRegistryGlobalInit(void);
+
+/*
+ * Release one reference to the global registry. Destroys when refcount
+ * reaches 0.
+ *
+ * Precondition: all data-path operations (Lookup, LookupById, PrepareDesc)
+ * must be quiesced before the final release. The caller is responsible for
+ * ensuring no concurrent registry access is in flight when the last
+ * reference is dropped. Violating this is undefined behavior (use-after-free).
+ */
 flagcxResult_t flagcxMrRegistryGlobalRelease(void);
 
 #ifdef __cplusplus

test/unittest/core/test_mr_registry.cpp

@@ -525,9 +525,8 @@ TEST_F(MrRegistryTest, ConcurrentWriters) {
   std::atomic<int> errors{0};
   auto worker = [&](int tid) {
     for (int i = 0; i < kOpsPerThread; i++) {
-      // Each thread uses a unique address range: tid * large_gap + i * page
-      uintptr_t addr =
-          (uintptr_t)(tid + 1) * 0x100000000ULL + (uintptr_t)i * kPageSize;
+      // Each thread uses a unique address range: tid * (1<<20) + i * page
+      uintptr_t addr = ((uintptr_t)(tid + 1) << 20) + (uintptr_t)i * kPageSize;
       size_t size = kPageSize;
 
       struct flagcxMrP2pExt *ext =