Skip to content

Commit 8ad0fea

Browse files
MC952-archMC952-arch
committed
mr-registry: make lookup APIs safe by construction with flagcxMrExtension
1 parent 9a1dc50 commit 8ad0fea

4 files changed

Lines changed: 424 additions & 106 deletions

File tree

‎flagcx/core/flagcx_mr_registry.cc‎

Lines changed: 186 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,49 @@ static flagcxResult_t ensureCapacity(struct flagcxMrRegistry *reg) {
105105
return flagcxSuccess;
106106
}
107107

108+
/*
109+
* Copy extension data into caller-provided flagcxMrExtension slots.
110+
* Each non-NULL slot gets a by-value copy tagged with FLAGCX_MR_OWNER_*.
111+
* NULL slots in the array are skipped. NULL array means "no extensions needed."
112+
*/
113+
static inline void copyExtensions(const struct flagcxMrEntry *entry,
114+
struct flagcxMrExtension *outExts[]) {
115+
if (outExts == NULL)
116+
return;
117+
if (outExts[FLAGCX_MR_OWNER_IDX_P2P]) {
118+
if (entry->p2p) {
119+
outExts[FLAGCX_MR_OWNER_IDX_P2P]->type = FLAGCX_MR_OWNER_P2P;
120+
outExts[FLAGCX_MR_OWNER_IDX_P2P]->p2p = *entry->p2p;
121+
} else {
122+
outExts[FLAGCX_MR_OWNER_IDX_P2P]->type = FLAGCX_MR_OWNER_NONE;
123+
}
124+
}
125+
if (outExts[FLAGCX_MR_OWNER_IDX_COLL]) {
126+
if (entry->coll) {
127+
outExts[FLAGCX_MR_OWNER_IDX_COLL]->type = FLAGCX_MR_OWNER_COLL;
128+
outExts[FLAGCX_MR_OWNER_IDX_COLL]->coll = *entry->coll;
129+
} else {
130+
outExts[FLAGCX_MR_OWNER_IDX_COLL]->type = FLAGCX_MR_OWNER_NONE;
131+
}
132+
}
133+
if (outExts[FLAGCX_MR_OWNER_IDX_RMA]) {
134+
if (entry->rma) {
135+
outExts[FLAGCX_MR_OWNER_IDX_RMA]->type = FLAGCX_MR_OWNER_RMA;
136+
outExts[FLAGCX_MR_OWNER_IDX_RMA]->rma = *entry->rma;
137+
} else {
138+
outExts[FLAGCX_MR_OWNER_IDX_RMA]->type = FLAGCX_MR_OWNER_NONE;
139+
}
140+
}
141+
}
142+
143+
static inline void sanitizeOutEntry(struct flagcxMrEntry *entry) {
144+
if (entry) {
145+
entry->p2p = NULL;
146+
entry->coll = NULL;
147+
entry->rma = NULL;
148+
}
149+
}
150+
108151
static void freeEntryExtensions(struct flagcxMrEntry *entry) {
109152
if (entry->p2p) {
110153
free(entry->p2p);
@@ -120,6 +163,78 @@ static void freeEntryExtensions(struct flagcxMrEntry *entry) {
120163
}
121164
}
122165

166+
/* ───── mrId index helpers ───── */
167+
168+
#define ID_INDEX_INITIAL_CAPACITY 16
169+
170+
/*
171+
* Append a new {mrId, baseAddr} pair to the end of idIndex.
172+
* Since mrIds are monotonically increasing, appending maintains sorted order.
173+
* Must be called under write lock.
174+
*/
175+
static flagcxResult_t idIndexAppend(struct flagcxMrRegistry *reg, uint64_t mrId,
176+
uintptr_t baseAddr) {
177+
if (reg->idCount >= reg->idCapacity) {
178+
int newCap =
179+
reg->idCapacity == 0 ? ID_INDEX_INITIAL_CAPACITY : reg->idCapacity * 2;
180+
struct flagcxMrIdEntry *newIdx = (struct flagcxMrIdEntry *)realloc(
181+
reg->idIndex, (size_t)newCap * sizeof(struct flagcxMrIdEntry));
182+
if (newIdx == NULL) {
183+
WARN("flagcxMrRegistry: idIndex realloc failed for capacity %d", newCap);
184+
return flagcxSystemError;
185+
}
186+
reg->idIndex = newIdx;
187+
reg->idCapacity = newCap;
188+
}
189+
reg->idIndex[reg->idCount].mrId = mrId;
190+
reg->idIndex[reg->idCount].baseAddr = baseAddr;
191+
reg->idCount++;
192+
return flagcxSuccess;
193+
}
194+
195+
/*
196+
* Binary search idIndex for mrId, remove entry with memmove.
197+
* Must be called under write lock.
198+
*/
199+
static void idIndexRemove(struct flagcxMrRegistry *reg, uint64_t mrId) {
200+
int lo = 0, hi = reg->idCount - 1;
201+
while (lo <= hi) {
202+
int mid = lo + (hi - lo) / 2;
203+
if (reg->idIndex[mid].mrId == mrId) {
204+
if (mid < reg->idCount - 1) {
205+
memmove(&reg->idIndex[mid], &reg->idIndex[mid + 1],
206+
(size_t)(reg->idCount - 1 - mid) *
207+
sizeof(struct flagcxMrIdEntry));
208+
}
209+
reg->idCount--;
210+
return;
211+
} else if (reg->idIndex[mid].mrId < mrId) {
212+
lo = mid + 1;
213+
} else {
214+
hi = mid - 1;
215+
}
216+
}
217+
}
218+
219+
/*
220+
* Binary search idIndex for mrId, return associated baseAddr.
221+
* Returns 0 if not found (valid baseAddr is never 0 for real registrations).
222+
*/
223+
static uintptr_t idIndexFindBaseAddr(const struct flagcxMrRegistry *reg,
224+
uint64_t mrId) {
225+
int lo = 0, hi = reg->idCount - 1;
226+
while (lo <= hi) {
227+
int mid = lo + (hi - lo) / 2;
228+
if (reg->idIndex[mid].mrId == mrId)
229+
return reg->idIndex[mid].baseAddr;
230+
else if (reg->idIndex[mid].mrId < mrId)
231+
lo = mid + 1;
232+
else
233+
hi = mid - 1;
234+
}
235+
return 0;
236+
}
237+
123238
/* ───── Lifecycle ───── */
124239

125240
flagcxResult_t flagcxMrRegistryCreate(struct flagcxMrRegistry **reg) {
@@ -132,6 +247,9 @@ flagcxResult_t flagcxMrRegistryCreate(struct flagcxMrRegistry **reg) {
132247
r->count = 0;
133248
r->capacity = 0;
134249
r->nextId = 1;
250+
r->idIndex = NULL;
251+
r->idCount = 0;
252+
r->idCapacity = 0;
135253

136254
if (pthread_rwlock_init(&r->rwlock, NULL) != 0) {
137255
free(r);
@@ -152,6 +270,7 @@ flagcxResult_t flagcxMrRegistryDestroy(struct flagcxMrRegistry *reg) {
152270
}
153271

154272
free(reg->entries);
273+
free(reg->idIndex);
155274
pthread_rwlock_destroy(&reg->rwlock);
156275
free(reg);
157276
return flagcxSuccess;
@@ -164,7 +283,7 @@ flagcxResult_t flagcxMrRegistryRegister(struct flagcxMrRegistry *reg,
164283
int ptrType, uint32_t ownerBit,
165284
void *mhandle, void *ext,
166285
uint64_t *outId) {
167-
if (reg == NULL || size == 0)
286+
if (reg == NULL || size == 0 || addr == 0)
168287
return flagcxInternalError;
169288

170289
int ownerIdx = ownerBitToIdx(ownerBit);
@@ -200,10 +319,20 @@ flagcxResult_t flagcxMrRegistryRegister(struct flagcxMrRegistry *reg,
200319
if (ext != NULL) {
201320
switch (ownerBit) {
202321
case FLAGCX_MR_OWNER_P2P: {
203-
/* Preserve the assigned mrId across ext replacement */
322+
/* Preserve the assigned mrId across ext replacement.
323+
* Reject if caller provides a conflicting non-zero mrId. */
204324
uint64_t prevMrId = existing->p2p ? existing->p2p->mrId : 0;
325+
struct flagcxMrP2pExt *newP2p = (struct flagcxMrP2pExt *)ext;
326+
if (newP2p->mrId != 0 && prevMrId != 0 &&
327+
newP2p->mrId != prevMrId) {
328+
WARN("flagcxMrRegistry: mrId conflict on ext replacement: "
329+
"existing %lu vs new %lu",
330+
(unsigned long)prevMrId, (unsigned long)newP2p->mrId);
331+
pthread_rwlock_unlock(&reg->rwlock);
332+
return flagcxInternalError;
333+
}
205334
free(existing->p2p);
206-
existing->p2p = (struct flagcxMrP2pExt *)ext;
335+
existing->p2p = newP2p;
207336
if (existing->p2p->mrId == 0)
208337
existing->p2p->mrId = prevMrId;
209338
break;
@@ -234,6 +363,14 @@ flagcxResult_t flagcxMrRegistryRegister(struct flagcxMrRegistry *reg,
234363
if (existing->p2p) {
235364
if (existing->p2p->mrId == 0)
236365
existing->p2p->mrId = reg->nextId++;
366+
if (idIndexAppend(reg, existing->p2p->mrId, addr) != flagcxSuccess) {
367+
/* Roll back: remove P2P ownership, caller retains ext */
368+
existing->p2p = NULL;
369+
existing->ownerMask &= ~ownerBit;
370+
existing->mhandles[ownerIdx] = NULL;
371+
pthread_rwlock_unlock(&reg->rwlock);
372+
return flagcxSystemError;
373+
}
237374
if (outId)
238375
*outId = existing->p2p->mrId;
239376
}
@@ -313,6 +450,16 @@ flagcxResult_t flagcxMrRegistryRegister(struct flagcxMrRegistry *reg,
313450
/* Assign mrId from registry's monotonic counter if not pre-set */
314451
if (entry->p2p->mrId == 0)
315452
entry->p2p->mrId = reg->nextId++;
453+
if (idIndexAppend(reg, entry->p2p->mrId, addr) != flagcxSuccess) {
454+
/* Roll back: remove the entry we just inserted */
455+
entry->p2p = NULL; /* caller retains ext ownership */
456+
if (pos < reg->count) {
457+
memmove(&reg->entries[pos], &reg->entries[pos + 1],
458+
(size_t)(reg->count - pos) * sizeof(struct flagcxMrEntry));
459+
}
460+
pthread_rwlock_unlock(&reg->rwlock);
461+
return flagcxSystemError;
462+
}
316463
if (outId)
317464
*outId = entry->p2p->mrId;
318465
}
@@ -361,14 +508,18 @@ flagcxResult_t flagcxMrRegistryDeregister(struct flagcxMrRegistry *reg,
361508
}
362509

363510
/* Copy out before modification */
364-
if (outEntry)
511+
if (outEntry) {
365512
*outEntry = *entry;
513+
sanitizeOutEntry(outEntry);
514+
}
366515

367516
/* Extract subsystem extension */
368517
void *ext = NULL;
369518
switch (ownerBit) {
370519
case FLAGCX_MR_OWNER_P2P:
371520
ext = entry->p2p;
521+
if (entry->p2p && entry->p2p->mrId != 0)
522+
idIndexRemove(reg, entry->p2p->mrId);
372523
entry->p2p = NULL;
373524
break;
374525
case FLAGCX_MR_OWNER_COLL:
@@ -406,7 +557,8 @@ flagcxResult_t flagcxMrRegistryDeregister(struct flagcxMrRegistry *reg,
406557

407558
flagcxResult_t flagcxMrRegistryLookup(struct flagcxMrRegistry *reg,
408559
uintptr_t addr,
409-
struct flagcxMrEntry *outEntry) {
560+
struct flagcxMrEntry *outEntry,
561+
struct flagcxMrExtension *outExts[]) {
410562
if (reg == NULL || outEntry == NULL)
411563
return flagcxInternalError;
412564

@@ -426,6 +578,8 @@ flagcxResult_t flagcxMrRegistryLookup(struct flagcxMrRegistry *reg,
426578
struct flagcxMrEntry *entry = &reg->entries[idx];
427579
if (addr >= entry->baseAddr && (addr - entry->baseAddr) < entry->size) {
428580
*outEntry = *entry;
581+
copyExtensions(entry, outExts);
582+
sanitizeOutEntry(outEntry);
429583
pthread_rwlock_unlock(&reg->rwlock);
430584
return flagcxSuccess;
431585
}
@@ -436,7 +590,8 @@ flagcxResult_t flagcxMrRegistryLookup(struct flagcxMrRegistry *reg,
436590

437591
flagcxResult_t flagcxMrRegistryFindExact(struct flagcxMrRegistry *reg,
438592
uintptr_t addr,
439-
struct flagcxMrEntry *outEntry) {
593+
struct flagcxMrEntry *outEntry,
594+
struct flagcxMrExtension *outExts[]) {
440595
if (reg == NULL || outEntry == NULL)
441596
return flagcxInternalError;
442597

@@ -449,33 +604,45 @@ flagcxResult_t flagcxMrRegistryFindExact(struct flagcxMrRegistry *reg,
449604
}
450605

451606
*outEntry = reg->entries[idx];
607+
copyExtensions(&reg->entries[idx], outExts);
608+
sanitizeOutEntry(outEntry);
452609
pthread_rwlock_unlock(&reg->rwlock);
453610
return flagcxSuccess;
454611
}
455612

456613
flagcxResult_t flagcxMrRegistryLookupById(struct flagcxMrRegistry *reg,
457614
uint64_t mrId,
458-
struct flagcxMrEntry *outEntry) {
615+
struct flagcxMrEntry *outEntry,
616+
struct flagcxMrExtension *outExts[]) {
459617
if (reg == NULL || outEntry == NULL)
460618
return flagcxInternalError;
461619

462620
pthread_rwlock_rdlock(&reg->rwlock);
463621

464-
for (int i = 0; i < reg->count; i++) {
465-
if (reg->entries[i].p2p && reg->entries[i].p2p->mrId == mrId) {
466-
*outEntry = reg->entries[i];
467-
pthread_rwlock_unlock(&reg->rwlock);
468-
return flagcxSuccess;
469-
}
622+
/* O(log n) lookup via mrId index → baseAddr → main entries[] */
623+
uintptr_t baseAddr = idIndexFindBaseAddr(reg, mrId);
624+
if (baseAddr == 0) {
625+
pthread_rwlock_unlock(&reg->rwlock);
626+
return flagcxInternalError;
627+
}
628+
629+
int idx = bsearchExact(reg->entries, reg->count, baseAddr);
630+
if (idx < 0 || !reg->entries[idx].p2p) {
631+
pthread_rwlock_unlock(&reg->rwlock);
632+
return flagcxInternalError;
470633
}
471634

635+
*outEntry = reg->entries[idx];
636+
copyExtensions(&reg->entries[idx], outExts);
637+
sanitizeOutEntry(outEntry);
472638
pthread_rwlock_unlock(&reg->rwlock);
473-
return flagcxInternalError;
639+
return flagcxSuccess;
474640
}
475641

476-
flagcxResult_t flagcxMrRegistryFindByHandle(struct flagcxMrRegistry *reg,
477-
int ownerIdx, void *mhandle,
478-
struct flagcxMrEntry *outEntry) {
642+
flagcxResult_t
643+
flagcxMrRegistryFindByHandle(struct flagcxMrRegistry *reg, int ownerIdx,
644+
void *mhandle, struct flagcxMrEntry *outEntry,
645+
struct flagcxMrExtension *outExts[]) {
479646
if (reg == NULL || outEntry == NULL || mhandle == NULL)
480647
return flagcxInternalError;
481648
if (ownerIdx < 0 || ownerIdx >= FLAGCX_MR_OWNER_COUNT)
@@ -486,6 +653,8 @@ flagcxResult_t flagcxMrRegistryFindByHandle(struct flagcxMrRegistry *reg,
486653
for (int i = 0; i < reg->count; i++) {
487654
if (reg->entries[i].mhandles[ownerIdx] == mhandle) {
488655
*outEntry = reg->entries[i];
656+
copyExtensions(&reg->entries[i], outExts);
657+
sanitizeOutEntry(outEntry);
489658
pthread_rwlock_unlock(&reg->rwlock);
490659
return flagcxSuccess;
491660
}

0 commit comments

Comments
 (0)