CUDA: allocation streams are handed back, not leaked; a handled pool failure leaves no error

A long-running broker began cancelling every data collection with

    Device decoding failed (CUDA (GPU) error (out of memory)), falling back to host decompression
    CUDA (GPU) error (out of memory)

while nvidia-smi showed the cards less than a fifth full. Two defects, both in CUDAMemHelpers.h and
both from the pooled allocator (cudaMallocAsync) that came with rc.162.

The leak. cuda_allocation_stream() kept one stream per (thread, device) in a thread_local map of raw
cudaStream_t and never destroyed them. That was written against rugnux, where the worker threads live
as long as the process. The broker starts fresh std::async threads for every data collection - 16 or
64 of them - so every collection left that many streams behind. Measured: 0.56 MB of device memory
per leaked stream, linear to 4928 streams, never returned, with nothing on the host side growing.

The RAII wrapper (CudaStream) was there but not used at this site, and using it as-is - a stream
destroyed when its thread exits - would not have been safe: ShadowFinder builds its GPU accumulator
on a throw-away std::async thread and frees it from another thread long after, and that free is
ordered on the allocating thread's stream. So the streams are still never destroyed, but a thread
now only borrows one: CudaStream objects live in a process-wide per-device idle list, a thread takes
one on first use and hands it back when it exits. Their number is bounded by the threads that were
ever alive at once instead of by the threads ever started. The list itself is deliberately leaked, so
that nothing calls into CUDA during static destruction.

Replaying the broker's pattern against the real header, 60 collections of 64 threads:
    before   3840 streams, 260 -> 2424 MB of device memory
    after      64 streams, 260 ->  358 MB

The stale error. Every helper here throws a named message, yet the log carried the raw CUDA string,
so the failure came through a cuda_err() and not from an allocation. CudaDevicePtr falls back to
cudaMalloc when cudaMallocAsync fails, silently - but the failed call stays behind as the thread's
last error, and the cudaGetLastError() that follows the next kernel launch reports it. The buffers
were all allocated; the frame was lost anyway, once on the device-decode route (caught, hence the
warning) and once more on the host fallback (fatal). The pooled attempt failing, and a stream that
cannot be created, are both handled by falling back, so both now clear the error they leave.

What finite resource the production cards ran out of at under 4 GB used was not established - no
cap on the number of streams was found up to 4928 on the card this was measured on. The leak is the
only thing on this path that grows with uptime.

tests/CUDAMemHelpersTest.cpp: later threads end up on the same stream, concurrent threads on
different ones, a buffer is freed cleanly after its allocating thread has exited (and another has
borrowed its stream), and a pool that cannot serve a request leaves no error behind - the last by
capping a memory pool at 4 MB so that the pooled attempt fails and the fallback succeeds.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

docs/CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## 1.0.0
 
+### 1.0.0-rc.172
+
+* Fixed `jfjoch_broker` cancelling every data collection with a CUDA "out of memory" error after long operation: GPU memory no longer leaks with each collection.
+
 ### 1.0.0-rc.171
 
 * Rugnux: basic support for CCD images (marCCD, SMV) and for gzipped miniCBF.

image_analysis/indexing/CUDAMemHelpers.h

@@ -6,6 +6,7 @@
 #include <cuda_runtime.h>
 #include <cufft.h>
 #include <map>
+#include <mutex>
 #include <stdexcept>
 #include <vector>
 #include "../common/JFJochException.h"
@@ -126,8 +127,8 @@ public:
 // the image loop, so do not hold everything.
 constexpr uint64_t CUDA_MEM_POOL_RELEASE_THRESHOLD = 1ull << 30;
 
-// The stream device allocations are ordered on: one per (thread, device), created on first use and
-// never destroyed.
+// The stream device allocations are ordered on: one per (thread, device), borrowed on first use and
+// handed back when the thread exits. The streams themselves are never destroyed.
 //
 // cudaMalloc and cudaFree are on CUDA's implicit-synchronization list - each one synchronises the
 // device across every stream - so one analysis engine per worker thread, each making a few dozen of
@@ -137,24 +138,70 @@ constexpr uint64_t CUDA_MEM_POOL_RELEASE_THRESHOLD = 1ull << 30;
 // The stream is deliberately not owned by the engine: cudaFreeAsync must be ordered on a stream that
 // is still alive, and an engine's own stream can be destroyed before the buffers it allocated. It is
 // keyed by device because a thread pinned to one GPU must not order work on another's stream.
