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()