fix(rma): add peer bounds checks and harden bootstrap/test robustness

Author: MC952-arch

flagcx/core/flagcx_hetero.cc

@@ -14,6 +14,7 @@
 #include <sched.h>
 #include <stdlib.h>
 #include <string.h>
+#include <time.h>
 #include <unistd.h>
 
 #ifndef FLAGCX_RMA_QUEUE_SIZE
@@ -51,17 +52,29 @@ struct flagcxRmaDoneWaitCtx {
   volatile uint64_t *doneSeqsCpu; // pointer to proxy's doneSeqsCpu[peer]
   uint64_t opSeq;                 // sequence to wait for
   volatile int *rmaError;         // proxy error flag
+  pthread_mutex_t *doneMutex;     // proxy done condvar mutex
+  pthread_cond_t *doneCond;       // proxy done condvar
 };
 
 // Host-func callback: blocks stream until proxy signals completion.
+// Uses condvar to sleep instead of spinning — zero CPU while waiting.
 static void flagcxRmaDoneWaitHostFunc(void *arg) {
   struct flagcxRmaDoneWaitCtx *ctx = (struct flagcxRmaDoneWaitCtx *)arg;
+  pthread_mutex_lock(ctx->doneMutex);
   while (__atomic_load_n(ctx->doneSeqsCpu, __ATOMIC_ACQUIRE) < ctx->opSeq) {
     if (__atomic_load_n(ctx->rmaError, __ATOMIC_ACQUIRE)) {
       break;
     }
-    sched_yield();
+    struct timespec ts;
+    clock_gettime(CLOCK_REALTIME, &ts);
+    ts.tv_nsec += 1000000; // 1ms timeout as safety net
+    if (ts.tv_nsec >= 1000000000) {
+      ts.tv_sec++;
+      ts.tv_nsec -= 1000000000;
+    }
+    pthread_cond_timedwait(ctx->doneCond, ctx->doneMutex, &ts);
   }
+  pthread_mutex_unlock(ctx->doneMutex);
   free(ctx);
 }
 
@@ -96,14 +109,16 @@ static flagcxResult_t flagcxRmaWaitDone(struct flagcxRmaProxyState *proxy,
     return deviceAdaptor->streamWaitValue64(stream, &proxy->doneSeqsDev[peer],
                                             opSeq, 0);
   } else {
-    // HOST_FUNC: launch callback that polls doneSeqsCpu until done
+    // HOST_FUNC: launch callback that waits on doneCond until done
     struct flagcxRmaDoneWaitCtx *ctx =
         (struct flagcxRmaDoneWaitCtx *)malloc(sizeof(*ctx));
     if (ctx == NULL)
       return flagcxSystemError;
     ctx->doneSeqsCpu = &proxy->doneSeqsCpu[peer];
     ctx->opSeq = opSeq;
     ctx->rmaError = &proxy->rmaError;
+    ctx->doneMutex = &proxy->doneMutex;
+    ctx->doneCond = &proxy->doneCond;
     return deviceAdaptor->launchHostFunc(stream, flagcxRmaDoneWaitHostFunc,
                                          ctx);
   }
@@ -313,6 +328,12 @@ flagcxRmaProxyPollNonPersistCompletion(struct flagcxRmaProxyState *proxy,
         __atomic_store_n(&proxy->doneSeqsCpu[peer], desc->opSeq,
                          __ATOMIC_RELEASE);
       __atomic_fetch_add(&proxy->completionCount, 1ULL, __ATOMIC_RELEASE);
+      // Wake HOST_FUNC waiters sleeping on doneCond
+      if (!proxy->useStreamOps) {
+        pthread_mutex_lock(&proxy->doneMutex);
+        pthread_cond_broadcast(&proxy->doneCond);
+        pthread_mutex_unlock(&proxy->doneMutex);
+      }
     }
     free(desc);
     did = true;
@@ -570,6 +591,9 @@ flagcxResult_t flagcxHeteroRmaProxyStart(flagcxHeteroComm_t comm) {
     flagcxIntruQueueConstruct(&proxy->inProgressQueues[p]);
   }
 
+  pthread_mutex_init(&proxy->doneMutex, NULL);
+  pthread_cond_init(&proxy->doneCond, NULL);
+
   // Allocate device memory for stream-based synchronization.
   proxy->doneSeqsDev = NULL;
   proxy->doneSeqsCpu = NULL;