+//
+// Nor is it owned by the thread, for the same reason: a buffer can outlive the thread that allocated
+// it - the shadow finder builds its accumulator on a throw-away thread and frees it from another,
+// long after - so the thread only borrows the stream. But it does hand it back. The broker starts
+// fresh worker threads for every data collection, and a stream holds about half a megabyte of device
+// memory, so a stream per thread ever started is a leak that a long-running broker does not survive.
+// Handed back, the streams number as many as the threads that were ever alive at once.
+namespace jfjoch_cuda_allocation_streams {
+    struct Idle {
+        std::mutex m;
+        std::map<int, std::vector<CudaStream>> streams;   // by device
+    };
+
+    // Never destroyed: a thread can still be handing a stream back while the process exits, and
+    // destroying streams during static destruction would call into a runtime that may be gone.
+    inline Idle &idle() {
+        static Idle *i = new Idle;
+        return *i;
+    }
+
+    struct Borrowed {
+        std::map<int, CudaStream> streams;   // by device
+        ~Borrowed() {
+            std::lock_guard lock(idle().m);
+            for (auto &[device, stream] : streams)
+                idle().streams[device].push_back(std::move(stream));
+        }
+    };
+}
+
 inline cudaStream_t cuda_allocation_stream() {
     int device = 0;
     if (cudaGetDevice(&device) != cudaSuccess)
         return nullptr;
-    thread_local std::map<int, cudaStream_t> streams;
-    const auto it = streams.find(device);
-    if (it != streams.end())
+    thread_local jfjoch_cuda_allocation_streams::Borrowed borrowed;
+    const auto it = borrowed.streams.find(device);
+    if (it != borrowed.streams.end())
         return it->second;
-    cudaStream_t stream = nullptr;
-    if (cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking) != cudaSuccess)
+
+    auto &idle = jfjoch_cuda_allocation_streams::idle();
+    {
+        std::lock_guard lock(idle.m);
+        auto &streams = idle.streams[device];
+        if (!streams.empty()) {
+            CudaStream stream = std::move(streams.back());
+            streams.pop_back();
+            return borrowed.streams.emplace(device, std::move(stream)).first->second;
+        }
+    }
+
+    try {
+        CudaStream stream;
+        cudaMemPool_t pool = nullptr;
+        if (cudaDeviceGetDefaultMemPool(&pool, device) == cudaSuccess) {
+            uint64_t threshold = CUDA_MEM_POOL_RELEASE_THRESHOLD;
+            cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold);
+        }
+        return borrowed.streams.emplace(device, std::move(stream)).first->second;
+    } catch (const JFJochException &) {
+        // No stream: the caller allocates synchronously instead, so this is handled - see
+        // CudaDevicePtr for why it must then not stay behind as the thread's last error.
+        cudaGetLastError();
         return nullptr;
-    cudaMemPool_t pool = nullptr;
-    if (cudaDeviceGetDefaultMemPool(&pool, device) == cudaSuccess) {
-        uint64_t threshold = CUDA_MEM_POOL_RELEASE_THRESHOLD;
-        cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold);
     }
-    streams.emplace(device, stream);
-    return stream;
 }
 
 // Which allocator a device buffer uses. Pooled is what every per-worker buffer wants. Synchronous is
