Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-09-22 08:53:40

0001 // Protocol Buffers - Google's data interchange format
0002 // Copyright 2022 Google Inc.  All rights reserved.
0003 //
0004 // Use of this source code is governed by a BSD-style
0005 // license that can be found in the LICENSE file or at
0006 // https://developers.google.com/open-source/licenses/bsd
0007 //
0008 // This file defines the internal class ThreadSafeArena
0009 
0010 #ifndef GOOGLE_PROTOBUF_THREAD_SAFE_ARENA_H__
0011 #define GOOGLE_PROTOBUF_THREAD_SAFE_ARENA_H__
0012 
0013 #include <atomic>
0014 #include <cstddef>
0015 #include <cstdint>
0016 #include <type_traits>
0017 #include <vector>
0018 
0019 #include "absl/base/attributes.h"
0020 #include "absl/base/optimization.h"
0021 #include "absl/synchronization/mutex.h"
0022 #include "google/protobuf/arena_align.h"
0023 #include "google/protobuf/arena_allocation_policy.h"
0024 #include "google/protobuf/arena_cleanup.h"
0025 #include "google/protobuf/arenaz_sampler.h"
0026 #include "google/protobuf/port.h"
0027 #include "google/protobuf/serial_arena.h"
0028 
0029 // Must be included last.
0030 #include "google/protobuf/port_def.inc"
0031 
0032 namespace google {
0033 namespace protobuf {
0034 namespace internal {
0035 
0036 // This class provides the core Arena memory allocation library. Different
0037 // implementations only need to implement the public interface below.
0038 // Arena is not a template type as that would only be useful if all protos
0039 // in turn would be templates, which will/cannot happen. However separating
0040 // the memory allocation part from the cruft of the API users expect we can
0041 // use #ifdef the select the best implementation based on hardware / OS.
0042 class PROTOBUF_EXPORT ThreadSafeArena {
0043  public:
0044   ThreadSafeArena();
0045 
0046   ThreadSafeArena(char* mem, size_t size);
0047 
0048   explicit ThreadSafeArena(void* mem, size_t size,
0049                            const AllocationPolicy& policy);
0050 
0051   // All protos have pointers back to the arena hence Arena must have
0052   // pointer stability.
0053   ThreadSafeArena(const ThreadSafeArena&) = delete;
0054   ThreadSafeArena& operator=(const ThreadSafeArena&) = delete;
0055   ThreadSafeArena(ThreadSafeArena&&) = delete;
0056   ThreadSafeArena& operator=(ThreadSafeArena&&) = delete;
0057 
0058   // Destructor deletes all owned heap allocated objects, and destructs objects
0059   // that have non-trivial destructors, except for proto2 message objects whose
0060   // destructors can be skipped. Also, frees all blocks except the initial block
0061   // if it was passed in.
0062   ~ThreadSafeArena();
0063 
0064   uint64_t Reset();
0065 
0066   uint64_t SpaceAllocated() const;
0067   uint64_t SpaceUsed() const;
0068 
0069   template <AllocationClient alloc_client = AllocationClient::kDefault>
0070   void* AllocateAligned(size_t n) {
0071     SerialArena* arena;
0072     if (ABSL_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
0073       return arena->AllocateAligned<alloc_client>(n);
0074     } else {
0075       return AllocateAlignedFallback<alloc_client>(n);
0076     }
0077   }
0078 
0079   void ReturnArrayMemory(void* p, size_t size) {
0080     SerialArena* arena = nullptr;
0081     if (ABSL_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
0082       arena->ReturnArrayMemory(p, size);
0083     }
0084   }
0085 
0086   // This function allocates n bytes if the common happy case is true and
0087   // returns true. Otherwise does nothing and returns false. This strange
0088   // semantics is necessary to allow callers to program functions that only
0089   // have fallback function calls in tail position. This substantially improves
0090   // code for the happy path.
0091   PROTOBUF_NDEBUG_INLINE bool MaybeAllocateAligned(size_t n, void** out) {
0092     SerialArena* arena = nullptr;
0093     if (ABSL_PREDICT_TRUE(GetSerialArenaFast(&arena))) {
0094       return arena->MaybeAllocateAligned(n, out);
0095     }
0096     return false;
0097   }
0098 
0099   void* AllocateAlignedWithCleanup(size_t n, size_t align,
0100                                    void (*destructor)(void*));
0101 
0102   // Add object pointer and cleanup function pointer to the list.
0103   void AddCleanup(void* elem, void (*cleanup)(void*));
0104 
0105   void* AllocateFromStringBlock();
0106 
0107   std::vector<void*> PeekCleanupListForTesting();
0108 
0109  private:
0110   friend class ArenaBenchmark;
0111   friend class TcParser;
0112   friend class SerialArena;
0113   friend struct SerialArenaChunkHeader;
0114   friend class cleanup::ChunkList;
0115   static uint64_t GetNextLifeCycleId();
0116 
0117   class SerialArenaChunk;
0118 
0119   // Returns a new SerialArenaChunk that has {id, serial} at slot 0. It may
0120   // grow based on "prev_num_slots".
0121   static SerialArenaChunk* NewSerialArenaChunk(uint32_t prev_capacity, void* id,
0122                                                SerialArena* serial);
0123   static SerialArenaChunk* SentrySerialArenaChunk();
0124 
0125   // Returns the first ArenaBlock* for the first SerialArena. If users provide
0126   // one, use it if it's acceptable. Otherwise returns a sentry block.
0127   ArenaBlock* FirstBlock(void* buf, size_t size);
0128   // Same as the above but returns a valid block if "policy" is not default.
0129   ArenaBlock* FirstBlock(void* buf, size_t size,
0130                          const AllocationPolicy& policy);
0131 
0132   // Adds SerialArena to the chunked list. May create a new chunk.
0133   void AddSerialArena(void* id, SerialArena* serial);
0134 
0135   void UnpoisonAllArenaBlocks() const;
0136 
0137   // Members are declared here to track sizeof(ThreadSafeArena) and hotness
0138   // centrally.
0139 
0140   // Unique for each arena. Changes on Reset().
0141   uint64_t tag_and_id_ = 0;
0142 
0143   TaggedAllocationPolicyPtr alloc_policy_;  // Tagged pointer to AllocPolicy.
0144   ThreadSafeArenaStatsHandle arena_stats_;
0145 
0146   // Adding a new chunk to head_ must be protected by mutex_.
0147   absl::Mutex mutex_;
0148   // Pointer to a linked list of SerialArenaChunk.
0149   std::atomic<SerialArenaChunk*> head_{nullptr};
0150 
0151   void* first_owner_;
0152   // Must be declared after alloc_policy_; otherwise, it may lose info on
0153   // user-provided initial block.
0154   SerialArena first_arena_;
0155 
0156   static_assert(std::is_trivially_destructible<SerialArena>{},
0157                 "SerialArena needs to be trivially destructible.");
0158 
0159   const AllocationPolicy* AllocPolicy() const { return alloc_policy_.get(); }
0160   void InitializeWithPolicy(const AllocationPolicy& policy);
0161   void* AllocateAlignedWithCleanupFallback(size_t n, size_t align,
0162                                            void (*destructor)(void*));
0163 
0164   void Init();
0165 
0166   // Delete or Destruct all objects owned by the arena.
0167   void CleanupList();
0168 
0169   inline void CacheSerialArena(SerialArena* serial) {
0170     thread_cache().last_serial_arena = serial;
0171     thread_cache().last_lifecycle_id_seen = tag_and_id_;
0172   }
0173 
0174   PROTOBUF_NDEBUG_INLINE bool GetSerialArenaFast(SerialArena** arena) {
0175     // If this thread already owns a block in this arena then try to use that.
0176     // This fast path optimizes the case where multiple threads allocate from
0177     // the same arena.
0178     ThreadCache* tc = &thread_cache();
0179     if (ABSL_PREDICT_TRUE(tc->last_lifecycle_id_seen == tag_and_id_)) {
0180       *arena = tc->last_serial_arena;
0181       return true;
0182     }
0183     return false;
0184   }
0185 
0186   // Finds SerialArena or creates one if not found. When creating a new one,
0187   // create a big enough block to accommodate n bytes.
0188   SerialArena* GetSerialArenaFallback(size_t n);
0189 
0190   SerialArena* GetSerialArena();
0191 
0192   template <AllocationClient alloc_client = AllocationClient::kDefault>
0193   void* AllocateAlignedFallback(size_t n);
0194 
0195   // Executes callback function over SerialArenaChunk. Passes const
0196   // SerialArenaChunk*.
0197   template <typename Callback>
0198   void WalkConstSerialArenaChunk(Callback fn) const;
0199 
0200   // Executes callback function over SerialArenaChunk.
0201   template <typename Callback>
0202   void WalkSerialArenaChunk(Callback fn);
0203 
0204   // Visits SerialArena and calls "fn", including "first_arena" and ones on
0205   // chunks. Do not rely on the order of visit. The callback function should
0206   // accept `const SerialArena*`.
0207   template <typename Callback>
0208   void VisitSerialArena(Callback fn) const;
0209 
0210   // Releases all memory except the first block which it returns. The first
0211   // block might be owned by the user and thus need some extra checks before
0212   // deleting.
0213   SizedPtr Free();
0214 
0215   // ThreadCache is accessed very frequently, so we align it such that it's
0216   // located within a single cache line.
0217   static constexpr size_t kThreadCacheAlignment = 32;
0218 
0219 #ifdef _MSC_VER
0220 #pragma warning(disable : 4324)
0221 #endif
0222   struct alignas(kThreadCacheAlignment) ThreadCache {
0223     // Number of per-thread lifecycle IDs to reserve. Must be power of two.
0224     // To reduce contention on a global atomic, each thread reserves a batch of
0225     // IDs.  The following number is calculated based on a stress test with
0226     // ~6500 threads all frequently allocating a new arena.
0227     static constexpr size_t kPerThreadIds = 256;
0228     // Next lifecycle ID available to this thread. We need to reserve a new
0229     // batch, if `next_lifecycle_id & (kPerThreadIds - 1) == 0`.
0230     uint64_t next_lifecycle_id{0};
0231     // The ThreadCache is considered valid as long as this matches the
0232     // lifecycle_id of the arena being used.
0233     uint64_t last_lifecycle_id_seen{static_cast<uint64_t>(-1)};
0234     SerialArena* last_serial_arena{nullptr};
0235   };
0236   static_assert(sizeof(ThreadCache) <= kThreadCacheAlignment,
0237                 "ThreadCache may span several cache lines");
0238 
0239   // Lifecycle_id can be highly contended variable in a situation of lots of
0240   // arena creation. Make sure that other global variables are not sharing the
0241   // cacheline.
0242 #ifdef _MSC_VER
0243 #pragma warning(disable : 4324)
0244 #endif
0245   using LifecycleId = uint64_t;
0246   alignas(kCacheAlignment) ABSL_CONST_INIT
0247       static std::atomic<LifecycleId> lifecycle_id_;
0248 #if defined(PROTOBUF_NO_THREADLOCAL)
0249   // iOS does not support __thread keyword so we use a custom thread local
0250   // storage class we implemented.
0251   static ThreadCache& thread_cache();
0252 #elif defined(PROTOBUF_USE_DLLS) && defined(_WIN32)
0253   // Thread local variables cannot be exposed through MSVC DLL interface but we
0254   // can wrap them in static functions.
0255   static ThreadCache& thread_cache();
0256 #else
0257   PROTOBUF_CONSTINIT static PROTOBUF_THREAD_LOCAL ThreadCache thread_cache_;
0258   static ThreadCache& thread_cache() { return thread_cache_; }
0259 #endif
0260 
0261  public:
0262   // kBlockHeaderSize is sizeof(ArenaBlock), aligned up to the nearest multiple
0263   // of 8 to protect the invariant that pos is always at a multiple of 8.
0264   static constexpr size_t kBlockHeaderSize = SerialArena::kBlockHeaderSize;
0265   static constexpr size_t kSerialArenaSize =
0266       (sizeof(SerialArena) + 7) & static_cast<size_t>(-8);
0267   static constexpr size_t kAllocPolicySize =
0268       ArenaAlignDefault::Ceil(sizeof(AllocationPolicy));
0269   static constexpr size_t kMaxCleanupNodeSize = 16;
0270   static_assert(kBlockHeaderSize % 8 == 0,
0271                 "kBlockHeaderSize must be a multiple of 8.");
0272   static_assert(kSerialArenaSize % 8 == 0,
0273                 "kSerialArenaSize must be a multiple of 8.");
0274 };
0275 
0276 }  // namespace internal
0277 }  // namespace protobuf
0278 }  // namespace google
0279 
0280 #include "google/protobuf/port_undef.inc"
0281 
0282 #endif  // GOOGLE_PROTOBUF_THREAD_SAFE_ARENA_H__