grpc 1.64.0 → 1.65.0.pre1

data/src/core/lib/resource_quota/arena.cc

@@ -21,17 +21,24 @@
 #include <atomic>
 #include <new>
 
+#include "absl/log/log.h"
+
 #include <grpc/support/alloc.h>
 #include <grpc/support/port_platform.h>
 
-#include "src/core/lib/gpr/alloc.h"
+#include "src/core/lib/resource_quota/resource_quota.h"
+#include "src/core/util/alloc.h"
+namespace grpc_core {
 
 namespace {
 
-void* ArenaStorage(size_t initial_size) {
+void* ArenaStorage(size_t& initial_size) {
   static constexpr size_t base_size =
-      GPR_ROUND_UP_TO_ALIGNMENT_SIZE(sizeof(grpc_core::Arena));
+      GPR_ROUND_UP_TO_ALIGNMENT_SIZE(sizeof(Arena));
   initial_size = GPR_ROUND_UP_TO_ALIGNMENT_SIZE(initial_size);
+  initial_size = std::max(
+      initial_size, GPR_ROUND_UP_TO_ALIGNMENT_SIZE(
+                        arena_detail::BaseArenaContextTraits::ContextSize()));
   size_t alloc_size = base_size + initial_size;
   static constexpr size_t alignment =
       (GPR_CACHELINE_SIZE > GPR_MAX_ALIGNMENT &&
@@ -43,9 +50,15 @@ void* ArenaStorage(size_t initial_size) {
 
 }  // namespace
 
-namespace grpc_core {
-
 Arena::~Arena() {
+  for (size_t i = 0; i < arena_detail::BaseArenaContextTraits::NumContexts();
+       ++i) {
+    arena_detail::BaseArenaContextTraits::Destroy(i, contexts()[i]);
+  }
+  DestroyManagedNewObjects();
+  arena_factory_->FinalizeArena(this);
+  arena_factory_->allocator().Release(
+      total_allocated_.load(std::memory_order_relaxed));
   Zone* z = last_zone_;
   while (z) {
     Zone* prev_z = z->prev;
@@ -53,24 +66,26 @@ Arena::~Arena() {
     gpr_free_aligned(z);
     z = prev_z;
   }
-#ifdef GRPC_ARENA_TRACE_POOLED_ALLOCATIONS
-  gpr_log(GPR_ERROR, "DESTRUCT_ARENA %p", this);
-#endif
 }
 
-Arena* Arena::Create(size_t initial_size, MemoryAllocator* memory_allocator) {
-  return new (ArenaStorage(initial_size))
-      Arena(initial_size, 0, memory_allocator);
+RefCountedPtr<Arena> Arena::Create(size_t initial_size,
+                                   RefCountedPtr<ArenaFactory> arena_factory) {
+  void* p = ArenaStorage(initial_size);
+  return RefCountedPtr<Arena>(
+      new (p) Arena(initial_size, std::move(arena_factory)));
 }
 
-std::pair<Arena*, void*> Arena::CreateWithAlloc(
-    size_t initial_size, size_t alloc_size, MemoryAllocator* memory_allocator) {
-  static constexpr size_t base_size =
-      GPR_ROUND_UP_TO_ALIGNMENT_SIZE(sizeof(Arena));
-  auto* new_arena = new (ArenaStorage(initial_size))
-      Arena(initial_size, alloc_size, memory_allocator);
-  void* first_alloc = reinterpret_cast<char*>(new_arena) + base_size;
-  return std::make_pair(new_arena, first_alloc);
+Arena::Arena(size_t initial_size, RefCountedPtr<ArenaFactory> arena_factory)
+    : initial_zone_size_(initial_size),
+      total_used_(GPR_ROUND_UP_TO_ALIGNMENT_SIZE(
+          arena_detail::BaseArenaContextTraits::ContextSize())),
+      arena_factory_(std::move(arena_factory)) {
+  for (size_t i = 0; i < arena_detail::BaseArenaContextTraits::NumContexts();
+       ++i) {
+    contexts()[i] = nullptr;
+  }
+  CHECK_GE(initial_size, arena_detail::BaseArenaContextTraits::ContextSize());
+  arena_factory_->allocator().Reserve(initial_size);
 }
 
 void Arena::DestroyManagedNewObjects() {
@@ -86,11 +101,9 @@ void Arena::DestroyManagedNewObjects() {
   }
 }
 
-void Arena::Destroy() {
-  DestroyManagedNewObjects();
-  memory_allocator_->Release(total_allocated_.load(std::memory_order_relaxed));
+void Arena::Destroy() const {
   this->~Arena();
-  gpr_free_aligned(this);
+  gpr_free_aligned(const_cast<Arena*>(this));
 }
 
 void* Arena::AllocZone(size_t size) {
@@ -102,7 +115,7 @@ void* Arena::AllocZone(size_t size) {
   static constexpr size_t zone_base_size =
       GPR_ROUND_UP_TO_ALIGNMENT_SIZE(sizeof(Zone));
   size_t alloc_size = zone_base_size + size;
-  memory_allocator_->Reserve(alloc_size);
+  arena_factory_->allocator().Reserve(alloc_size);
   total_allocated_.fetch_add(alloc_size, std::memory_order_relaxed);
   Zone* z = new (gpr_malloc_aligned(alloc_size, GPR_MAX_ALIGNMENT)) Zone();
   auto* prev = last_zone_.load(std::memory_order_relaxed);
@@ -120,63 +133,27 @@ void Arena::ManagedNewObject::Link(std::atomic<ManagedNewObject*>* head) {
   }
 }
 
-#ifndef GRPC_ARENA_POOLED_ALLOCATIONS_USE_MALLOC
-void* Arena::AllocPooled(size_t obj_size, size_t alloc_size,
-                         std::atomic<FreePoolNode*>* head) {
-  // ABA mitigation:
-  // AllocPooled may be called by multiple threads, and to remove a node from
-  // the free list we need to manipulate the next pointer, which may be done
-  // differently by each thread in a naive implementation.
-  // The literature contains various ways of dealing with this. Here we expect
-  // to be mostly single threaded - Arena's are owned by calls and calls don't
-  // do a lot of concurrent work with the pooled allocator. The place that they
-  // do is allocating metadata batches for decoding HPACK headers in chttp2.
-  // So we adopt an approach that is simple and fast for the single threaded
-  // case, and that is also correct in the multi threaded case.
-
-  // First, take ownership of the entire free list. At this point we know that
-  // no other thread can see free nodes and will be forced to allocate.
-  // We think we're mostly single threaded and so that's ok.
-  FreePoolNode* p = head->exchange(nullptr, std::memory_order_acquire);
-  // If there are no nodes in the free list, then go ahead and allocate from the
-  // arena.
-  if (p == nullptr) {
-    void* r = Alloc(alloc_size);
-    TracePoolAlloc(obj_size, r);
-    return r;
-  }
-  // We had a non-empty free list... but we own the *entire* free list.
-  // We only want one node, so if there are extras we'd better give them back.
-  if (p->next != nullptr) {
-    // We perform an exchange to do so, but if there were concurrent frees with
-    // this allocation then there'll be a free list that needs to be merged with
-    // ours.
-    FreePoolNode* extra = head->exchange(p->next, std::memory_order_acq_rel);
-    // If there was a free list concurrently created, we merge it into the
-    // overall free list here by simply freeing each node in turn. This is O(n),
-    // but only O(n) in the number of nodes that were freed concurrently, and
-    // again: we think real world use cases are going to see this as mostly
-    // single threaded.
-    while (extra != nullptr) {
-      FreePoolNode* next = extra->next;
-      FreePooled(extra, head);
-      extra = next;
+RefCountedPtr<ArenaFactory> SimpleArenaAllocator(size_t initial_size) {
+  class Allocator : public ArenaFactory {
+   public:
+    explicit Allocator(size_t initial_size)
+        : ArenaFactory(
+              ResourceQuota::Default()->memory_quota()->CreateMemoryAllocator(
+                  "simple-arena-allocator")),
+          initial_size_(initial_size) {}
+
+    RefCountedPtr<Arena> MakeArena() override {
+      return Arena::Create(initial_size_, Ref());
     }
-  }
-  TracePoolAlloc(obj_size, p);
-  return p;
-}
 
-void Arena::FreePooled(void* p, std::atomic<FreePoolNode*>* head) {
-  // May spuriously trace a free of an already freed object - see AllocPooled
-  // ABA mitigation.
-  TracePoolFree(p);
-  FreePoolNode* node = static_cast<FreePoolNode*>(p);
-  node->next = head->load(std::memory_order_acquire);
-  while (!head->compare_exchange_weak(
-      node->next, node, std::memory_order_acq_rel, std::memory_order_relaxed)) {
-  }
+    void FinalizeArena(Arena*) override {
+      // No-op.
+    }
+
+   private:
+    size_t initial_size_;
+  };
+  return MakeRefCounted<Allocator>(initial_size);
 }
-#endif
 
 }  // namespace grpc_core

data/src/core/lib/resource_quota/arena.h

@@ -35,81 +35,73 @@
 #include <grpc/event_engine/memory_allocator.h>
 #include <grpc/support/port_platform.h>
 
-#include "src/core/lib/gpr/alloc.h"
 #include "src/core/lib/gprpp/construct_destruct.h"
 #include "src/core/lib/promise/context.h"
 #include "src/core/lib/resource_quota/memory_quota.h"
-
-#define GRPC_ARENA_POOLED_ALLOCATIONS_USE_MALLOC
-// #define GRPC_ARENA_TRACE_POOLED_ALLOCATIONS
+#include "src/core/util/alloc.h"
 
 namespace grpc_core {
 
-namespace arena_detail {
+class Arena;
 
-#ifndef GRPC_ARENA_POOLED_ALLOCATIONS_USE_MALLOC
-struct PoolAndSize {
-  size_t alloc_size;
-  size_t pool_index;
-};
+template <typename T>
+struct ArenaContextType;
 
-template <typename Void, size_t kIndex, size_t kObjectSize,
-          size_t... kBucketSize>
-struct PoolIndexForSize;
-
-template <size_t kObjectSize, size_t kIndex, size_t kSmallestRemainingBucket,
-          size_t... kBucketSizes>
-struct PoolIndexForSize<
-    absl::enable_if_t<kObjectSize <= kSmallestRemainingBucket>, kIndex,
-    kObjectSize, kSmallestRemainingBucket, kBucketSizes...> {
-  static constexpr size_t kPool = kIndex;
-  static constexpr size_t kSize = kSmallestRemainingBucket;
-};
+namespace arena_detail {
 
-template <size_t kObjectSize, size_t kIndex, size_t kSmallestRemainingBucket,
-          size_t... kBucketSizes>
-struct PoolIndexForSize<
-    absl::enable_if_t<(kObjectSize > kSmallestRemainingBucket)>, kIndex,
-    kObjectSize, kSmallestRemainingBucket, kBucketSizes...>
-    : public PoolIndexForSize<void, kIndex + 1, kObjectSize, kBucketSizes...> {
-};
+// Tracks all registered arena context types (these should only be registered
+// via ArenaContextTraits at static initialization time).
+class BaseArenaContextTraits {
+ public:
+  // Count of number of contexts that have been allocated.
+  static uint16_t NumContexts() {
+    return static_cast<uint16_t>(RegisteredTraits().size());
+  }
 
-template <size_t kObjectSize, size_t... kBucketSizes>
-constexpr size_t PoolFromObjectSize(
-    absl::integer_sequence<size_t, kBucketSizes...>) {
-  return PoolIndexForSize<void, 0, kObjectSize, kBucketSizes...>::kPool;
-}
+  // Number of bytes required to store the context pointers on an arena.
+  static size_t ContextSize() { return NumContexts() * sizeof(void*); }
 
-template <size_t kObjectSize, size_t... kBucketSizes>
-constexpr size_t AllocationSizeFromObjectSize(
-    absl::integer_sequence<size_t, kBucketSizes...>) {
-  return PoolIndexForSize<void, 0, kObjectSize, kBucketSizes...>::kSize;
-}
+  // Call the registered destruction function for a context.
+  static void Destroy(uint16_t id, void* ptr) {
+    if (ptr == nullptr) return;
+    RegisteredTraits()[id](ptr);
+  }
 
-template <size_t kIndex, size_t... kBucketSizes>
-struct ChoosePoolForAllocationSizeImpl;
+ protected:
+  // Allocate a new context id and register the destruction function.
+  static uint16_t MakeId(void (*destroy)(void* ptr)) {
+    auto& traits = RegisteredTraits();
+    const uint16_t id = static_cast<uint16_t>(traits.size());
+    traits.push_back(destroy);
+    return id;
+  }
 
-template <size_t kIndex, size_t kBucketSize, size_t... kBucketSizes>
-struct ChoosePoolForAllocationSizeImpl<kIndex, kBucketSize, kBucketSizes...> {
-  static PoolAndSize Fn(size_t n) {
-    if (n <= kBucketSize) return {kBucketSize, kIndex};
-    return ChoosePoolForAllocationSizeImpl<kIndex + 1, kBucketSizes...>::Fn(n);
+ private:
+  static std::vector<void (*)(void*)>& RegisteredTraits() {
+    static NoDestruct<std::vector<void (*)(void*)>> registered_traits;
+    return *registered_traits;
   }
 };
 
-template <size_t kIndex>
-struct ChoosePoolForAllocationSizeImpl<kIndex> {
-  static PoolAndSize Fn(size_t n) {
-    return PoolAndSize{n, std::numeric_limits<size_t>::max()};
-  }
+// Traits for a specific context type.
+template <typename T>
+class ArenaContextTraits : public BaseArenaContextTraits {
+ public:
+  static uint16_t id() { return id_; }
+
+ private:
+  static const uint16_t id_;
 };
 
-template <size_t... kBucketSizes>
-PoolAndSize ChoosePoolForAllocationSize(
-    size_t n, absl::integer_sequence<size_t, kBucketSizes...>) {
-  return ChoosePoolForAllocationSizeImpl<0, kBucketSizes...>::Fn(n);
+template <typename T>
+void DestroyArenaContext(void* p) {
+  ArenaContextType<T>::Destroy(static_cast<T*>(p));
 }
-#else
+
+template <typename T>
+const uint16_t ArenaContextTraits<T>::id_ =
+    BaseArenaContextTraits::MakeId(DestroyArenaContext<T>);
+
 template <typename T, typename A, typename B>
 struct IfArray {
   using Result = A;
@@ -119,30 +111,36 @@ template <typename T, typename A, typename B>
 struct IfArray<T[], A, B> {
   using Result = B;
 };
-#endif
+
+struct UnrefDestroy {
+  void operator()(const Arena* arena) const;
+};
 
 }  // namespace arena_detail
 
-class Arena {
-#ifndef GRPC_ARENA_POOLED_ALLOCATIONS_USE_MALLOC
-  // Selected pool sizes.
-  // How to tune: see tools/codegen/core/optimize_arena_pool_sizes.py
-  using PoolSizes = absl::integer_sequence<size_t, 80, 304, 528, 1024>;
-  struct FreePoolNode {
-    FreePoolNode* next;
-  };
-#endif
+class ArenaFactory : public RefCounted<ArenaFactory> {
+ public:
+  virtual RefCountedPtr<Arena> MakeArena() = 0;
+  virtual void FinalizeArena(Arena* arena) = 0;
+
+  MemoryAllocator& allocator() { return allocator_; }
 
+ protected:
+  explicit ArenaFactory(MemoryAllocator allocator)
+      : allocator_(std::move(allocator)) {}
+
+ private:
+  MemoryAllocator allocator_;
+};
+
+RefCountedPtr<ArenaFactory> SimpleArenaAllocator(size_t initial_size = 1024);
+
+class Arena final : public RefCounted<Arena, NonPolymorphicRefCount,
+                                      arena_detail::UnrefDestroy> {
  public:
   // Create an arena, with \a initial_size bytes in the first allocated buffer.
-  static Arena* Create(size_t initial_size, MemoryAllocator* memory_allocator);
-
-  // Create an arena, with \a initial_size bytes in the first allocated buffer,
-  // and return both a void pointer to the returned arena and a void* with the
-  // first allocation.
-  static std::pair<Arena*, void*> CreateWithAlloc(
-      size_t initial_size, size_t alloc_size,
-      MemoryAllocator* memory_allocator);
+  static RefCountedPtr<Arena> Create(size_t initial_size,
+                                     RefCountedPtr<ArenaFactory> arena_factory);
 
   // Destroy all `ManagedNew` allocated objects.
   // Allows safe destruction of these objects even if they need context held by
@@ -151,9 +149,6 @@ class Arena {
   // TODO(ctiller): eliminate ManagedNew.
   void DestroyManagedNewObjects();
 
-  // Destroy an arena.
-  void Destroy();
-
   // Return the total amount of memory allocated by this arena.
   size_t TotalUsedBytes() const {
     return total_used_.load(std::memory_order_relaxed);
@@ -194,95 +189,6 @@ class Arena {
     return &p->t;
   }
 
-#ifndef GRPC_ARENA_POOLED_ALLOCATIONS_USE_MALLOC
-  class PooledDeleter {
-   public:
-    explicit PooledDeleter(std::atomic<FreePoolNode*>* free_list)
-        : free_list_(free_list) {}
-    PooledDeleter() = default;
-    template <typename T>
-    void operator()(T* p) {
-      // TODO(ctiller): promise based filter hijacks ownership of some pointers
-      // to make them appear as PoolPtr without really transferring ownership,
-      // by setting the arena to nullptr.
-      // This is a transitional hack and should be removed once promise based
-      // filter is removed.
-      if (free_list_ != nullptr) {
-        p->~T();
-        FreePooled(p, free_list_);
-      }
-    }
-
-    bool has_freelist() const { return free_list_ != nullptr; }
-
-   private:
-    std::atomic<FreePoolNode*>* free_list_;
-  };
-
-  template <typename T>
-  using PoolPtr = std::unique_ptr<T, PooledDeleter>;
-
-  // Make a unique_ptr to T that is allocated from the arena.
-  // When the pointer is released, the memory may be reused for other
-  // MakePooled(.*) calls.
-  // CAUTION: The amount of memory allocated is rounded up to the nearest
-  //          value in Arena::PoolSizes, and so this may pessimize total
-  //          arena size.
-  template <typename T, typename... Args>
-  PoolPtr<T> MakePooled(Args&&... args) {
-    auto* free_list =
-        &pools_[arena_detail::PoolFromObjectSize<sizeof(T)>(PoolSizes())];
-    return PoolPtr<T>(
-        new (AllocPooled(
-            sizeof(T),
-            arena_detail::AllocationSizeFromObjectSize<sizeof(T)>(PoolSizes()),
-            free_list)) T(std::forward<Args>(args)...),
-        PooledDeleter(free_list));
-  }
-
-  // Make a unique_ptr to an array of T that is allocated from the arena.
-  // When the pointer is released, the memory may be reused for other
-  // MakePooled(.*) calls.
-  // One can use MakePooledArray<char> to allocate a buffer of bytes.
-  // CAUTION: The amount of memory allocated is rounded up to the nearest
-  //          value in Arena::PoolSizes, and so this may pessimize total
-  //          arena size.
-  template <typename T>
-  PoolPtr<T[]> MakePooledArray(size_t n) {
-    auto where =
-        arena_detail::ChoosePoolForAllocationSize(n * sizeof(T), PoolSizes());
-    if (where.pool_index == std::numeric_limits<size_t>::max()) {
-      return PoolPtr<T[]>(new (Alloc(where.alloc_size)) T[n],
-                          PooledDeleter(nullptr));
-    } else {
-      return PoolPtr<T[]>(new (AllocPooled(where.alloc_size, where.alloc_size,
-                                           &pools_[where.pool_index])) T[n],
-                          PooledDeleter(&pools_[where.pool_index]));
-    }
-  }
-
-  // Like MakePooled, but with manual memory management.
-  // The caller is responsible for calling DeletePooled() on the returned
-  // pointer, and expected to call it with the same type T as was passed to this
-  // function (else the free list returned to the arena will be corrupted).
-  template <typename T, typename... Args>
-  T* NewPooled(Args&&... args) {
-    auto* free_list =
-        &pools_[arena_detail::PoolFromObjectSize<sizeof(T)>(PoolSizes())];
-    return new (AllocPooled(
-        sizeof(T),
-        arena_detail::AllocationSizeFromObjectSize<sizeof(T)>(PoolSizes()),
-        free_list)) T(std::forward<Args>(args)...);
-  }
-
-  template <typename T>
-  void DeletePooled(T* p) {
-    auto* free_list =
-        &pools_[arena_detail::PoolFromObjectSize<sizeof(T)>(PoolSizes())];
-    p->~T();
-    FreePooled(p, free_list);
-  }
-#else
   class PooledDeleter {
    public:
     PooledDeleter() = default;
@@ -364,9 +270,30 @@ class Arena {
   void DeletePooled(T* p) {
     delete p;
   }
-#endif
+
+  // Context accessors
+  // Prefer to use the free-standing `GetContext<>` and `SetContext<>` functions
+  // for modern promise-based code -- however legacy filter stack based code
+  // often needs to access these directly.
+  template <typename T>
+  T* GetContext() {
+    return static_cast<T*>(
+        contexts()[arena_detail::ArenaContextTraits<T>::id()]);
+  }
+
+  template <typename T>
+  void SetContext(T* context) {
+    void*& slot = contexts()[arena_detail::ArenaContextTraits<T>::id()];
+    if (slot != nullptr) {
+      ArenaContextType<T>::Destroy(static_cast<T*>(slot));
+    }
+    slot = context;
+    DCHECK_EQ(GetContext<T>(), context);
+  }
 
  private:
+  friend struct arena_detail::UnrefDestroy;
+
   struct Zone {
     Zone* prev;
   };
@@ -396,41 +323,20 @@ class Arena {
   //   quick optimization (avoiding an atomic fetch-add) for the common case
   //   where we wish to create an arena and then perform an immediate
   //   allocation.
-  explicit Arena(size_t initial_size, size_t initial_alloc,
-                 MemoryAllocator* memory_allocator)
-      : total_used_(GPR_ROUND_UP_TO_ALIGNMENT_SIZE(initial_alloc)),
-        initial_zone_size_(initial_size),
-        memory_allocator_(memory_allocator) {}
+  explicit Arena(size_t initial_size,
+                 RefCountedPtr<ArenaFactory> arena_factory);
 
   ~Arena();
 
   void* AllocZone(size_t size);
-
-#ifndef GRPC_ARENA_POOLED_ALLOCATIONS_USE_MALLOC
-  void* AllocPooled(size_t obj_size, size_t alloc_size,
-                    std::atomic<FreePoolNode*>* head);
-  static void FreePooled(void* p, std::atomic<FreePoolNode*>* head);
-#endif
-
-  void TracePoolAlloc(size_t size, void* ptr) {
-    (void)size;
-    (void)ptr;
-#ifdef GRPC_ARENA_TRACE_POOLED_ALLOCATIONS
-    gpr_log(GPR_ERROR, "ARENA %p ALLOC %" PRIdPTR " @ %p", this, size, ptr);
-#endif
-  }
-  static void TracePoolFree(void* ptr) {
-    (void)ptr;
-#ifdef GRPC_ARENA_TRACE_POOLED_ALLOCATIONS
-    gpr_log(GPR_ERROR, "FREE %p", ptr);
-#endif
-  }
+  void Destroy() const;
+  void** contexts() { return reinterpret_cast<void**>(this + 1); }
 
   // Keep track of the total used size. We use this in our call sizing
   // hysteresis.
-  std::atomic<size_t> total_used_{0};
-  std::atomic<size_t> total_allocated_{0};
   const size_t initial_zone_size_;
+  std::atomic<size_t> total_used_;
+  std::atomic<size_t> total_allocated_{initial_zone_size_};
   // If the initial arena allocation wasn't enough, we allocate additional zones
   // in a reverse linked list. Each additional zone consists of (1) a pointer to
   // the zone added before this zone (null if this is the first additional zone)
@@ -438,27 +344,30 @@ class Arena {
   // last zone; the zone list is reverse-walked during arena destruction only.
   std::atomic<Zone*> last_zone_{nullptr};
   std::atomic<ManagedNewObject*> managed_new_head_{nullptr};
-#ifndef GRPC_ARENA_POOLED_ALLOCATIONS_USE_MALLOC
-  std::atomic<FreePoolNode*> pools_[PoolSizes::size()]{};
-#endif
-  // The backing memory quota
-  MemoryAllocator* const memory_allocator_;
-};
-
-// Smart pointer for arenas when the final size is not required.
-struct ScopedArenaDeleter {
-  void operator()(Arena* arena) { arena->Destroy(); }
+  RefCountedPtr<ArenaFactory> arena_factory_;
 };
-using ScopedArenaPtr = std::unique_ptr<Arena, ScopedArenaDeleter>;
-inline ScopedArenaPtr MakeScopedArena(size_t initial_size,
-                                      MemoryAllocator* memory_allocator) {
-  return ScopedArenaPtr(Arena::Create(initial_size, memory_allocator));
-}
 
 // Arenas form a context for activities
 template <>
 struct ContextType<Arena> {};
 
+namespace arena_detail {
+inline void UnrefDestroy::operator()(const Arena* arena) const {
+  arena->Destroy();
+}
+}  // namespace arena_detail
+
+namespace promise_detail {
+
+template <typename T>
+class Context<T, absl::void_t<decltype(ArenaContextType<T>::Destroy)>> {
+ public:
+  static T* get() { return GetContext<Arena>()->GetContext<T>(); }
+  static void set(T* value) { GetContext<Arena>()->SetContext(value); }
+};
+
+}  // namespace promise_detail
+
 }  // namespace grpc_core
 
 #endif  // GRPC_SRC_CORE_LIB_RESOURCE_QUOTA_ARENA_H