Make promise arena part of the context (#28313)

* Make promise arena part of the context

* Automated change: Fix sanity tests

Co-authored-by: ctiller <ctiller@users.noreply.github.com>

BUILD

@@ -1509,6 +1509,7 @@ grpc_cc_library(
         "src/core/lib/resource_quota/arena.h",
     ],
     deps = [
+        "context",
         "gpr_base",
         "memory_quota",
     ],

src/core/lib/promise/arena_promise.h

@@ -118,8 +118,8 @@ struct ChooseImplForCallable;
 template <typename T, typename Callable>
 struct ChooseImplForCallable<
     T, Callable, absl::enable_if_t<!std::is_empty<Callable>::value>> {
-  static ImplInterface<T>* Make(Arena* arena, Callable&& callable) {
-    return arena->template New<CallableImpl<T, Callable>>(
+  static ImplInterface<T>* Make(Callable&& callable) {
+    return GetContext<Arena>()->template New<CallableImpl<T, Callable>>(
         std::forward<Callable>(callable));
   }
 };
@@ -127,16 +127,16 @@ struct ChooseImplForCallable<
 template <typename T, typename Callable>
 struct ChooseImplForCallable<
     T, Callable, absl::enable_if_t<std::is_empty<Callable>::value>> {
-  static ImplInterface<T>* Make(Arena*, Callable&& callable) {
+  static ImplInterface<T>* Make(Callable&& callable) {
     return SharedImpl<T, Callable>::Get(std::forward<Callable>(callable));
   }
 };
 
 // Wrap ChooseImplForCallable with a friend approachable syntax.
 template <typename T, typename Callable>
-ImplInterface<T>* MakeImplForCallable(Arena* arena, Callable&& callable) {
+ImplInterface<T>* MakeImplForCallable(Callable&& callable) {
   return ChooseImplForCallable<T, Callable>::Make(
-      arena, std::forward<Callable>(callable));
+      std::forward<Callable>(callable));
 }
 
 }  // namespace arena_promise_detail
@@ -149,10 +149,12 @@ class ArenaPromise {
   ArenaPromise() = default;
 
   // Construct an ArenaPromise that will call the given callable when polled.
-  template <typename Callable>
-  ArenaPromise(Arena* arena, Callable&& callable)
+  template <typename Callable,
+            typename Ignored =
+                absl::enable_if_t<!std::is_same<Callable, ArenaPromise>::value>>
+  explicit ArenaPromise(Callable&& callable)
       : impl_(arena_promise_detail::MakeImplForCallable<T>(
-            arena, std::forward<Callable>(callable))) {}
+            std::forward<Callable>(callable))) {}
 
   // ArenaPromise is not copyable.
   ArenaPromise(const ArenaPromise&) = delete;

src/core/lib/resource_quota/arena.h

@@ -38,6 +38,7 @@
 #include <grpc/support/sync.h>
 
 #include "src/core/lib/gpr/alloc.h"
+#include "src/core/lib/promise/context.h"
 #include "src/core/lib/resource_quota/memory_quota.h"
 
 namespace grpc_core {
@@ -132,6 +133,10 @@ inline ScopedArenaPtr MakeScopedArena(size_t initial_size,
   return ScopedArenaPtr(Arena::Create(initial_size, memory_allocator));
 }
 
+// Arenas form a context for activities
+template <>
+struct ContextType<Arena> {};
+
 }  // namespace grpc_core
 
 #endif /* GRPC_CORE_LIB_RESOURCE_QUOTA_ARENA_H */

test/core/promise/arena_promise_test.cc

@@ -27,17 +27,19 @@ static auto* g_memory_allocator = new MemoryAllocator(
 
 TEST(ArenaPromiseTest, AllocatedWorks) {
   auto arena = MakeScopedArena(1024, g_memory_allocator);
+  promise_detail::Context<Arena> context(arena.get());
   int x = 42;
-  ArenaPromise<int> p(arena.get(), [x] { return Poll<int>(x); });
+  ArenaPromise<int> p([x] { return Poll<int>(x); });
   EXPECT_EQ(p(), Poll<int>(42));
-  p = ArenaPromise<int>(arena.get(), [] { return Poll<int>(43); });
+  p = ArenaPromise<int>([] { return Poll<int>(43); });
   EXPECT_EQ(p(), Poll<int>(43));
 }
 
 TEST(ArenaPromiseTest, DestructionWorks) {
   auto arena = MakeScopedArena(1024, g_memory_allocator);
+  promise_detail::Context<Arena> context(arena.get());
   auto x = std::make_shared<int>(42);
-  auto p = ArenaPromise<int>(arena.get(), [x] { return Poll<int>(*x); });
+  auto p = ArenaPromise<int>([x] { return Poll<int>(*x); });
   ArenaPromise<int> q(std::move(p));
   EXPECT_EQ(q(), Poll<int>(42));
 }