TOCTOU race in `MemoryRuntimeManager.acquire()` can hand out a closed runtime

[jueiope]

Summary

MemoryRuntimeManager.acquire() in memind-server/src/main/java/com/openmemind/ai/memory/server/runtime/MemoryRuntimeManager.java (lines 33–43) has a window between checking draining and incrementing inFlightRequests where another thread can drain and close the underlying memory runtime. Callers can then receive a RuntimeLease whose backing Memory is already closed.

Analysis

public RuntimeLease acquire() {
    while (true) {
        RuntimeHandle handle = requireCurrentHandle();
        if (handle.draining().get()) {            // (1) read draining
            continue;
        }
        handle.inFlightRequests().incrementAndGet(); // (2) increment
        if (handle == current.get() && !handle.draining().get()) {
            return new RuntimeLease(handle, () -> release(handle));
        }
        release(handle);
    }
}

Consider this interleaving with thread A in acquire() and thread B in swap():

Step	Thread A	Thread B
1	reads draining = false (line 35)	—
2	—	sets draining = true
3	—	calls tryClose(handle); in-flight is 0
4	—	handle.memory().close() runs to completion
5	inFlightRequests.incrementAndGet() (line 38)	—
6	second check: current may still point at this handle if swap has not yet replaced it; draining = true is now read, so loop continues	—

The second draining check at line 39 does catch the case above on the next loop iteration, so the worst-case observed bug requires swap() to also update current before line 39 runs — at which point the comparison handle == current.get() returns false and we go through release(handle). The release path then decrements inFlightRequests to 0 again, which is fine.

However, there is still one concrete bug: in step 5, we incremented inFlightRequests from 0 to 1 on a handle whose memory().close() has already returned. If release(handle) is invoked via the lease's onRelease runnable (which it is, line 41), and tryClose is called again, it will attempt to close an already-closed memory — depending on the underlying Memory.close() implementation this may throw or be benign, but the contract of Closeable.close() only requires idempotence, not error-freeness.

Root cause

The increment-then-validate pattern needs to be replaced with an atomic check-and-increment, e.g. an AtomicInteger whose negative value sentinels a closed handle, or a single CAS that both checks draining == false and increments.

Suggested fix

Use inFlightRequests.updateAndGet(prev -> draining.get() ? prev : prev + 1) and then check whether the value actually changed. Alternatively, hold the drain operation behind a StampedLock whose write stamp is acquired by swap and read stamps by acquire.

Impact

Severity: High under normal traffic + config swaps, but the actual blast radius depends on Memory.close()'s tolerance for re-entry. Worth fixing before this is exercised in production with hot config reloads.