google/protobuf/serial_arena.h

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 SerialArena
0009
0010 #ifndef GOOGLE_PROTOBUF_SERIAL_ARENA_H__
0011 #define GOOGLE_PROTOBUF_SERIAL_ARENA_H__
0012
0013 #include <algorithm>
0014 #include <atomic>
0015 #include <cstddef>
0016 #include <cstdint>
0017 #include <string>
0018 #include <vector>
0019
0020 #include "absl/base/attributes.h"
0021 #include "absl/base/optimization.h"
0022 #include "absl/base/prefetch.h"
0023 #include "absl/log/absl_check.h"
0024 #include "absl/numeric/bits.h"
0025 #include "google/protobuf/arena_align.h"
0026 #include "google/protobuf/arena_cleanup.h"
0027 #include "google/protobuf/port.h"
0028 #include "google/protobuf/string_block.h"
0029
0030 // Must be included last.
0031 #include "google/protobuf/port_def.inc"
0032
0033 namespace google {
0034 namespace protobuf {
0035 namespace internal {
0036
0037 // Arena blocks are variable length malloc-ed objects.  The following structure
0038 // describes the common header for all blocks.
0039 struct ArenaBlock {
0040   // For the sentry block with zero-size where ptr_/limit_ both point to `this`.
0041   constexpr ArenaBlock() : next(nullptr), size(0) {}
0042
0043   ArenaBlock(ArenaBlock* next, size_t size) : next(next), size(size) {
0044     ABSL_DCHECK_GT(size, sizeof(ArenaBlock));
0045   }
0046
0047   char* Pointer(size_t n) {
0048     ABSL_DCHECK_LE(n, size);
0049     return reinterpret_cast<char*>(this) + n;
0050   }
0051   char* Limit() { return Pointer(size & static_cast<size_t>(-8)); }
0052
0053   bool IsSentry() const { return size == 0; }
0054
0055   ArenaBlock* const next;
0056   const size_t size;
0057   // data follows
0058 };
0059
0060 enum class AllocationClient { kDefault, kArray };
0061
0062 class ThreadSafeArena;
0063
0064 // Tag type used to invoke the constructor of the first SerialArena.
0065 struct FirstSerialArena {
0066   explicit FirstSerialArena() = default;
0067 };
0068
0069 // A simple arena allocator. Calls to allocate functions must be properly
0070 // serialized by the caller, hence this class cannot be used as a general
0071 // purpose allocator in a multi-threaded program. It serves as a building block
0072 // for ThreadSafeArena, which provides a thread-safe arena allocator.
0073 //
0074 // This class manages
0075 // 1) Arena bump allocation + owning memory blocks.
0076 // 2) Maintaining a cleanup list.
0077 // It delegates the actual memory allocation back to ThreadSafeArena, which
0078 // contains the information on block growth policy and backing memory allocation
0079 // used.
0080 class PROTOBUF_EXPORT SerialArena {
0081  public:
0082   static constexpr size_t kBlockHeaderSize =
0083       ArenaAlignDefault::Ceil(sizeof(ArenaBlock));
0084
0085   void CleanupList() { cleanup_list_.Cleanup(*this); }
0086   uint64_t SpaceAllocated() const {
0087     return space_allocated_.load(std::memory_order_relaxed);
0088   }
0089   uint64_t SpaceUsed() const;
0090
0091   // See comments on `cached_blocks_` member for details.
0092   PROTOBUF_ALWAYS_INLINE void* TryAllocateFromCachedBlock(size_t size) {
0093     if (ABSL_PREDICT_FALSE(size < 16)) return nullptr;
0094     // We round up to the next larger block in case the memory doesn't match
0095     // the pattern we are looking for.
0096     const size_t index = absl::bit_width(size - 1) - 4;
0097
0098     if (ABSL_PREDICT_FALSE(index >= cached_block_length_)) return nullptr;
0099     auto& cached_head = cached_blocks_[index];
0100     if (cached_head == nullptr) return nullptr;
0101
0102     void* ret = cached_head;
0103     internal::UnpoisonMemoryRegion(ret, size);
0104     cached_head = cached_head->next;
0105     return ret;
0106   }
0107
0108   // In kArray mode we look through cached blocks.
0109   // We do not do this by default because most non-array allocations will not
0110   // have the right size and will fail to find an appropriate cached block.
0111   //
0112   // TODO: Evaluate if we should use cached blocks for message types of
0113   // the right size. We can statically know if the allocation size can benefit
0114   // from it.
0115   template <AllocationClient alloc_client = AllocationClient::kDefault>
0116   void* AllocateAligned(size_t n) {
0117     ABSL_DCHECK(internal::ArenaAlignDefault::IsAligned(n));
0118     ABSL_DCHECK_GE(limit_, ptr());
0119
0120     if (alloc_client == AllocationClient::kArray) {
0121       if (void* res = TryAllocateFromCachedBlock(n)) {
0122         return res;
0123       }
0124     }
0125
0126     void* ptr;
0127     if (ABSL_PREDICT_TRUE(MaybeAllocateAligned(n, &ptr))) {
0128       return ptr;
0129     }
0130     return AllocateAlignedFallback(n);
0131   }
0132
0133  private:
0134   static PROTOBUF_ALWAYS_INLINE constexpr size_t AlignUpTo(size_t n, size_t a) {
0135     // We are wasting space by over allocating align - 8 bytes. Compared to a
0136     // dedicated function that takes current alignment in consideration.  Such a
0137     // scheme would only waste (align - 8)/2 bytes on average, but requires a
0138     // dedicated function in the outline arena allocation functions. Possibly
0139     // re-evaluate tradeoffs later.
0140     return a <= 8 ? ArenaAlignDefault::Ceil(n) : ArenaAlignAs(a).Padded(n);
0141   }
0142
0143   static PROTOBUF_ALWAYS_INLINE void* AlignTo(void* p, size_t a) {
0144     return (a <= ArenaAlignDefault::align)
0145                ? ArenaAlignDefault::CeilDefaultAligned(p)
0146                : ArenaAlignAs(a).CeilDefaultAligned(p);
0147   }
0148
0149   // See comments on `cached_blocks_` member for details.
0150   void ReturnArrayMemory(void* p, size_t size) {
0151     // We only need to check for 32-bit platforms.
0152     // In 64-bit platforms the minimum allocation size from Repeated*Field will
0153     // be 16 guaranteed.
0154     if (sizeof(void*) < 8) {
0155       if (ABSL_PREDICT_FALSE(size < 16)) return;
0156     } else {
0157       PROTOBUF_ASSUME(size >= 16);
0158     }
0159
0160     // We round down to the next smaller block in case the memory doesn't match
0161     // the pattern we are looking for. eg, someone might have called Reserve()
0162     // on the repeated field.
0163     const size_t index = absl::bit_width(size) - 5;
0164
0165     if (ABSL_PREDICT_FALSE(index >= cached_block_length_)) {
0166       // We can't put this object on the freelist so make this object the
0167       // freelist. It is guaranteed it is larger than the one we have, and
0168       // large enough to hold another allocation of `size`.
0169       CachedBlock** new_list = static_cast<CachedBlock**>(p);
0170       size_t new_size = size / sizeof(CachedBlock*);
0171
0172       std::copy(cached_blocks_, cached_blocks_ + cached_block_length_,
0173                 new_list);
0174
0175       // We need to unpoison this memory before filling it in case it has been
0176       // poisoned by another sanitizer client.
0177       internal::UnpoisonMemoryRegion(
0178           new_list + cached_block_length_,
0179           (new_size - cached_block_length_) * sizeof(CachedBlock*));
0180
0181       std::fill(new_list + cached_block_length_, new_list + new_size, nullptr);
0182
0183       cached_blocks_ = new_list;
0184       // Make the size fit in uint8_t. This is the power of two, so we don't
0185       // need anything larger.
0186       cached_block_length_ =
0187           static_cast<uint8_t>(std::min(size_t{64}, new_size));
0188
0189       return;
0190     }
0191
0192     auto& cached_head = cached_blocks_[index];
0193     auto* new_node = static_cast<CachedBlock*>(p);
0194     new_node->next = cached_head;
0195     cached_head = new_node;
0196     internal::PoisonMemoryRegion(p, size);
0197   }
0198
0199  public:
0200   // Allocate space if the current region provides enough space.
0201   bool MaybeAllocateAligned(size_t n, void** out) {
0202     ABSL_DCHECK(internal::ArenaAlignDefault::IsAligned(n));
0203     ABSL_DCHECK_GE(limit_, ptr());
0204     char* ret = ptr();
0205     if (ABSL_PREDICT_FALSE(limit_ - ret < static_cast<ptrdiff_t>(n))) {
0206       return false;
0207     }
0208     internal::UnpoisonMemoryRegion(ret, n);
0209     *out = ret;
0210     char* next = ret + n;
0211     set_ptr(next);
0212     MaybePrefetchData(next);
0213     return true;
0214   }
0215
0216   // If there is enough space in the current block, allocate space for one
0217   // std::string object and register for destruction. The object has not been
0218   // constructed and the memory returned is uninitialized.
0219   PROTOBUF_ALWAYS_INLINE void* MaybeAllocateStringWithCleanup() {
0220     void* p;
0221     return MaybeAllocateString(p) ? p : nullptr;
0222   }
0223
0224   PROTOBUF_ALWAYS_INLINE
0225   void* AllocateAlignedWithCleanup(size_t n, size_t align,
0226                                    void (*destructor)(void*)) {
0227     n = ArenaAlignDefault::Ceil(n);
0228     char* ret = ArenaAlignAs(align).CeilDefaultAligned(ptr());
0229     // See the comment in MaybeAllocateAligned re uintptr_t.
0230     if (ABSL_PREDICT_FALSE(reinterpret_cast<uintptr_t>(ret) + n >
0231                            reinterpret_cast<uintptr_t>(limit_))) {
0232       return AllocateAlignedWithCleanupFallback(n, align, destructor);
0233     }
0234     internal::UnpoisonMemoryRegion(ret, n);
0235     char* next = ret + n;
0236     set_ptr(next);
0237     AddCleanup(ret, destructor);
0238     ABSL_DCHECK_GE(limit_, ptr());
0239     MaybePrefetchData(next);
0240     return ret;
0241   }
0242
0243   PROTOBUF_ALWAYS_INLINE
0244   void AddCleanup(void* elem, void (*destructor)(void*)) {
0245     cleanup_list_.Add(elem, destructor, *this);
0246     MaybePrefetchCleanup();
0247   }
0248
0249   ABSL_ATTRIBUTE_RETURNS_NONNULL void* AllocateFromStringBlock();
0250
0251   std::vector<void*> PeekCleanupListForTesting();
0252
0253  private:
0254   friend class ThreadSafeArena;
0255   friend class cleanup::ChunkList;
0256
0257   // See comments for cached_blocks_.
0258   struct CachedBlock {
0259     // Simple linked list.
0260     CachedBlock* next;
0261   };
0262
0263   static constexpr ptrdiff_t kPrefetchDataDegree = ABSL_CACHELINE_SIZE * 16;
0264   static constexpr ptrdiff_t kPrefetchCleanupDegree = ABSL_CACHELINE_SIZE * 6;
0265
0266   // Constructor is private as only New() should be used.
0267   inline SerialArena(ArenaBlock* b, ThreadSafeArena& parent);
0268
0269   // Constructors to handle the first SerialArena.
0270   inline explicit SerialArena(ThreadSafeArena& parent);
0271   inline SerialArena(FirstSerialArena, ArenaBlock* b, ThreadSafeArena& parent);
0272
0273   bool MaybeAllocateString(void*& p);
0274   ABSL_ATTRIBUTE_RETURNS_NONNULL void* AllocateFromStringBlockFallback();
0275
0276   // Prefetch the next prefetch_degree bytes after `prefetch_ptr` and
0277   // up to `limit`, if `next` is within prefetch_degree bytes of `prefetch_ptr`.
0278   PROTOBUF_ALWAYS_INLINE
0279   static const char* MaybePrefetchImpl(const ptrdiff_t prefetch_degree,
0280                                        const char* next, const char* limit,
0281                                        const char* prefetch_ptr) {
0282     if (ABSL_PREDICT_TRUE(prefetch_ptr - next > prefetch_degree))
0283       return prefetch_ptr;
0284     if (ABSL_PREDICT_TRUE(prefetch_ptr < limit)) {
0285       prefetch_ptr = std::max(next, prefetch_ptr);
0286       ABSL_DCHECK(prefetch_ptr != nullptr);
0287       const char* end = std::min(limit, prefetch_ptr + prefetch_degree);
0288       for (; prefetch_ptr < end; prefetch_ptr += ABSL_CACHELINE_SIZE) {
0289         absl::PrefetchToLocalCacheForWrite(prefetch_ptr);
0290       }
0291     }
0292     return prefetch_ptr;
0293   }
0294   PROTOBUF_ALWAYS_INLINE
0295   void MaybePrefetchData(const char* next) {
0296     ABSL_DCHECK(static_cast<const void*>(prefetch_ptr_) == ptr() ||
0297                 static_cast<const void*>(prefetch_ptr_) >= head());
0298     prefetch_ptr_ =
0299         MaybePrefetchImpl(kPrefetchDataDegree, next, limit_, prefetch_ptr_);
0300   }
0301   PROTOBUF_ALWAYS_INLINE
0302   void MaybePrefetchCleanup() {
0303     ABSL_DCHECK(static_cast<const void*>(cleanup_list_.prefetch_ptr_) ==
0304                     nullptr ||
0305                 static_cast<const void*>(cleanup_list_.prefetch_ptr_) >=
0306                     cleanup_list_.head_);
0307     cleanup_list_.prefetch_ptr_ = MaybePrefetchImpl(
0308         kPrefetchCleanupDegree, reinterpret_cast<char*>(cleanup_list_.next_),
0309         reinterpret_cast<char*>(cleanup_list_.limit_),
0310         cleanup_list_.prefetch_ptr_);
0311   }
0312
0313   // Creates a new SerialArena inside mem using the remaining memory as for
0314   // future allocations.
0315   // The `parent` arena must outlive the serial arena, which is guaranteed
0316   // because the parent manages the lifetime of the serial arenas.
0317   static SerialArena* New(SizedPtr mem, ThreadSafeArena& parent);
0318   // Free SerialArena returning the memory passed in to New.
0319   template <typename Deallocator>
0320   SizedPtr Free(Deallocator deallocator);
0321
0322   size_t FreeStringBlocks() {
0323     // On the active block delete all strings skipping the unused instances.
0324     size_t unused_bytes = string_block_unused_.load(std::memory_order_relaxed);
0325     if (StringBlock* sb = string_block_.load(std::memory_order_relaxed)) {
0326       return FreeStringBlocks(sb, unused_bytes);
0327     }
0328     return 0;
0329   }
0330   static size_t FreeStringBlocks(StringBlock* string_block, size_t unused);
0331
0332   // Adds 'used` to space_used_ in relaxed atomic order.
0333   void AddSpaceUsed(size_t space_used) {
0334     space_used_.store(space_used_.load(std::memory_order_relaxed) + space_used,
0335                       std::memory_order_relaxed);
0336   }
0337
0338   // Adds 'allocated` to space_allocated_ in relaxed atomic order.
0339   void AddSpaceAllocated(size_t space_allocated) {
0340     space_allocated_.store(
0341         space_allocated_.load(std::memory_order_relaxed) + space_allocated,
0342         std::memory_order_relaxed);
0343   }
0344
0345   // Helper getters/setters to handle relaxed operations on atomic variables.
0346   ArenaBlock* head() { return head_.load(std::memory_order_relaxed); }
0347   const ArenaBlock* head() const {
0348     return head_.load(std::memory_order_relaxed);
0349   }
0350
0351   char* ptr() { return ptr_.load(std::memory_order_relaxed); }
0352   const char* ptr() const { return ptr_.load(std::memory_order_relaxed); }
0353   void set_ptr(char* ptr) { return ptr_.store(ptr, std::memory_order_relaxed); }
0354   PROTOBUF_ALWAYS_INLINE void set_range(char* ptr, char* limit) {
0355     set_ptr(ptr);
0356     prefetch_ptr_ = ptr;
0357     limit_ = limit;
0358   }
0359
0360   void* AllocateAlignedFallback(size_t n);
0361   void* AllocateAlignedWithCleanupFallback(size_t n, size_t align,
0362                                            void (*destructor)(void*));
0363   void AddCleanupFallback(void* elem, void (*destructor)(void*));
0364   inline void AllocateNewBlock(size_t n);
0365   inline void Init(ArenaBlock* b, size_t offset);
0366
0367   // Members are declared here to track sizeof(SerialArena) and hotness
0368   // centrally. They are (roughly) laid out in descending order of hotness.
0369
0370   // We initialize ptr/limit with an arbitrary valid pointer.
0371   // This allows Allocate to always return non-null even when asking for zero
0372   // bytes.
0373   static char* ArbitraryInternalPointerForInit() {
0374     // Use rodata to detect potential bugs. No one should be writing here.
0375     alignas(8) static constexpr char dummy{};
0376     return const_cast<char*>(&dummy);
0377   }
0378
0379   // Next pointer to allocate from.  Always 8-byte aligned.  Points inside
0380   // head_ (and head_->pos will always be non-canonical).  We keep these
0381   // here to reduce indirection.
0382   std::atomic<char*> ptr_{ArbitraryInternalPointerForInit()};
0383   // Limiting address up to which memory can be allocated from the head block.
0384   char* limit_ = ArbitraryInternalPointerForInit();
0385   // Current prefetch positions. Data from `ptr_` up to but not including
0386   // `prefetch_ptr_` is software prefetched.
0387   const char* prefetch_ptr_ = ArbitraryInternalPointerForInit();
0388
0389   // Chunked linked list for managing cleanup for arena elements.
0390   cleanup::ChunkList cleanup_list_;
0391
0392   // The active string block.
0393   std::atomic<StringBlock*> string_block_{nullptr};
0394
0395   // The number of unused bytes in string_block_.
0396   // We allocate from `effective_size()` down to 0 inside `string_block_`.
0397   // `unused  == 0` means that `string_block_` is exhausted. (or null).
0398   std::atomic<size_t> string_block_unused_{0};
0399
0400   std::atomic<ArenaBlock*> head_{nullptr};  // Head of linked list of blocks.
0401   std::atomic<size_t> space_used_{0};       // Necessary for metrics.
0402   std::atomic<size_t> space_allocated_{0};
0403   ThreadSafeArena& parent_;
0404
0405   // Repeated*Field and Arena play together to reduce memory consumption by
0406   // reusing blocks. Currently, natural growth of the repeated field types makes
0407   // them allocate blocks of size `8 + 2^N, N>=3`.
0408   // When the repeated field grows returns the previous block and we put it in
0409   // this free list.
0410   // `cached_blocks_[i]` points to the free list for blocks of size `8+2^(i+3)`.
0411   // The array of freelists is grown when needed in `ReturnArrayMemory()`.
0412   uint8_t cached_block_length_ = 0;
0413   CachedBlock** cached_blocks_ = nullptr;
0414 };
0415
0416 PROTOBUF_ALWAYS_INLINE bool SerialArena::MaybeAllocateString(void*& p) {
0417   // Check how many unused instances are in the current block.
0418   size_t unused_bytes = string_block_unused_.load(std::memory_order_relaxed);
0419   if (ABSL_PREDICT_TRUE(unused_bytes != 0)) {
0420     unused_bytes -= sizeof(std::string);
0421     string_block_unused_.store(unused_bytes, std::memory_order_relaxed);
0422     p = string_block_.load(std::memory_order_relaxed)->AtOffset(unused_bytes);
0423     return true;
0424   }
0425   return false;
0426 }
0427
0428 ABSL_ATTRIBUTE_RETURNS_NONNULL PROTOBUF_ALWAYS_INLINE void*
0429 SerialArena::AllocateFromStringBlock() {
0430   void* p;
0431   if (ABSL_PREDICT_TRUE(MaybeAllocateString(p))) return p;
0432   return AllocateFromStringBlockFallback();
0433 }
0434
0435 }  // namespace internal
0436 }  // namespace protobuf
0437 }  // namespace google
0438
0439 #include "google/protobuf/port_undef.inc"
0440
0441 #endif  // GOOGLE_PROTOBUF_SERIAL_ARENA_H__