@@ -659,6 +683,8 @@ flagcxResult_t flagcxHeteroRmaProxyStart(flagcxHeteroComm_t comm) {
       deviceAdaptor->gdrMemFree(proxy->readySeqsDev, NULL);
     for (int p = 0; p < nRanks; p++)
       pthread_mutex_destroy(&proxy->peerProducerMutexes[p]);
+    pthread_cond_destroy(&proxy->doneCond);
+    pthread_mutex_destroy(&proxy->doneMutex);
     free(proxy->circularBuffers);
     free((void *)proxy->pis);
     free((void *)proxy->cis);
@@ -691,6 +717,9 @@ flagcxResult_t flagcxHeteroRmaProxyStop(flagcxHeteroComm_t comm) {
   for (int p = 0; p < proxy->nRanks; p++)
     pthread_mutex_destroy(&proxy->peerProducerMutexes[p]);
 
+  pthread_cond_destroy(&proxy->doneCond);
+  pthread_mutex_destroy(&proxy->doneMutex);
+
   // Free CPU mappings / host memory
   if (proxy->useStreamOps && deviceAdaptor->gdrPtrMunmap != NULL) {
     if (proxy->doneSeqsCpu != NULL)

flagcx/core/include/flagcx_hetero.h

@@ -97,6 +97,11 @@ struct flagcxRmaProxyState {
   bool ipcInitFailed; // true if IpcInit was attempted and failed (prevents
                       // retries)
 
+  // Condition variable for HOST_FUNC done-wait: proxy thread broadcasts
+  // after updating doneSeqsCpu so host-func callbacks wake without spinning.
+  pthread_mutex_t doneMutex;
+  pthread_cond_t doneCond;
+
   pthread_t thread;
   volatile int stop;
 };

flagcx/flagcx.cc

@@ -2592,6 +2592,11 @@ flagcxResult_t flagcxPutSignal(const void *localbuff, size_t count,
     return flagcxInvalidArgument;
   if (peerWin == NULL)
     return flagcxInvalidArgument;
+  if (peer < 0 || peer >= comm->heteroComm->nRanks) {
+    WARN("flagcxPutSignal: peer %d out of range (nRanks=%d)", peer,
+         comm->heteroComm->nRanks);
+    return flagcxInvalidArgument;
+  }
 
   // Resolve window to MR index and compute byte size
   size_t byteSize = count * getFlagcxDataTypeSize(datatype);
@@ -2638,6 +2643,11 @@ flagcxResult_t flagcxSignal(int peer, unsigned int flags, flagcxComm_t comm,
   (void)flags;
   if (comm == NULL || comm->heteroComm == NULL)
     return flagcxInvalidArgument;
+  if (peer < 0 || peer >= comm->heteroComm->nRanks) {
+    WARN("flagcxSignal: peer %d out of range (nRanks=%d)", peer,
+         comm->heteroComm->nRanks);
+    return flagcxInvalidArgument;
+  }
 
   // Sender writes to its own slot in receiver's signal buffer
   size_t signalOffset = (size_t)comm->heteroComm->rank * sizeof(uint64_t);
@@ -2663,6 +2673,11 @@ flagcxResult_t flagcxWaitSignal(int nDesc,
 
   for (int i = 0; i < nDesc; i++) {
     int peer = signalDescs[i].peer;
+    if (peer < 0 || peer >= hetero->nRanks) {
+      WARN("flagcxWaitSignal: peer %d out of range (nRanks=%d)", peer,
+           hetero->nRanks);
+      return flagcxInvalidArgument;
+    }
     uint64_t opCnt = (uint64_t)signalDescs[i].opCnt;
 
     // Signal offset: peer's slot in my local signal buffer

flagcx/service/bootstrap.cc

@@ -139,15 +139,15 @@ static flagcxResult_t bootstrapNetSendRecv(struct flagcxSocket *sendSocket,
   int recvExpectedSize;
   FLAGCXCHECK(flagcxSocketSendRecv(sendSocket, &sendSize, sizeof(int),
                                    recvSocket, &recvExpectedSize, sizeof(int)));
-  if (recvExpectedSize > recvSize) {
-    WARN("bootstrapNetSendRecv: message truncated: received %d bytes instead "
-         "of %d",
+  if (recvExpectedSize < 0 || recvExpectedSize > recvSize) {
+    WARN("bootstrapNetSendRecv: invalid recvExpectedSize %d (recvSize=%d)",
          recvExpectedSize, recvSize);
     return flagcxInternalError;
   }
   // Exchange payload (interleaved, handles arbitrary sizes)
   FLAGCXCHECK(flagcxSocketSendRecv(sendSocket, sendData, sendSize, recvSocket,
-                                   recvData, recvExpectedSize));
+                                   recvData,
+                                   std::min(recvSize, recvExpectedSize)));
   return flagcxSuccess;
 }
 

test/unittest/rma/rma_test.cpp

@@ -17,7 +17,7 @@ void RmaTest::SetUpTestSuite() {
   MPI_Comm_size(MPI_COMM_WORLD, &nranks);
 
   size = RMA_TEST_SIZE;
-  signalSize = sizeof(uint64_t) * 64; // room for multiple signals
+  signalSize = sizeof(uint64_t) * nranks;
 
   flagcxDeviceHandleInit(&devHandle);
 
@@ -56,7 +56,14 @@ void RmaTest::SetUpTestSuite() {
   // Allocate and register signal buffer
   flagcxMemAlloc(&signalBuff, signalSize);
   devHandle->deviceMemset(signalBuff, 0, signalSize, flagcxMemDevice, nullptr);
-  flagcxOneSideSignalRegister(comm, signalBuff, signalSize, FLAGCX_PTR_CUDA);
+  res = flagcxOneSideSignalRegister(comm, signalBuff, signalSize,
+                                    FLAGCX_PTR_CUDA);
+  if (res != flagcxSuccess) {
+    flagcxMemFree(signalBuff);
+    signalBuff = nullptr;
+    dataWin = nullptr;
+    return;
+  }
 }
 
 void RmaTest::TearDownTestSuite() {