From bb0e4fdd6f4960d1d5125256dc147f28d83db560 Mon Sep 17 00:00:00 2001
From: Michael Schellenberger Costa <miscco@nvidia.com>
Date: Wed, 6 Mar 2024 06:17:24 -0800
Subject: [PATCH] Add `get_upstream_resource` method to
 `stream_checking_resource_adaptor` (#15203)

Also deprecate `get_upstream` as we want to get away from raw upstreams

Authors:
  - Michael Schellenberger Costa (https://github.com/miscco)

Approvers:
  - Mark Harris (https://github.com/harrism)
  - David Wendt (https://github.com/davidwendt)

URL: https://github.com/rapidsai/cudf/pull/15203
---
 .../cudf_test/stream_checking_resource_adaptor.hpp | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/cpp/include/cudf_test/stream_checking_resource_adaptor.hpp b/cpp/include/cudf_test/stream_checking_resource_adaptor.hpp
index d1841ff42a1..cafde6ca7d5 100644
--- a/cpp/include/cudf_test/stream_checking_resource_adaptor.hpp
+++ b/cpp/include/cudf_test/stream_checking_resource_adaptor.hpp
@@ -20,6 +20,7 @@
 #include <cudf/detail/utilities/stacktrace.hpp>
 
 #include <rmm/mr/device/device_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <iostream>
 
@@ -58,11 +59,14 @@ class stream_checking_resource_adaptor final : public rmm::mr::device_memory_res
     default;
 
   /**
-   * @brief Return pointer to the upstream resource.
+   * @brief Returns the wrapped upstream resource
    *
-   * @return Pointer to the upstream resource.
+   * @return The wrapped upstream resource
    */
-  Upstream* get_upstream() const noexcept { return upstream_; }
+  [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
+  {
+    return upstream_;
+  }
 
  private:
   /**
@@ -110,8 +114,8 @@ class stream_checking_resource_adaptor final : public rmm::mr::device_memory_res
   {
     if (this == &other) { return true; }
     auto cast = dynamic_cast<stream_checking_resource_adaptor<Upstream> const*>(&other);
-    return cast != nullptr ? upstream_->is_equal(*cast->get_upstream())
-                           : upstream_->is_equal(other);
+    if (cast == nullptr) { return upstream_->is_equal(other); }
+    return get_upstream_resource() == cast->get_upstream_resource();
   }
 
   /**