tensor-pool: free_tensor_pool destroys the caller's own pool when freeing a peer's pool (missing owner guard on PINNED_POOL cleanup)

[phil-opp]

Note: This issue was created by a scheduled automated Claude check that randomly selects a component and reviews it for correctness bugs. Please verify before acting on it.

Summary

In the opt-in tensor-pool extension, TensorPoolTransport::free_tensor_pool performs its sender-side cleanup of the process-global PINNED_POOL (and TRANSIT_META, plus the CUDA GPU/transit buffers) by extracting only the bare counter from the pool id, with no check that the pool is owned by the calling node. Because PINNED_POOL is keyed by a per-process monotonic counter (not namespaced by node id), a node that frees a peer's pool ends up freeing its own same-counter pool instead — freeing a live GPU buffer whose IPC handle is already exported to a downstream node.

The sibling cleanup path, process_pending_tensor_pool_frees, does the identical work but correctly guards on ownership, with a comment explicitly explaining why. free_tensor_pool is missing that guard.

File: libraries/extensions/tensor-pool/python/src/transport.rs

The two paths, side by side

process_pending_tensor_pool_frees — correct (has the owner guard): (lines ~1198-1219)

// Guard against cross-process counter aliasing: buffer ids are
// pool_{node_id}_{counter}.  Extract the owner segment ... and require
// an exact equality match ...
if let Some(owner_and_counter) = buffer_id.strip_prefix("pool_")
    && let Some((owner, counter_str)) = owner_and_counter.rsplit_once('_')
    && owner == self.node_id.as_ref()          // <-- owner check
    && let Ok(c) = counter_str.parse::<u64>()
    && let Some(slot) = PINNED_POOL.lock()...remove(&c)
{
    // _unregister_host(slot.base); _free_gpu_buf(c); _free_transit(...); TRANSIT_META.remove(&c);
}

free_tensor_pool — buggy (no owner check): (lines ~2554-2581)

// PINNED_POOL is sender-side (per-process), so bare counter is sufficient.   // <-- wrong assumption
{
    let counter = buffer_id
        .strip_prefix("pool_")
        .and_then(|s| s.rsplit_once('_').map(|(_, c)| c))   // owner segment discarded
        .and_then(|c| c.parse::<u64>().ok());
    if let Some(c) = counter
        && let Some(slot) = PINNED_POOL.lock()...remove(&c)  // <-- no owner == self check
    {
        if let Ok(helpers) = get_cuda_helpers(py) {
            let bound = helpers.bind(py);
            let _ = bound.call_method1("_unregister_host", (slot.base,));
            let _ = bound.call_method1("_free_gpu_buf", (c,));
            if slot.transit_ptr != 0 {
                let _ = bound.call_method1("_free_transit", (slot.transit_ptr,));
            }
        }
        TRANSIT_META.lock()...remove(&c);
    }
    // PoolSlot dropped here -> Shmem unmapped
}

The comment "PINNED_POOL is sender-side (per-process), so bare counter is sufficient" only holds if the freeing node is always the owner. But the public API explicitly supports a receiver freeing a pool it received (docstring on register_tensor_pool, lines ~1229-1231: "The returned pool ID can be shared across nodes ... so that a receiver can call read_tensor_pool and free_tensor_pool on it").

Why the counters collide

• PINNED_COUNTER is a process-global Mutex<u64> starting at 0, incremented before each registration (line ~1298). Every process's first pool is counter 1, the second is 2, and so on.
• PINNED_POOL is HashMap<u64, PoolSlot> keyed by that bare counter (insert at line ~1662), not namespaced by node id.
• The pool id / buffer_id is format!("pool_{}_{}", self.node_id, pool_counter) (line ~1684).

So counter-1 collisions across nodes are the norm, not a corner case.

Concrete failure scenario (a standard A → B → C pipeline)

1. Node B registers its own output pool → pool_B_1; in B's process PINNED_POOL[1] = B's live slot (GPU buffer + exported IPC handle).
2. Node A registers its output pool → pool_A_1 and sends the id to B.
3. B reads pool_A_1, then calls free_tensor_pool("pool_A_1") to release it promptly.
4. In B's process: counter = 1, so PINNED_POOL.remove(&1) returns B's own pool_B_1 slot, not anything belonging to A.
5. B then runs _unregister_host(slot.base), _free_gpu_buf(1) (cudaFrees B's own pooled GPU buffer), _free_transit(...), removes TRANSIT_META[1], and drops the PoolSlot (munmapping B's own shmem) — all against its own still-live output pool.
6. Downstream node C holds an IPC handle to B's now-freed GPU buffer → reads freed-then-reallocated GPU memory (silent data corruption / use-after-free). On B's next write_tensor_pool("pool_B_1"), the cache miss reallocates a new GPU buffer, so C's imported handle is permanently stale.

Note that the receiver-side caches in the same function (RECV_GPU_VA, RECV_CPU_SHMEM, GPU_BUF_SIZES, lines ~2585-2620) are correctly keyed by the full namespaced buffer_id, so they target pool_A_1 correctly. It is specifically the sender-side PINNED_POOL/TRANSIT_META block that mis-fires on the wrong pool.

Suggested fix

Mirror the guard already used in process_pending_tensor_pool_frees: only touch PINNED_POOL/TRANSIT_META/GPU buffers when the pool's owner segment equals self.node_id. Extract (owner, counter) from buffer_id, require owner == self.node_id.as_ref(), and only then remove(&counter).

Severity

tensor-pool is an opt-in extension outside the 1.0 compatibility guarantees (#3152), so this only affects builds/wheels compiled with the tensor-pool feature. Within that feature it is a real, reachable data-corruption bug in the middle-node topology of any multi-stage pipeline that frees received pools.

Confidence

Medium-high. The defect is verified by direct code reading; the divergence from the sibling path (whose own comment documents exactly this "cross-process counter aliasing" hazard) is strong corroboration. It requires the freeing node to also own a pool at a colliding counter — which is the common case, since counters restart at 1 per process.