@@ -183,14 +230,21 @@ public:
     explicit CudaDevicePtr(size_t count, CudaAlloc alloc = CudaAlloc::Pooled) {
         if (alloc == CudaAlloc::Pooled)
             stream_ = cuda_allocation_stream();
-        if (stream_ && cudaMallocAsync(&ptr_, count * sizeof(T), stream_) == cudaSuccess) {
-            // The allocation is ordered on this stream, so anything that wants to use the memory from
-            // another stream has to be ordered after it. Waiting for it once here is that ordering,
-            // and it leaves the pointer as freely usable as cudaMalloc's would have been.
-            cudaStreamSynchronize(stream_);
-            return;
+        if (stream_) {
+            if (cudaMallocAsync(&ptr_, count * sizeof(T), stream_) == cudaSuccess) {
+                // The allocation is ordered on this stream, so anything that wants to use the memory
+                // from another stream has to be ordered after it. Waiting for it once here is that
+                // ordering, and it leaves the pointer as freely usable as cudaMalloc's would have been.
+                cudaStreamSynchronize(stream_);
+                return;
+            }
+            // The pool failing is handled - by the synchronous allocation below - so it must not stay
+            // behind as the thread's last error. Every kernel launch is followed by cudaGetLastError(),
+            // which would report it there: an "out of memory" thrown over an image whose buffers were
+            // all allocated and whose kernel launched.
+            cudaGetLastError();
+            stream_ = nullptr;
         }
-        stream_ = nullptr;
         if (cudaMalloc(&ptr_, count * sizeof(T)) != cudaSuccess)
             throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                                   "Failed to allocate device memory");

tests/CMakeLists.txt

@@ -76,6 +76,7 @@ ADD_EXECUTABLE(jfjoch_test
         ImagePreprocessorGPUFusedTest.cpp
         AdaptiveSpotFinderGPUTest.cpp
         SpotExtractorGPUParityTest.cpp
+        CUDAMemHelpersTest.cpp
         CalcBraggPredictionTest.cpp
         SpotUtilsTest.cpp
         LatticeSearchTest.cpp

tests/CUDAMemHelpersTest.cpp

@@ -0,0 +1,103 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute
+// SPDX-License-Identifier: GPL-3.0-only
+
+#include <catch2/catch_all.hpp>
+#include "../common/CUDAWrapper.h"
+
+#ifdef JFJOCH_USE_CUDA
+
+#include <future>
+#include <latch>
+#include <set>
+#include <thread>
+
+#include "../image_analysis/indexing/CUDAMemHelpers.h"
+
+// The broker starts fresh worker threads for every data collection. A stream per thread ever started
+// is a device-memory leak that ends a long-running broker, so threads that follow one another have to
+// end up on the same stream.
+TEST_CASE("CudaAllocationStream_ReusedByLaterThreads", "[CUDAMemHelpers]") {
+    if (get_gpu_count() == 0)
+        SKIP("No CUDA GPU present");
+
+    std::set<cudaStream_t> seen;
+    for (int i = 0; i < 100; i++)
+        std::thread([&seen] {
+            CudaDevicePtr<uint8_t> buffer(1 << 16);
+            seen.insert(cuda_allocation_stream());
+        }).join();
+
+    CHECK(seen.size() == 1);
+    CHECK(*seen.begin() != nullptr);
+}
+
+TEST_CASE("CudaAllocationStream_DistinctForConcurrentThreads", "[CUDAMemHelpers]") {
+    if (get_gpu_count() == 0)
+        SKIP("No CUDA GPU present");
+
+    // Neither thread may exit - and hand its stream back - before the other has taken its own.
+    std::latch both_have_one(2);
+    auto take = [&both_have_one] {
+        const cudaStream_t stream = cuda_allocation_stream();
+        both_have_one.count_down();
+        both_have_one.wait();
+        return stream;
+    };
+    auto a = std::async(std::launch::async, take);
+    auto b = std::async(std::launch::async, take);
+    const cudaStream_t stream_a = a.get();
+    const cudaStream_t stream_b = b.get();
+
+    CHECK(stream_a != nullptr);
+    CHECK(stream_b != nullptr);
+    CHECK(stream_a != stream_b);
+}
+
+// The shadow finder builds its GPU accumulator on a throw-away thread and frees it from another one
+// long after. The free is ordered on the allocating thread's stream, so that stream has to outlive
+// the thread - including while a later thread has borrowed it.
+TEST_CASE("CudaDevicePtr_FreedAfterAllocatingThreadExited", "[CUDAMemHelpers]") {
+    if (get_gpu_count() == 0)
+        SKIP("No CUDA GPU present");
+
+    auto buffer = std::async(std::launch::async, [] {
+        return std::make_unique<CudaDevicePtr<uint8_t>>(1 << 20);
+    }).get();
+    std::thread([] { CudaDevicePtr<uint8_t> other(1 << 16); }).join();
+
+    REQUIRE(cudaMemset(buffer->get(), 0, 1 << 20) == cudaSuccess);
+    buffer.reset();
+    CHECK(cudaDeviceSynchronize() == cudaSuccess);
+    CHECK(cudaGetLastError() == cudaSuccess);
+}
+
+// A pool that cannot serve the request is not an error: the buffer comes from cudaMalloc instead. It
+// must then not stay behind as the thread's last error, or the cudaGetLastError() after the next
+// kernel launch throws "out of memory" over work that went fine.
+TEST_CASE("CudaDevicePtr_HandledPoolFailureLeavesNoError", "[CUDAMemHelpers]") {
+    if (get_gpu_count() == 0)
+        SKIP("No CUDA GPU present");
+
+    int device = 0;
+    REQUIRE(cudaGetDevice(&device) == cudaSuccess);
+
+    cudaMemPoolProps props{};
+    props.allocType = cudaMemAllocationTypePinned;
+    props.location.type = cudaMemLocationTypeDevice;
+    props.location.id = device;
+    props.maxSize = 4 << 20;
+
+    cudaMemPool_t previous = nullptr, capped = nullptr;
+    REQUIRE(cudaDeviceGetMemPool(&previous, device) == cudaSuccess);
+    REQUIRE(cudaMemPoolCreate(&capped, &props) == cudaSuccess);
+    REQUIRE(cudaDeviceSetMemPool(device, capped) == cudaSuccess);
+    {
+        CudaDevicePtr<uint8_t> buffer(64 << 20);
+        CHECK(buffer.get() != nullptr);
+        CHECK(cudaGetLastError() == cudaSuccess);
+    }
+    REQUIRE(cudaDeviceSetMemPool(device, previous) == cudaSuccess);
+    cudaMemPoolDestroy(capped);
+}
+
+#endif