Replace OwningResourceAdaptor with a BackRef mixin

cpp/benchmarks/utils/rmm_utils.hpp

@@ -13,7 +13,6 @@
 #include <rmm/mr/pool_memory_resource.hpp>
 
 #include <rapidsmpf/error.hpp>
-#include <rapidsmpf/rmm_resource_adaptor.hpp>
 
 /**
  * @brief Create and set a RMM memory resource as the current device resource.
@@ -43,15 +42,3 @@ inline void set_current_rmm_resource(std::string const& name) {
         RAPIDSMPF_FAIL("unknown RMM resource name: " + name);
     }
 }
-
-/**
- * @brief Create a statistics-enabled device memory resource wrapping the current
- * device resource, and set it as the current device resource.
- *
- * @return A RmmResourceAdaptor (shared ownership) for accessing statistics.
- */
-[[nodiscard]] inline rapidsmpf::RmmResourceAdaptor set_device_mem_resource_with_stats() {
-    rapidsmpf::RmmResourceAdaptor adaptor{rmm::mr::get_current_device_resource_ref()};
-    rmm::mr::set_current_device_resource(adaptor);
-    return adaptor;
-}

cpp/include/rapidsmpf/memory/back_ref_mixin.hpp

@@ -0,0 +1,125 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <memory>
+#include <utility>
+
+#include <rapidsmpf/error.hpp>
+#include <rapidsmpf/utils/misc.hpp>
+
+namespace rapidsmpf {
+
+/**
+ * @brief Mixin that lets copies of the this object keep an external object reference of
+ * type @p BackRef alive.
+ *
+ * @note Copying an instance without calling `set_backref()` throws `std::bad_weak_ptr`.
+ *
+ * @tparam BackRef Type of the external object reference.
+ */
+template <typename BackRef>
+class BackRefMixin {
+  public:
+    /// @brief Construct a mixin with no installed back-reference.
+    BackRefMixin() noexcept = default;
+
+    /**
+     * @brief Acquire shared ownership of @p other's back-referenced owner.
+     *
+     * After construction, *this* references the same owner as @p other and
+     * keeps it alive for as long as *this* lives.
+     *
+     * @param other Mixin to copy from.
+     * @throws std::bad_weak_ptr if @p other is uninstalled, or is installed
+     * but its owner has been destroyed.
+     */
+    BackRefMixin(BackRefMixin const& other)
+        : weak_{other.weak_},
+          strong_{other.strong_ ? other.strong_ : std::shared_ptr<BackRef>{other.weak_}} {
+    }
+
+    /**
+     * @brief Copy assignment with the same semantics as the copy constructor.
+     *
+     * If the assignment throws, *this* is left unchanged.
+     *
+     * @param other Mixin to copy from.
+     * @return Reference to this.
+     * @throws std::bad_weak_ptr if @p other is uninstalled, or is installed
+     * but its owner has been destroyed.
+     */
+    BackRefMixin& operator=(BackRefMixin const& other) {
+        if (this != &other) {
+            // Promote first so a throw leaves *this unchanged.
+            auto promoted =
+                other.strong_ ? other.strong_ : std::shared_ptr<BackRef>{other.weak_};
+            weak_ = other.weak_;
+            strong_ = std::move(promoted);
+        }
+        return *this;
+    }
+
+    /// @brief Move constructor.
+    BackRefMixin(BackRefMixin&&) noexcept = default;
+
+    /**
+     * @brief Move assignment operator.
+     * @return Reference to this mixin.
+     */
+    BackRefMixin& operator=(BackRefMixin&&) noexcept = default;
+
+    ~BackRefMixin() = default;
+
+    /**
+     * @brief Owner-based equality.
+     *
+     * @param other Mixin to compare against.
+     * @return `true` if both mixins reference the same back-referenced
+     * owner, or are both uninstalled.
+     */
+    [[nodiscard]] bool operator==(BackRefMixin const& other) const noexcept {
+        return owner_equal(weak_, other.weak_);
+    }
+
+    /**
+     * @brief Install a back-reference on this instance.
+     *
+     * After this call, every subsequent copy of *this* will hold shared
+     * ownership of @p backref's referent for the lifetime of the copy.
+     *
+     * Installation only affects *this* and copies made from it afterwards;
+     * copies that already exist are not retroactively back-referenced. To
+     * guarantee the back-reference is honored by every observable copy,
+     * install it before the host object becomes reachable to any code that
+     * may copy it.
+     *
+     * @param backref Non-empty weak reference to the back-referenced owner.
+     * @throws std::invalid_argument if @p backref is empty, or if a
+     * back-reference is already installed on this instance.
+     */
+    void set_backref(std::weak_ptr<BackRef> backref) {
+        std::weak_ptr<BackRef> const empty_weak{};
+        RAPIDSMPF_EXPECTS(
+            !owner_equal(backref, empty_weak),
+            "set_backref: backref must not be empty",
+            std::invalid_argument
+        );
+        RAPIDSMPF_EXPECTS(
+            owner_equal(weak_, empty_weak),
+            "set_backref: backref is already set",
+            std::invalid_argument
+        );
+        weak_ = std::move(backref);
+        strong_.reset();
+    }
+
+  private:
+    std::weak_ptr<BackRef> weak_{};
+    std::shared_ptr<BackRef> strong_{};
+};
+
+}  // namespace rapidsmpf

