fix: prevent cross-process GPU memory overcommit on concurrent cuMemAlloc.

src/allocator/allocator.c

@@ -33,7 +33,8 @@ size_t round_up(size_t size, size_t unit) {
     return size;
 }
 
-int oom_check(const int dev, size_t addon) {
+/* already_locked: caller already holds lock_shrreg (not reentrant). */
+static int oom_check_impl(const int dev, size_t addon, int already_locked) {
     CUdevice d;
     if (dev==-1)
         cuCtxGetDevice(&d);
@@ -49,18 +50,45 @@ int oom_check(const int dev, size_t addon) {
     size_t new_allocated = _usage + addon;
     LOG_INFO("_usage=%lu limit=%lu new_allocated=%lu",_usage,limit,new_allocated);
     if (new_allocated > limit) {
+        int cleared;
+
         LOG_ERROR("Device %d OOM %lu / %lu", d, new_allocated, limit);
 
-        lock_shrreg();
-        int cleared = clear_proc_slot_nolock(1);
-        unlock_shrreg();
+        if (already_locked) {
+            cleared = clear_proc_slot_nolock(1);
+        } else {
+            lock_shrreg();
+            cleared = clear_proc_slot_nolock(1);
+            unlock_shrreg();
+        }
         if (cleared > 0)
-            return oom_check(dev,addon);
+            return oom_check_impl(dev, addon, already_locked);
         return 1;
     }
     return 0;
 }
 
+int oom_check(const int dev, size_t addon) {
+    return oom_check_impl(dev, addon, 0);
+}
+
+int reserve_device_memory(CUdevice dev, size_t size) {
+    lock_shrreg();
+    if (oom_check_impl(dev, size, 1)) {
+        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();
+}
+
 CUresult view_vgpu_allocator() {
     allocated_list_entry *al;
     size_t total;
@@ -109,14 +137,28 @@ void allocator_init() {
     pthread_mutex_init(&mutex,NULL);
 }
 
+/* Wrap INIT_ALLOCATED_LIST_ENTRY so QUIT_WITH_ERROR returns here, not from
+ * callers that hold mutex / CUDA memory / shared reservations. */
+static int new_allocated_list_entry(allocated_list_entry **out,
+                                    CUdeviceptr address, size_t size,
+                                    CUdevice dev) {
+    allocated_list_entry *e;
+    INIT_ALLOCATED_LIST_ENTRY(e, address, size, dev);
+    *out = e;
+    return 0;
+}
+
 int add_chunk(CUdeviceptr *address, size_t size) {
     CUdevice dev;
     CUresult res;
+    allocated_list_entry *e;
 
     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;
 
     /* GPU allocation outside lock, the expensive part */
@@ -127,44 +169,36 @@ 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 */
+    if (new_allocated_list_entry(&e, 0, size, dev) != 0) {
         pthread_mutex_unlock(&mutex);
-        CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemFree_v2, *address);
-        return CUDA_ERROR_OUT_OF_MEMORY;
+        cuMemoryFree(*address);
+        release_device_memory(dev, size);
+        return -1;
     }
-
-    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) {
+    allocated_list_entry *e;
+
     pthread_mutex_lock(&mutex);
-    size_t addr=0;
-    size_t allocsize;
-    if (oom_check(dev,size)){
+    if (new_allocated_list_entry(&e, 0, size, dev) != 0) {
         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);
+    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,18 @@ 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);
+        if (add_chunk_only(*dptr, bytesize, dev) != 0) {
+            CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemFree_v2, *dptr);
+            release_device_memory(dev, bytesize);
+            return CUDA_ERROR_OUT_OF_MEMORY;
+        }
+    } else {
+        release_device_memory(dev, bytesize);
     }
     return res;
 }
@@ -161,17 +167,40 @@ CUresult cuMemAllocPitch_v2(CUdeviceptr* dptr, size_t* pPitch, size_t WidthInByt
     size_t guess_pitch = (ElementSizeBytes == 0 || WidthInBytes == 0) ? 0 :
         (((WidthInBytes - 1) / ElementSizeBytes) + 1) * ElementSizeBytes;
     size_t bytesize = guess_pitch * Height;
+    size_t actual;
     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);
+    if (res != CUDA_SUCCESS) {
+        release_device_memory(dev, bytesize);
+        return res;
     }
-    return res;
+    /* Driver pitch may exceed guess_pitch due to alignment; account for real size. */
+    if (Height != 0 && pPitch != NULL && *pPitch > SIZE_MAX / Height) {
+        CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemFree_v2, *dptr);
+        release_device_memory(dev, bytesize);
+        return CUDA_ERROR_OUT_OF_MEMORY;
+    }
+    actual = (pPitch != NULL) ? (*pPitch * Height) : bytesize;
+    if (actual > bytesize) {
+        if (reserve_device_memory(dev, actual - bytesize) != 0) {
+            CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemFree_v2, *dptr);
+            release_device_memory(dev, bytesize);
+            return CUDA_ERROR_OUT_OF_MEMORY;
+        }
+    } else if (actual < bytesize) {
+        release_device_memory(dev, bytesize - actual);
+    }
+    if (add_chunk_only(*dptr, actual, dev) != 0) {
+        CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemFree_v2, *dptr);
+        release_device_memory(dev, actual);
+        return CUDA_ERROR_OUT_OF_MEMORY;
+    }
+    return CUDA_SUCCESS;
 }
 
 CUresult cuMemFree_v2(CUdeviceptr dptr) {
@@ -595,13 +624,21 @@ 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) {
+            if (add_chunk_only(*handle, size, dev) != 0) {
+                CUDA_OVERRIDE_CALL(cuda_library_entry, cuMemRelease, *handle);
+                release_device_memory(dev, size);
+                return CUDA_ERROR_OUT_OF_MEMORY;
+            }
+        } else {
+            release_device_memory(dev, size);
+        }
     }
     return res;
 }

