From 3562094389a13fbcaa96db0456bfbec581c92b4f Mon Sep 17 00:00:00 2001 From: Craig Tiller Date: Thu, 30 May 2024 15:46:56 -0700 Subject: [PATCH] [arena+context] Begin merging the concepts of context and arenas for calls (#36773) Add a dynamic registration mechanism for a new kind of context: one that lives in an arena. An upcoming set of changes will move all of the legacy context types into this mechanism. It's likely we'll move other promise based context types to this mechanism also, until the only promise-based context type is `Arena`. Closes #36773 COPYBARA_INTEGRATE_REVIEW=https://github.com/grpc/grpc/pull/36773 from ctiller:transport-refs-7.8 e6892031ca32baf54881e47e8561eba85a4a08e6 PiperOrigin-RevId: 638810722 --- src/core/lib/resource_quota/arena.cc | 38 ++++++---- src/core/lib/resource_quota/arena.h | 84 ++++++++++++++++++++- src/core/lib/resource_quota/memory_quota.cc | 6 +- test/core/resource_quota/arena_test.cc | 18 +++++ 4 files changed, 125 insertions(+), 21 deletions(-) diff --git a/src/core/lib/resource_quota/arena.cc b/src/core/lib/resource_quota/arena.cc index 2f9e3c851e5..e30bd1047ca 100644 --- a/src/core/lib/resource_quota/arena.cc +++ b/src/core/lib/resource_quota/arena.cc @@ -21,18 +21,23 @@ #include #include +#include "absl/log/log.h" + #include #include #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, + arena_detail::BaseArenaContextTraits::ContextSize()); size_t alloc_size = base_size + initial_size; static constexpr size_t alignment = (GPR_CACHELINE_SIZE > GPR_MAX_ALIGNMENT && @@ -44,9 +49,11 @@ 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( @@ -58,24 +65,25 @@ Arena::~Arena() { gpr_free_aligned(z); z = prev_z; } -#ifdef GRPC_ARENA_TRACE_POOLED_ALLOCATIONS - gpr_log(GPR_ERROR, "DESTRUCT_ARENA %p", this); -#endif } RefCountedPtr Arena::Create(size_t initial_size, RefCountedPtr arena_factory) { - return RefCountedPtr(new (ArenaStorage(initial_size)) Arena( - initial_size, 0, std::move(arena_factory))); + void* p = ArenaStorage(initial_size); + return RefCountedPtr( + new (p) Arena(initial_size, std::move(arena_factory))); } -Arena::Arena(size_t initial_size, size_t initial_alloc, - RefCountedPtr arena_factory) - : total_used_(GPR_ROUND_UP_TO_ALIGNMENT_SIZE(initial_alloc)), - initial_zone_size_(initial_size), +Arena::Arena(size_t initial_size, RefCountedPtr arena_factory) + : initial_zone_size_(initial_size), + total_used_(arena_detail::BaseArenaContextTraits::ContextSize()), arena_factory_(std::move(arena_factory)) { - arena_factory_->allocator().Reserve( - GPR_ROUND_UP_TO_ALIGNMENT_SIZE(initial_alloc)); + 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() { diff --git a/src/core/lib/resource_quota/arena.h b/src/core/lib/resource_quota/arena.h index 61fc917badf..0b92e4b5b74 100644 --- a/src/core/lib/resource_quota/arena.h +++ b/src/core/lib/resource_quota/arena.h @@ -44,8 +44,59 @@ namespace grpc_core { class Arena; +template +struct ArenaContextType; + namespace arena_detail { +// 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(RegisteredTraits().size()); + } + + // Number of bytes required to store the context pointers on an arena. + static size_t ContextSize() { return NumContexts() * sizeof(void*); } + + // Call the registered destruction function for a context. + static void Destroy(uint16_t id, void* ptr) { + if (ptr == nullptr) return; + RegisteredTraits()[id](ptr); + } + + 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(traits.size()); + traits.push_back(destroy); + return id; + } + + private: + static std::vector& RegisteredTraits() { + static NoDestruct> registered_traits; + return *registered_traits; + } +}; + +// Traits for a specific context type. +template +class ArenaContextTraits : public BaseArenaContextTraits { + public: + static uint16_t id() { return id_; } + + private: + static const uint16_t id_; +}; + +template +const uint16_t ArenaContextTraits::id_ = BaseArenaContextTraits::MakeId( + [](void* ptr) { ArenaContextType::Destroy(static_cast(ptr)); }); + template struct IfArray { using Result = A; @@ -215,6 +266,21 @@ class Arena final : public RefCounted + T* GetContext() { + return static_cast( + contexts()[arena_detail::ArenaContextTraits::id()]); + } + + template + void SetContext(T* context) { + void*& slot = contexts()[arena_detail::ArenaContextTraits::id()]; + if (slot != nullptr) { + ArenaContextType::Destroy(static_cast(slot)); + } + slot = context; + } + private: friend struct arena_detail::UnrefDestroy; @@ -247,19 +313,20 @@ class Arena final : public RefCounted arena_factory); ~Arena(); void* AllocZone(size_t size); void Destroy() const; + void** contexts() { return reinterpret_cast(this + 1); } // Keep track of the total used size. We use this in our call sizing // hysteresis. - std::atomic total_used_{0}; - std::atomic total_allocated_{0}; const size_t initial_zone_size_; + std::atomic total_used_; + std::atomic 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) @@ -280,6 +347,17 @@ inline void UnrefDestroy::operator()(const Arena* arena) const { } } // namespace arena_detail +namespace promise_detail { + +template +class Context::Destroy)>> { + public: + static T* get() { return GetContext()->GetContext(); } + static void set(T* value) { GetContext()->SetContext(value); } +}; + +} // namespace promise_detail + } // namespace grpc_core #endif // GRPC_SRC_CORE_LIB_RESOURCE_QUOTA_ARENA_H diff --git a/src/core/lib/resource_quota/memory_quota.cc b/src/core/lib/resource_quota/memory_quota.cc index 79702ce7fff..458597da5da 100644 --- a/src/core/lib/resource_quota/memory_quota.cc +++ b/src/core/lib/resource_quota/memory_quota.cc @@ -253,9 +253,9 @@ GrpcMemoryAllocatorImpl::GrpcMemoryAllocatorImpl( } GrpcMemoryAllocatorImpl::~GrpcMemoryAllocatorImpl() { - CHECK(free_bytes_.load(std::memory_order_acquire) + - sizeof(GrpcMemoryAllocatorImpl) == - taken_bytes_.load(std::memory_order_relaxed)); + CHECK_EQ(free_bytes_.load(std::memory_order_acquire) + + sizeof(GrpcMemoryAllocatorImpl), + taken_bytes_.load(std::memory_order_relaxed)); memory_quota_->Return(taken_bytes_.load(std::memory_order_relaxed)); } diff --git a/test/core/resource_quota/arena_test.cc b/test/core/resource_quota/arena_test.cc index cd3d3f55f8d..83e9721ec8e 100644 --- a/test/core/resource_quota/arena_test.cc +++ b/test/core/resource_quota/arena_test.cc @@ -274,6 +274,24 @@ TEST(ArenaTest, ConcurrentMakePooled) { } } +struct Foo { + explicit Foo(int x) : p(std::make_unique(x)) {} + std::unique_ptr p; +}; + +template <> +struct ArenaContextType { + static void Destroy(Foo* p) { p->~Foo(); } +}; + +TEST(ArenaTest, FooContext) { + auto arena = SimpleArenaAllocator()->MakeArena(); + EXPECT_EQ(arena->GetContext(), nullptr); + arena->SetContext(arena->New(42)); + ASSERT_NE(arena->GetContext(), nullptr); + EXPECT_EQ(*arena->GetContext()->p, 42); +} + class MockArenaFactory : public ArenaFactory { public: MockArenaFactory()