From 8e022f815814b901711f1d4eca41d40afdca71d6 Mon Sep 17 00:00:00 2001
From: Martijn Vels <mvels@google.com>
Date: Thu, 26 Jan 2023 11:15:17 -0800
Subject: [PATCH] Add test coverage for `AllocateAtLeast()`

PiperOrigin-RevId: 504893114
---
 src/google/protobuf/arena_unittest.cc |  4 ++++
 src/google/protobuf/port.h            | 34 +++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/src/google/protobuf/arena_unittest.cc b/src/google/protobuf/arena_unittest.cc
index a1be800221..4832cff0e9 100644
--- a/src/google/protobuf/arena_unittest.cc
+++ b/src/google/protobuf/arena_unittest.cc
@@ -31,10 +31,13 @@
 #include "google/protobuf/arena.h"
 
 #include <algorithm>
+#include <array>
+#include <atomic>
 #include <cstddef>
 #include <cstring>
 #include <memory>
 #include <string>
+#include <thread>
 #include <type_traits>
 #include <typeinfo>
 #include <vector>
@@ -51,6 +54,7 @@
 #include "google/protobuf/io/zero_copy_stream_impl_lite.h"
 #include "google/protobuf/message.h"
 #include "google/protobuf/message_lite.h"
+#include "google/protobuf/port.h"
 #include "google/protobuf/repeated_field.h"
 #include "google/protobuf/test_util.h"
 #include "google/protobuf/unittest.pb.h"
diff --git a/src/google/protobuf/port.h b/src/google/protobuf/port.h
index 7c8ce83da6..0abb04455d 100644
--- a/src/google/protobuf/port.h
+++ b/src/google/protobuf/port.h
@@ -63,11 +63,45 @@ struct SizedPtr {
   size_t n;
 };
 
+// Debug hook allowing setting up test scenarios for AllocateAtLeast usage.
+using AllocateAtLeastHookFn = SizedPtr (*)(size_t, void*);
+
+// `AllocAtLeastHook` API
+constexpr bool HaveAllocateAtLeastHook();
+void SetAllocateAtLeastHook(AllocateAtLeastHookFn fn, void* context = nullptr);
+
+#if !defined(NDEBUG) && defined(ABSL_HAVE_THREAD_LOCAL) && \
+    defined(__cpp_inline_variables)
+
+// Hook data for current thread. These vars must not be accessed directly, use
+// the 'HaveAllocateAtLeastHook()` and `SetAllocateAtLeastHook()` API instead.
+inline thread_local AllocateAtLeastHookFn allocate_at_least_hook = nullptr;
+inline thread_local void* allocate_at_least_hook_context = nullptr;
+
+constexpr bool HaveAllocateAtLeastHook() { return true; }
+inline void SetAllocateAtLeastHook(AllocateAtLeastHookFn fn, void* context) {
+  allocate_at_least_hook = fn;
+  allocate_at_least_hook_context = context;
+}
+
+#else  // !NDEBUG && ABSL_HAVE_THREAD_LOCAL && __cpp_inline_variables
+
+constexpr bool HaveAllocateAtLeastHook() { return false; }
+inline void SetAllocateAtLeastHook(AllocateAtLeastHookFn fn, void* context) {}
+
+#endif  // !NDEBUG && ABSL_HAVE_THREAD_LOCAL && __cpp_inline_variables
+
 // Allocates at least `size` bytes. This function follows the c++ language
 // proposal from D0901R10 (http://wg21.link/D0901R10) and will be implemented
 // in terms of the new operator new semantics when available. The allocated
 // memory should be released by a call to `SizedDelete` or `::operator delete`.
 inline SizedPtr AllocateAtLeast(size_t size) {
+#if !defined(NDEBUG) && defined(ABSL_HAVE_THREAD_LOCAL) && \
+    defined(__cpp_inline_variables)
+  if (allocate_at_least_hook != nullptr) {
+    return allocate_at_least_hook(size, allocate_at_least_hook_context);
+  }
+#endif  // !NDEBUG && ABSL_HAVE_THREAD_LOCAL && __cpp_inline_variables
   return {::operator new(size), size};
 }