[arena] Add tests around newer arena code (#36933)

Ensure arena accounting is working, and add a test that a constant call size results in a constant call size estimate.

Closes #36933

COPYBARA_INTEGRATE_REVIEW=https://github.com/grpc/grpc/pull/36933 from ctiller:arena-accounts 116c805633
PiperOrigin-RevId: 644102412

CMakeLists.txt

@@ -964,6 +964,9 @@ if(gRPC_BUILD_TESTS)
   add_dependencies(buildtests_cxx buffer_list_test)
   add_dependencies(buildtests_cxx byte_buffer_test)
   add_dependencies(buildtests_cxx c_slice_buffer_test)
+  if(_gRPC_PLATFORM_LINUX OR _gRPC_PLATFORM_MAC OR _gRPC_PLATFORM_POSIX)
+    add_dependencies(buildtests_cxx call_arena_allocator_test)
+  endif()
   add_dependencies(buildtests_cxx call_creds_test)
   add_dependencies(buildtests_cxx call_filters_test)
   add_dependencies(buildtests_cxx call_finalization_test)
@@ -8500,6 +8503,50 @@ target_link_libraries(c_slice_buffer_test
 )
 
 
+endif()
+if(gRPC_BUILD_TESTS)
+if(_gRPC_PLATFORM_LINUX OR _gRPC_PLATFORM_MAC OR _gRPC_PLATFORM_POSIX)
+
+  add_executable(call_arena_allocator_test
+    test/core/transport/call_arena_allocator_test.cc
+  )
+  if(WIN32 AND MSVC)
+    if(BUILD_SHARED_LIBS)
+      target_compile_definitions(call_arena_allocator_test
+      PRIVATE
+        "GPR_DLL_IMPORTS"
+        "GRPC_DLL_IMPORTS"
+      )
+    endif()
+  endif()
+  target_compile_features(call_arena_allocator_test PUBLIC cxx_std_14)
+  target_include_directories(call_arena_allocator_test
+    PRIVATE
+      ${CMAKE_CURRENT_SOURCE_DIR}
+      ${CMAKE_CURRENT_SOURCE_DIR}/include
+      ${_gRPC_ADDRESS_SORTING_INCLUDE_DIR}
+      ${_gRPC_RE2_INCLUDE_DIR}
+      ${_gRPC_SSL_INCLUDE_DIR}
+      ${_gRPC_UPB_GENERATED_DIR}
+      ${_gRPC_UPB_GRPC_GENERATED_DIR}
+      ${_gRPC_UPB_INCLUDE_DIR}
+      ${_gRPC_XXHASH_INCLUDE_DIR}
+      ${_gRPC_ZLIB_INCLUDE_DIR}
+      third_party/googletest/googletest/include
+      third_party/googletest/googletest
+      third_party/googletest/googlemock/include
+      third_party/googletest/googlemock
+      ${_gRPC_PROTO_GENS_DIR}
+  )
+
+  target_link_libraries(call_arena_allocator_test
+    ${_gRPC_ALLTARGETS_LIBRARIES}
+    gtest
+    grpc_test_util
+  )
+
+
+endif()
 endif()
 if(gRPC_BUILD_TESTS)
 

build_autogenerated.yaml

@@ -6331,6 +6331,21 @@ targets:
   - gtest
   - grpc_test_util
   uses_polling: false
+- name: call_arena_allocator_test
+  gtest: true
+  build: test
+  language: c++
+  headers: []
+  src:
+  - test/core/transport/call_arena_allocator_test.cc
+  deps:
+  - gtest
+  - grpc_test_util
+  platforms:
+  - linux
+  - posix
+  - mac
+  uses_polling: false
 - name: call_creds_test
   gtest: true
   build: test

include/grpc/event_engine/memory_request.h

@@ -46,6 +46,24 @@ class MemoryRequest {
   size_t min() const { return min_; }
   size_t max() const { return max_; }
 
+  bool operator==(const MemoryRequest& other) const {
+    return min_ == other.min_ && max_ == other.max_;
+  }
+  bool operator!=(const MemoryRequest& other) const {
+    return !(*this == other);
+  }
+
+  template <typename Sink>
+  friend void AbslStringify(Sink& s, const MemoryRequest& r) {
+    if (r.min_ == r.max_) {
+      s.Append(r.min_);
+    } else {
+      s.Append(r.min_);
+      s.Append("..");
+      s.Append(r.max_);
+    }
+  }
+
  private:
   size_t min_;
   size_t max_;

src/core/lib/resource_quota/arena.cc

@@ -33,19 +33,17 @@ namespace grpc_core {
 namespace {
 
 void* ArenaStorage(size_t& initial_size) {
-  static constexpr size_t base_size =
-      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;
+  size_t base_size = Arena::ArenaOverhead() +
+                     GPR_ROUND_UP_TO_ALIGNMENT_SIZE(
+                         arena_detail::BaseArenaContextTraits::ContextSize());
+  initial_size =
+      std::max(GPR_ROUND_UP_TO_ALIGNMENT_SIZE(initial_size), base_size);
   static constexpr size_t alignment =
       (GPR_CACHELINE_SIZE > GPR_MAX_ALIGNMENT &&
        GPR_CACHELINE_SIZE % GPR_MAX_ALIGNMENT == 0)
           ? GPR_CACHELINE_SIZE
           : GPR_MAX_ALIGNMENT;
-  return gpr_malloc_aligned(alloc_size, alignment);
+  return gpr_malloc_aligned(initial_size, alignment);
 }
 
 }  // namespace
@@ -77,8 +75,9 @@ RefCountedPtr<Arena> Arena::Create(size_t initial_size,
 
 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())),
+      total_used_(ArenaOverhead() +
+                  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) {
@@ -133,14 +132,17 @@ void Arena::ManagedNewObject::Link(std::atomic<ManagedNewObject*>* head) {
   }
 }
 
-RefCountedPtr<ArenaFactory> SimpleArenaAllocator(size_t initial_size) {
+MemoryAllocator DefaultMemoryAllocatorForSimpleArenaAllocator() {
+  return ResourceQuota::Default()->memory_quota()->CreateMemoryAllocator(
+      "simple-arena-allocator");
+}
+
+RefCountedPtr<ArenaFactory> SimpleArenaAllocator(size_t initial_size,
+                                                 MemoryAllocator allocator) {
   class Allocator : public ArenaFactory {
    public:
-    explicit Allocator(size_t initial_size)
-        : ArenaFactory(
-              ResourceQuota::Default()->memory_quota()->CreateMemoryAllocator(
-                  "simple-arena-allocator")),
-          initial_size_(initial_size) {}
+    Allocator(size_t initial_size, MemoryAllocator allocator)
+        : ArenaFactory(std::move(allocator)), initial_size_(initial_size) {}
 
     RefCountedPtr<Arena> MakeArena() override {
       return Arena::Create(initial_size_, Ref());
@@ -153,7 +155,7 @@ RefCountedPtr<ArenaFactory> SimpleArenaAllocator(size_t initial_size) {
    private:
     size_t initial_size_;
   };
-  return MakeRefCounted<Allocator>(initial_size);
+  return MakeRefCounted<Allocator>(initial_size, std::move(allocator));
 }
 
 }  // namespace grpc_core

src/core/lib/resource_quota/arena.h

@@ -133,7 +133,11 @@ class ArenaFactory : public RefCounted<ArenaFactory> {
   MemoryAllocator allocator_;
 };
 
-RefCountedPtr<ArenaFactory> SimpleArenaAllocator(size_t initial_size = 1024);
+MemoryAllocator DefaultMemoryAllocatorForSimpleArenaAllocator();
+RefCountedPtr<ArenaFactory> SimpleArenaAllocator(
+    size_t initial_size = 1024,
+    MemoryAllocator allocator =
+        DefaultMemoryAllocatorForSimpleArenaAllocator());
 
 class Arena final : public RefCounted<Arena, NonPolymorphicRefCount,
                                       arena_detail::UnrefDestroy> {
@@ -156,12 +160,10 @@ class Arena final : public RefCounted<Arena, NonPolymorphicRefCount,
 
   // Allocate \a size bytes from the arena.
   void* Alloc(size_t size) {
-    static constexpr size_t base_size =
-        GPR_ROUND_UP_TO_ALIGNMENT_SIZE(sizeof(Arena));
     size = GPR_ROUND_UP_TO_ALIGNMENT_SIZE(size);
     size_t begin = total_used_.fetch_add(size, std::memory_order_relaxed);
     if (begin + size <= initial_zone_size_) {
-      return reinterpret_cast<char*>(this) + base_size + begin;
+      return reinterpret_cast<char*>(this) + begin;
     } else {
       return AllocZone(size);
     }
@@ -291,6 +293,13 @@ class Arena final : public RefCounted<Arena, NonPolymorphicRefCount,
     DCHECK_EQ(GetContext<T>(), context);
   }
 
+  static size_t ArenaOverhead() {
+    return GPR_ROUND_UP_TO_ALIGNMENT_SIZE(sizeof(Arena));
+  }
+  static size_t ArenaZoneOverhead() {
+    return GPR_ROUND_UP_TO_ALIGNMENT_SIZE(sizeof(Zone));
+  }
+
  private:
   friend struct arena_detail::UnrefDestroy;
 

src/core/lib/transport/call_arena_allocator.h

@@ -66,6 +66,8 @@ class CallArenaAllocator final : public ArenaFactory {
     call_size_estimator_.UpdateCallSizeEstimate(arena->TotalUsedBytes());
   }
 
+  size_t CallSizeEstimate() { return call_size_estimator_.CallSizeEstimate(); }
+
  private:
   CallSizeEstimator call_size_estimator_;
 };

test/core/resource_quota/arena_test.cc

@@ -41,6 +41,8 @@
 #include "src/core/lib/resource_quota/resource_quota.h"
 #include "test/core/test_util/test_config.h"
 
+using testing::Mock;
+using testing::Return;
 using testing::StrictMock;
 
 namespace grpc_core {
@@ -84,6 +86,48 @@ INSTANTIATE_TEST_SUITE_P(
                       AllocShape{1, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}},
                       AllocShape{6, {1, 2, 3}}));
 
+class MockMemoryAllocatorImpl
+    : public grpc_event_engine::experimental::internal::MemoryAllocatorImpl {
+ public:
+  MOCK_METHOD(size_t, Reserve, (MemoryRequest));
+  MOCK_METHOD(grpc_slice, MakeSlice, (MemoryRequest));
+  MOCK_METHOD(void, Release, (size_t));
+  MOCK_METHOD(void, Shutdown, ());
+};
+
+TEST(ArenaTest, InitialReservationCorrect) {
+  auto allocator_impl = std::make_shared<StrictMock<MockMemoryAllocatorImpl>>();
+  auto allocator = SimpleArenaAllocator(1024, MemoryAllocator(allocator_impl));
+  EXPECT_CALL(*allocator_impl, Reserve(MemoryRequest(1024, 1024)))
+      .WillOnce(Return(1024));
+  auto arena = allocator->MakeArena();
+  Mock::VerifyAndClearExpectations(allocator_impl.get());
+  EXPECT_CALL(*allocator_impl, Release(1024));
+  arena.reset();
+  Mock::VerifyAndClearExpectations(allocator_impl.get());
+  EXPECT_CALL(*allocator_impl, Shutdown());
+}
+
+TEST(ArenaTest, SubsequentReservationCorrect) {
+  auto allocator_impl = std::make_shared<StrictMock<MockMemoryAllocatorImpl>>();
+  auto allocator = SimpleArenaAllocator(1024, MemoryAllocator(allocator_impl));
+  EXPECT_CALL(*allocator_impl, Reserve(MemoryRequest(1024, 1024)))
+      .WillOnce(Return(1024));
+  auto arena = allocator->MakeArena();
+  Mock::VerifyAndClearExpectations(allocator_impl.get());
+  EXPECT_CALL(*allocator_impl,
+              Reserve(MemoryRequest(4096 + Arena::ArenaZoneOverhead(),
+                                    4096 + Arena::ArenaZoneOverhead())))
+      .WillOnce(Return(4096 + Arena::ArenaZoneOverhead()));
+  arena->Alloc(4096);
+  Mock::VerifyAndClearExpectations(allocator_impl.get());
+  EXPECT_CALL(*allocator_impl,
+              Release(1024 + 4096 + Arena::ArenaZoneOverhead()));
+  arena.reset();
+  Mock::VerifyAndClearExpectations(allocator_impl.get());
+  EXPECT_CALL(*allocator_impl, Shutdown());
+}
+
 #define CONCURRENT_TEST_THREADS 10
 
 size_t concurrent_test_iterations() {
@@ -313,6 +357,32 @@ TEST(ArenaTest, FinalizeArenaIsCalled) {
   arena.reset();
 }
 
+TEST(ArenaTest, AccurateBaseByteCount) {
+  auto factory = MakeRefCounted<StrictMock<MockArenaFactory>>();
+  auto arena = Arena::Create(1, factory);
+  EXPECT_CALL(*factory, FinalizeArena(arena.get())).WillOnce([](Arena* a) {
+    EXPECT_EQ(a->TotalUsedBytes(),
+              Arena::ArenaOverhead() +
+                  GPR_ROUND_UP_TO_ALIGNMENT_SIZE(
+                      arena_detail::BaseArenaContextTraits::ContextSize()));
+  });
+  arena.reset();
+}
+
+TEST(ArenaTest, AccurateByteCountWithAllocation) {
+  auto factory = MakeRefCounted<StrictMock<MockArenaFactory>>();
+  auto arena = Arena::Create(1, factory);
+  arena->Alloc(1000);
+  EXPECT_CALL(*factory, FinalizeArena(arena.get())).WillOnce([](Arena* a) {
+    EXPECT_EQ(a->TotalUsedBytes(),
+              Arena::ArenaOverhead() +
+                  GPR_ROUND_UP_TO_ALIGNMENT_SIZE(
+                      arena_detail::BaseArenaContextTraits::ContextSize()) +
+                  GPR_ROUND_UP_TO_ALIGNMENT_SIZE(1000));
+  });
+  arena.reset();
+}
+
 }  // namespace grpc_core
 
 int main(int argc, char* argv[]) {

test/core/transport/BUILD

@@ -36,6 +36,23 @@ grpc_cc_test(
     ],
 )
 
+grpc_cc_test(
+    name = "call_arena_allocator_test",
+    srcs = ["call_arena_allocator_test.cc"],
+    external_deps = [
+        "gtest",
+    ],
+    language = "C++",
+    tags = ["no_windows"],  # TODO(jtattermusch): investigate the timeout on windows
+    uses_event_engine = False,
+    uses_polling = False,
+    deps = [
+        "//:gpr",
+        "//:grpc",
+        "//test/core/test_util:grpc_test_util",
+    ],
+)
+
 grpc_cc_test(
     name = "interception_chain_test",
     srcs = ["interception_chain_test.cc"],

test/core/transport/call_arena_allocator_test.cc

@@ -0,0 +1,82 @@
+//
+//
+// Copyright 2017 gRPC authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//
+
+#include "src/core/lib/transport/call_arena_allocator.h"
+
+#include <inttypes.h>
+#include <string.h>
+
+#include <algorithm>
+#include <iosfwd>
+#include <memory>
+#include <ostream>
+#include <string>
+#include <vector>
+
+#include "absl/strings/str_join.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include <grpc/support/sync.h>
+#include <grpc/support/time.h>
+
+#include "src/core/lib/gprpp/ref_counted_ptr.h"
+#include "src/core/lib/gprpp/thd.h"
+#include "src/core/lib/iomgr/exec_ctx.h"
+#include "src/core/lib/resource_quota/resource_quota.h"
+#include "test/core/test_util/test_config.h"
+
+namespace grpc_core {
+
+TEST(CallArenaAllocatorTest, SettlesEmpty) {
+  auto allocator = MakeRefCounted<CallArenaAllocator>(
+      ResourceQuota::Default()->memory_quota()->CreateMemoryAllocator(
+          "test-allocator"),
+      1);
+  for (int i = 0; i < 10000; i++) {
+    allocator->MakeArena();
+  }
+  auto estimate = allocator->CallSizeEstimate();
+  for (int i = 0; i < 10000; i++) {
+    allocator->MakeArena();
+  }
+  EXPECT_EQ(estimate, allocator->CallSizeEstimate());
+}
+
+TEST(CallArenaAllocatorTest, SettlesWithAllocation) {
+  auto allocator = MakeRefCounted<CallArenaAllocator>(
+      ResourceQuota::Default()->memory_quota()->CreateMemoryAllocator(
+          "test-allocator"),
+      1);
+  for (int i = 0; i < 10000; i++) {
+    allocator->MakeArena()->Alloc(10000);
+  }
+  auto estimate = allocator->CallSizeEstimate();
+  for (int i = 0; i < 10000; i++) {
+    allocator->MakeArena()->Alloc(10000);
+  }
+  LOG(INFO) << estimate;
+}
+
+}  // namespace grpc_core
+
+int main(int argc, char* argv[]) {
+  grpc::testing::TestEnvironment give_me_a_name(&argc, argv);
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

tools/run_tests/generated/tests.json

@@ -1277,6 +1277,28 @@
     ],
     "uses_polling": false
   },
+  {
+    "args": [],
+    "benchmark": false,
+    "ci_platforms": [
+      "linux",
+      "mac",
+      "posix"
+    ],
+    "cpu_cost": 1.0,
+    "exclude_configs": [],
+    "exclude_iomgrs": [],
+    "flaky": false,
+    "gtest": true,
+    "language": "c++",
+    "name": "call_arena_allocator_test",
+    "platforms": [
+      "linux",
+      "mac",
+      "posix"
+    ],
+    "uses_polling": false
+  },
   {
     "args": [],
     "benchmark": false,