fix(issue-2394): key arena per-stream state by CUDA stream ID

The arena memory resource keyed per-stream arenas by the raw
cudaStream_t handle. CUDA can hand a destroyed stream's handle value
to a new stream (ABA reuse), so a new logical stream could inherit
stale arena state. Key stream_arenas_ by cudaStreamGetId instead,
mirroring the pattern already used by stream_ordered_memory_resource.

Closes #2394

Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>

Author: nethum529

cpp/include/rmm/mr/detail/arena_memory_resource_impl.hpp

@@ -22,6 +22,27 @@ namespace RMM_NAMESPACE {
 namespace mr {
 namespace detail {
 
+/// The type used to key per-stream arenas by CUDA stream ID.
+using stream_id_type = unsigned long long;
+
+/**
+ * @brief Return the CUDA stream ID used to key per-stream arenas.
+ *
+ * Normalizes the default stream to `cudaStreamLegacy` for consistency between
+ * PTDS and non-PTDS mode, then returns the stream's ID from `cudaStreamGetId`.
+ * Keying by ID (rather than the raw `cudaStream_t` handle) prevents a newly
+ * created stream that reuses a destroyed stream's handle value from inheriting
+ * stale arena state (ABA reuse).
+ *
+ * @param stream The stream for which to get the ID.
+ * @param may_throw If `true`, a `cudaStreamGetId` failure throws `rmm::cuda_error` (only safe to
+ * pass from a context that is not `noexcept`, e.g. the `allocate` path). If `false`, failure is
+ * only asserted in debug builds and silently ignored in release builds, matching the
+ * `noexcept`-safe requirements of the `deallocate`/`deallocate_sync` path.
+ * @return The CUDA stream ID.
+ */
+stream_id_type get_stream_id(cuda_stream_view stream, bool may_throw);
+
 /**
  * @brief Implementation class for arena_memory_resource.
  *
@@ -81,17 +102,17 @@ class arena_memory_resource_impl {
 
   void deallocate_from_other_arena(cuda_stream_view stream, void* ptr, std::size_t bytes);
 
-  arena& get_arena(cuda_stream_view stream);
+  arena& get_arena(cuda_stream_view stream, bool may_throw);
   arena& get_thread_arena();
-  arena& get_stream_arena(cuda_stream_view stream);
+  arena& get_stream_arena(cuda_stream_view stream, bool may_throw);
 
   void dump_memory_log(std::size_t bytes);
 
   static bool use_per_thread_arena(cuda_stream_view stream);
 
   global_arena global_arena_;
   std::map<std::thread::id, std::shared_ptr<arena>> thread_arenas_;
-  std::map<cudaStream_t, arena> stream_arenas_;
+  std::map<stream_id_type, arena> stream_arenas_;
   bool dump_log_on_failure_{};
   std::shared_ptr<rapids_logger::logger> logger_{};
   mutable std::shared_mutex map_mtx_;

cpp/src/mr/detail/arena_memory_resource_impl.cpp

@@ -15,6 +15,26 @@ namespace RMM_NAMESPACE {
 namespace mr {
 namespace detail {
 
+stream_id_type get_stream_id(cuda_stream_view stream, bool may_throw)
+{
+  // We use cudaStreamLegacy as the arena map key for the default stream for consistency between
+  // PTDS and non-PTDS mode. Keying by the stream ID (rather than the raw cudaStream_t handle)
+  // prevents a newly created stream that reuses a destroyed stream's handle value from inheriting
+  // stale arena state (ABA reuse).
+  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-cstyle-cast)
+  auto* const stream_to_store = stream.is_default() ? cudaStreamLegacy : stream.value();
+  stream_id_type stream_id{};
+  if (may_throw) {
+    // Safe to throw: only called from the non-noexcept allocate() path.
+    RMM_CUDA_TRY(cudaStreamGetId(stream_to_store, &stream_id));
+  } else {
+    // Called from the noexcept deallocate()/deallocate_sync() path: must not throw. A failure
+    // here is only surfaced via assertion in debug builds.
+    RMM_ASSERT_CUDA_SUCCESS(cudaStreamGetId(stream_to_store, &stream_id));
+  }
+  return stream_id;
+}
+
 arena_memory_resource_impl::arena_memory_resource_impl(
   cuda::mr::any_resource<cuda::mr::device_accessible> upstream_mr,
   std::optional<std::size_t> arena_size,
@@ -39,7 +59,7 @@ void* arena_memory_resource_impl::allocate(cuda::stream_ref stream,
 #else
   bytes = rmm::align_up(bytes, rmm::CUDA_ALLOCATION_ALIGNMENT);
 #endif
-  auto& arena = get_arena(sv);
+  auto& arena = get_arena(sv, /*may_throw=*/true);
 
   {
     std::shared_lock lock(mtx_);
@@ -73,7 +93,7 @@ void arena_memory_resource_impl::deallocate(cuda::stream_ref stream,
 #else
   bytes = rmm::align_up(bytes, rmm::CUDA_ALLOCATION_ALIGNMENT);
 #endif
-  auto& arena = get_arena(sv);
+  auto& arena = get_arena(sv, /*may_throw=*/false);
 
   {
     std::shared_lock lock(mtx_);
@@ -138,10 +158,11 @@ void arena_memory_resource_impl::deallocate_from_other_arena(cuda_stream_view st
   }
 }
 
-arena_memory_resource_impl::arena& arena_memory_resource_impl::get_arena(cuda_stream_view stream)
+arena_memory_resource_impl::arena& arena_memory_resource_impl::get_arena(cuda_stream_view stream,
+                                                                         bool may_throw)
 {
   if (use_per_thread_arena(stream)) { return get_thread_arena(); }
-  return get_stream_arena(stream);
+  return get_stream_arena(stream, may_throw);
 }
 
 arena_memory_resource_impl::arena& arena_memory_resource_impl::get_thread_arena()
@@ -162,18 +183,19 @@ arena_memory_resource_impl::arena& arena_memory_resource_impl::get_thread_arena(
 }
 
 arena_memory_resource_impl::arena& arena_memory_resource_impl::get_stream_arena(
-  cuda_stream_view stream)
+  cuda_stream_view stream, bool may_throw)
 {
   RMM_LOGGING_ASSERT(!use_per_thread_arena(stream));
+  auto const stream_id = get_stream_id(stream, may_throw);
   {
     std::shared_lock lock(map_mtx_);
-    auto const iter = stream_arenas_.find(stream.value());
+    auto const iter = stream_arenas_.find(stream_id);
     if (iter != stream_arenas_.end()) { return iter->second; }
   }
   {
     std::unique_lock lock(map_mtx_);
-    stream_arenas_.emplace(stream.value(), global_arena_);
-    return stream_arenas_.at(stream.value());
+    stream_arenas_.emplace(stream_id, global_arena_);
+    return stream_arenas_.at(stream_id);
   }
 }
 

cpp/tests/mr/arena_mr_tests.cpp

@@ -637,5 +637,68 @@ TEST_F(ArenaTest, DumpLogOnFailure)  // NOLINT
   EXPECT_GE(file_status.st_size, 0);
 }
 
+// Distinct live streams get distinct IDs, and a stream's ID is stable across calls.
+TEST(ArenaStreamIdTest, DistinctStreamsHaveDistinctStableIds)  // NOLINT
+{
+  rmm::cuda_stream stream_a{};
+  rmm::cuda_stream stream_b{};
+
+  EXPECT_NE(rmm::mr::detail::get_stream_id(stream_a.view(), /*may_throw=*/true),
+            rmm::mr::detail::get_stream_id(stream_b.view(), /*may_throw=*/true));
+  EXPECT_EQ(rmm::mr::detail::get_stream_id(stream_a.view(), /*may_throw=*/true),
+            rmm::mr::detail::get_stream_id(stream_a.view(), /*may_throw=*/true));
+}
+
+// The default stream is normalized via cudaStreamLegacy, so it maps to one consistent key.
+TEST(ArenaStreamIdTest, DefaultStreamNormalizesToLegacyId)  // NOLINT
+{
+  auto const id = rmm::mr::detail::get_stream_id(rmm::cuda_stream_view{}, /*may_throw=*/true);
+  EXPECT_EQ(id, rmm::mr::detail::get_stream_id(rmm::cuda_stream_view{}, /*may_throw=*/true));
+
+#ifndef CUDA_API_PER_THREAD_DEFAULT_STREAM
+  // Non-PTDS: the default stream normalizes to cudaStreamLegacy, so it keys to the legacy id.
+  EXPECT_EQ(rmm::mr::detail::get_stream_id(rmm::cuda_stream_view{}, /*may_throw=*/true),
+            rmm::mr::detail::get_stream_id(rmm::cuda_stream_legacy, /*may_throw=*/true));
+#endif
+
+  // In both modes, a real created stream must never collide with the legacy key. This catches a
+  // regression that collapses every stream to a single key.
+  rmm::cuda_stream stream{};
+  EXPECT_NE(rmm::mr::detail::get_stream_id(stream.view(), /*may_throw=*/true),
+            rmm::mr::detail::get_stream_id(rmm::cuda_stream_legacy, /*may_throw=*/true));
+}
+
+// Regression for #2394: CUDA can hand a destroyed stream's raw handle value to a new stream (ABA
+// reuse). Keying arenas by the raw handle would let the new stream inherit stale arena state.
+// Keying by cudaStreamGetId prevents this: a reused handle yields a distinct ID.
+//
+// This test cannot force reuse deterministically because CUDA does not guarantee handle-reuse
+// timing (see the issue's own caveat), so it probes up to a bounded number of iterations and
+// skips if reuse is never observed. It either proves the property or skips, never falsely fails.
+TEST(ArenaStreamIdTest, ReusedRawHandleYieldsDistinctId)  // NOLINT
+{
+  cudaStream_t raw{};
+  rmm::mr::detail::stream_id_type id{};
+  {
+    rmm::cuda_stream original{};
+    raw = original.view().value();
+    id  = rmm::mr::detail::get_stream_id(original.view(), /*may_throw=*/true);
+  }  // original destroyed here, freeing its raw handle for potential reuse
+
+  constexpr int max_iterations = 2048;
+  bool observed_reuse          = false;
+  for (int i = 0; i < max_iterations; ++i) {
+    rmm::cuda_stream candidate{};
+    if (candidate.view().value() == raw) {
+      // Same raw handle, but the ID must differ so no stale arena is inherited.
+      EXPECT_NE(rmm::mr::detail::get_stream_id(candidate.view(), /*may_throw=*/true), id);
+      observed_reuse = true;
+      break;
+    }
+  }
+
+  if (!observed_reuse) { GTEST_SKIP() << "CUDA did not reuse the stream handle within the bound"; }
+}
+
 }  // namespace
 }  // namespace rmm::test