cpp/include/rapidsmpf/memory/buffer_resource.hpp

@@ -23,7 +23,6 @@
 #include <rapidsmpf/memory/buffer.hpp>
 #include <rapidsmpf/memory/host_memory_resource.hpp>
 #include <rapidsmpf/memory/memory_reservation.hpp>
-#include <rapidsmpf/memory/owning_resource_adaptor.hpp>
 #include <rapidsmpf/memory/pinned_memory_resource.hpp>
 #include <rapidsmpf/memory/resource_types.hpp>
 #include <rapidsmpf/memory/spill_manager.hpp>
@@ -189,9 +188,31 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
      * rmm::device_buffer buf{1024, stream, br->device_mr()};
      * br.reset();  // safe: `buf` keeps the BufferResource alive internally
      * @endcode
+     *
+     * @note Device memory resource provided to the constructor is wrapped in an
+     * `RmmResourceAdaptor` for allocation tracking, and concretely the returned
+     * resource_ref points to that adaptor. See `device_mr_adaptor()` for a more
+     * convenient way to access the adaptor.
      */
     [[nodiscard]] rmm::device_async_resource_ref device_mr() noexcept;
 
+    /**
+     * @brief Access the concrete device memory resource adaptor.
+     *
+     * `BufferResource` wraps the device memory resource in an internal
+     * `RmmResourceAdaptor` for allocation tracking. This exposes that adaptor
+     * directly, e.g. to query allocation statistics via `get_main_record()` or
+     * `current_allocated()`.
+     *
+     * @return Reference to the internal device `RmmResourceAdaptor`. The
+     * reference is valid for as long as this `BufferResource` is alive.
+     *
+     * @note To ensure that the allocations are properly tracked, use `device_mr()` or
+     * `device_mr_adaptor()` instead of the original memory resource passed to the
+     * constructor.
+     */
+    [[nodiscard]] RmmResourceAdaptor& device_mr_adaptor() noexcept;
+
     /**
      * @brief Get the RMM host memory resource.
      *
@@ -495,6 +516,7 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
   private:
     /** @brief Private constructor, use `create()` or `from_options()`. */
     BufferResource(
+        cuda::mr::any_resource<cuda::mr::device_accessible> device_mr,
         std::optional<PinnedMemoryResource> pinned_mr,
         std::unordered_map<MemoryType, std::int64_t> memory_limits,
         std::optional<Duration> periodic_spill_check,
@@ -503,10 +525,7 @@ class BufferResource : public std::enable_shared_from_this<BufferResource> {
     );
 
     std::mutex mutex_;
-    // Stable storage for the device MR returned by `device_mr()`. See
-    // `OwningResourceAdaptor` for the lifetime semantics. Installed by
-    // `create()` after construction.
-    std::optional<OwningResourceAdaptor<RmmResourceAdaptor, BufferResource>> owning_mr_;
+    RmmResourceAdaptor owning_mr_;
     std::optional<PinnedMemoryResource> pinned_mr_;
     HostMemoryResource host_mr_;
     std::array<std::atomic<std::int64_t>, MEMORY_TYPES.size()> memory_limits_;