fix: reserve shared GPU memory before CUDA alloc

Serialize check-and-add under lock_shrreg so concurrent processes cannot
both pass oom_check when only the sum exceeds the limit. Keep the expensive
CUDA allocation outside the lock; roll back the reservation on failure.

Author: Liauuu

src/allocator/allocator.c

@@ -27,9 +27,9 @@ extern CUresult cuMemoryFree(CUdeviceptr dptr);
 pthread_once_t allocator_allocate_flag = PTHREAD_ONCE_INIT;
 pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
 
-/* Optional delay after the unlocked oom pre-check in add_chunk.
- * Used only by the race-reproduce harness to widen the TOCTOU window
- * between check and usage accounting across processes.
+/* Optional delay after reservation, before the expensive CUDA alloc.
+ * With reserve-then-alloc this only stretches the CUDA window; shared usage
+ * is already committed so a concurrent peer must OOM.
  * Env: HAMI_ALLOC_RACE_WINDOW_US=<microseconds>, unset/0 = no delay. */
 static void maybe_widen_alloc_race_window(void) {
     const char *env = getenv("HAMI_ALLOC_RACE_WINDOW_US");
@@ -52,6 +52,23 @@ size_t round_up(size_t size, size_t unit) {
     return size;
 }
 
+int reserve_device_memory(CUdevice dev, size_t size) {
+    lock_shrreg();
+    if (oom_check(dev, size)) {
+        unlock_shrreg();
+        return CUDA_ERROR_OUT_OF_MEMORY;
+    }
+    add_gpu_device_memory_usage(getpid(), dev, size, 2);
+    unlock_shrreg();
+    return 0;
+}
+
+void release_device_memory(CUdevice dev, size_t size) {
+    lock_shrreg();
+    rm_gpu_device_memory_usage(getpid(), dev, size, 2);
+    unlock_shrreg();
+}
+
 int oom_check(const int dev, size_t addon) {
     CUdevice d;
     if (dev==-1)
@@ -131,8 +148,10 @@ int add_chunk(CUdeviceptr *address, size_t size) {
 
     cuCtxGetDevice(&dev);
 
-    /* OOM pre-check without lock */
-    if (oom_check(dev, size))
+    /* Reserve under the shared-region lock so concurrent processes cannot
+     * both pass oom_check before either commits usage. CUDA alloc stays
+     * outside the lock. */
+    if (reserve_device_memory(dev, size) != 0)
         return CUDA_ERROR_OUT_OF_MEMORY;
 
     maybe_widen_alloc_race_window();
@@ -145,44 +164,28 @@ int add_chunk(CUdeviceptr *address, size_t size) {
     }
     if (res != CUDA_SUCCESS) {
         LOG_ERROR("cuMemoryAllocate failed res=%d", res);
+        release_device_memory(dev, size);
         return res;
     }
 
-    /* Tracking inside lock — pure in-memory ops, microseconds */
+    /* Local list tracking only — usage already reserved */
     pthread_mutex_lock(&mutex);
-
-    if (oom_check(dev, size)) {
-        /* Another process consumed memory between our pre-check and now */
-        pthread_mutex_unlock(&mutex);
-        CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemFree_v2, *address);
-        return CUDA_ERROR_OUT_OF_MEMORY;
-    }
-
     allocated_list_entry *e;
     INIT_ALLOCATED_LIST_ENTRY(e, 0, size, dev);
     e->entry->address = *address;
     LIST_ADD(device_overallocated, e);
-    add_gpu_device_memory_usage(getpid(), dev, size, 2);
-
     pthread_mutex_unlock(&mutex);
     return 0;
 }
 
+/* Track a pointer in the local list. Caller must already have reserved
+ * `size` via reserve_device_memory() (or equivalent usage accounting). */
 int add_chunk_only(CUdeviceptr address, size_t size, CUdevice dev) {
     pthread_mutex_lock(&mutex);
-    size_t addr=0;
-    size_t allocsize;
-    if (oom_check(dev,size)){
-        pthread_mutex_unlock(&mutex);
-        return -1;
-    }
     allocated_list_entry *e;
-    INIT_ALLOCATED_LIST_ENTRY(e, addr, size, dev);
-    LIST_ADD(device_overallocated,e);
-    //uint64_t t_size;
-    e->entry->address=address;
-    allocsize = size;
-    add_gpu_device_memory_usage(getpid(), dev, allocsize, 2);
+    INIT_ALLOCATED_LIST_ENTRY(e, 0, size, dev);
+    e->entry->address = address;
+    LIST_ADD(device_overallocated, e);
     pthread_mutex_unlock(&mutex);
     return 0;
 }

src/allocator/allocator.h

@@ -156,6 +156,13 @@ CUresult view_vgpu_allocator();
 // Checks if oom
 int oom_check(const int dev,size_t addon);
 
+/* Cross-process check-and-reserve under lock_shrreg.
+ * Returns 0 on success, CUDA_ERROR_OUT_OF_MEMORY if the limit would be exceeded.
+ * On success, size is already accounted in shared usage; call
+ * release_device_memory() if the subsequent CUDA alloc fails. */
+int reserve_device_memory(CUdevice dev, size_t size);
+void release_device_memory(CUdevice dev, size_t size);
+
 // Allocate and free device memory
 int allocate_raw(CUdeviceptr *dptr, size_t size);
 int free_raw(CUdeviceptr dptr);

src/cuda/memory.c

@@ -145,12 +145,14 @@ CUresult cuMemAllocManaged(CUdeviceptr* dptr, size_t bytesize, unsigned int flag
     ENSURE_RUNNING();
     CUdevice dev;
     CHECK_DRV_API(cuCtxGetDevice(&dev));
-    if (oom_check(dev,bytesize)){
+    if (reserve_device_memory(dev, bytesize) != 0) {
         return CUDA_ERROR_OUT_OF_MEMORY;
     }
     CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemAllocManaged, dptr, bytesize, flags);
     if (res == CUDA_SUCCESS) {
         add_chunk_only(*dptr, bytesize, dev);
+    } else {
+        release_device_memory(dev, bytesize);
     }
     return res;
 }
@@ -163,12 +165,14 @@ CUresult cuMemAllocPitch_v2(CUdeviceptr* dptr, size_t* pPitch, size_t WidthInByt
     ENSURE_RUNNING();
     CUdevice dev;
     CHECK_DRV_API(cuCtxGetDevice(&dev));
-    if (oom_check(dev,bytesize)){
+    if (reserve_device_memory(dev, bytesize) != 0) {
         return CUDA_ERROR_OUT_OF_MEMORY;
     }
     CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,cuMemAllocPitch_v2, dptr, pPitch, WidthInBytes, Height, ElementSizeBytes);
     if (res == CUDA_SUCCESS) {
         add_chunk_only(*dptr, bytesize, dev);
+    } else {
+        release_device_memory(dev, bytesize);
     }
     return res;
 }
@@ -590,13 +594,17 @@ CUresult cuMemCreate ( CUmemGenericAllocationHandle* handle, size_t size, const
     if (do_oom_check && cuCtxGetDevice(&dev) != CUDA_SUCCESS) {
         dev = prop->location.id;
     }
-    if (do_oom_check && oom_check(dev, size)) {
+    if (do_oom_check && reserve_device_memory(dev, size) != 0) {
         return CUDA_ERROR_OUT_OF_MEMORY;
     }
     CUresult res = CUDA_OVERRIDE_CALL(cuda_library_entry,
         cuMemCreate, handle, size, prop, flags);
-    if (do_oom_check && res == CUDA_SUCCESS) {
-        add_chunk_only(*handle, size, dev);
+    if (do_oom_check) {
+        if (res == CUDA_SUCCESS) {
+            add_chunk_only(*handle, size, dev);
+        } else {
+            release_device_memory(dev, size);
+        }
     }
     return res;
 }

test/run_concurrent_oom_race.sh

@@ -1,5 +1,6 @@
 #!/usr/bin/env bash
-# Run the cross-process oom_check TOCTOU reproducer against libvgpu.so.
+# Run the cross-process oom_check race test against libvgpu.so (race-fix).
+# Default: --expect-fixed (limit must hold under concurrent alloc + race window).
 set -euo pipefail
 
 ROOT="$(cd "$(dirname "$0")/.." && pwd)"
@@ -31,8 +32,14 @@ export HAMI_RACE_ROUNDS="$ROUNDS"
 export HAMI_ALLOC_RACE_WINDOW_US="$WINDOW_US"
 export LIBCUDA_LOG_LEVEL="${LIBCUDA_LOG_LEVEL:-1}"
 
-echo "Running: $BIN $*"
+# Default to verifying the fix; omit args only triggers --expect-fixed.
+ARGS=("$@")
+if [[ ${#ARGS[@]} -eq 0 ]]; then
+  ARGS=(--expect-fixed)
+fi
+
+echo "Running: $BIN ${ARGS[*]}"
 echo "  LD_PRELOAD=$LD_PRELOAD"
 echo "  LIMIT=$LIMIT ALLOC=$ALLOC ROUNDS=$ROUNDS WINDOW_US=$WINDOW_US"
 echo "  CACHE=$CACHE"
-exec "$BIN" "$@"
+exec "$BIN" "${ARGS[@]}"

test/test_concurrent_oom_race.c

@@ -24,6 +24,7 @@
  *   ./build/test/test_concurrent_oom_race
  *
  * Or: ./test/run_concurrent_oom_race.sh
+ *   (race-fix defaults to --expect-fixed)
  *
  * Suggested sizing: limit=1024m, alloc=600m so each request fits alone
  * (after two process contexts) but 2*alloc exceeds the limit.