test/CMakeLists.txt

@@ -69,6 +69,16 @@ if (TARGET vgpu)
         TIMEOUT 10)
 endif()
 
+# GPU + LD_PRELOAD regression for cross-process cuMemAlloc TOCTOU (#273).
+# Uses the wrapper so --expect-fixed and limit/cache env are set.
+# Wrapper exits 77 when no NVIDIA GPU is present (skip, not fail).
+add_test(NAME concurrent_oom_race
+    COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/run_concurrent_oom_race.sh --expect-fixed)
+set_tests_properties(concurrent_oom_race PROPERTIES
+    TIMEOUT 180
+    SKIP_RETURN_CODE 77
+    ENVIRONMENT "LIBVGPU=${CMAKE_BINARY_DIR}/libvgpu.so;TEST_BIN=${CMAKE_CURRENT_BINARY_DIR}/test_concurrent_oom_race")
+
 
 add_custom_target(python_test ALL
     COMMAND cp -r ${CMAKE_CURRENT_SOURCE_DIR}/python ${CMAKE_CURRENT_BINARY_DIR})

test/run_concurrent_oom_race.sh

@@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+# Run the cross-process oom_check race test against libvgpu.so.
+# Default: --expect-fixed (limit must hold under concurrent alloc).
+# Exit 77 = no GPU (CTest SKIP_RETURN_CODE).
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+LIB="${LIBVGPU:-$ROOT/build/libvgpu.so}"
+BIN="${TEST_BIN:-$ROOT/build/test/test_concurrent_oom_race}"
+CACHE="${CUDA_DEVICE_MEMORY_SHARED_CACHE:-/tmp/hami_oom_race.cache}"
+LIMIT="${CUDA_DEVICE_MEMORY_LIMIT:-1024m}"
+ALLOC="${HAMI_RACE_ALLOC:-600m}"
+ROUNDS="${HAMI_RACE_ROUNDS:-30}"
+
+# Detect a usable NVIDIA GPU before LD_PRELOAD so CPU-only hosts skip cleanly.
+if ! command -v nvidia-smi >/dev/null 2>&1 || ! nvidia-smi -L >/dev/null 2>&1; then
+  echo "SKIP: no NVIDIA GPU available (nvidia-smi)" >&2
+  exit 77
+fi
+
+if [[ ! -f "$LIB" ]]; then
+  echo "missing $LIB — build first: (cd \"$ROOT\" && ./build.sh)" >&2
+  exit 1
+fi
+if [[ ! -x "$BIN" ]]; then
+  echo "missing $BIN — rebuild so test/CMakeLists.txt picks up the new test" >&2
+  exit 1
+fi
+
+mkdir -p /tmp/vgpulock
+rm -f "$CACHE"
+
+export LD_PRELOAD="$LIB"
+export CUDA_DEVICE_MEMORY_SHARED_CACHE="$CACHE"
+export CUDA_DEVICE_MEMORY_LIMIT="$LIMIT"
+export HAMI_RACE_ALLOC="$ALLOC"
+export HAMI_RACE_ROUNDS="$ROUNDS"
+export LIBCUDA_LOG_LEVEL="${LIBCUDA_LOG_LEVEL:-1}"
+
+# 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"
+echo "  CACHE=$CACHE"
+exec "$BIN" "${ARGS[@]}"