| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you 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. |
| |
| |
| #ifndef IMPALA_RUNTIME_MEM_POOL_H |
| #define IMPALA_RUNTIME_MEM_POOL_H |
| |
| #include <stdio.h> |
| |
| #include <algorithm> |
| #include <cstddef> |
| #include <string> |
| #include <vector> |
| |
| #include "common/logging.h" |
| #include "gutil/dynamic_annotations.h" |
| #include "gutil/threading/thread_collision_warner.h" |
| #include "util/bit-util.h" |
| |
| namespace impala { |
| |
| class MemTracker; |
| |
| /// A MemPool maintains a list of memory chunks from which it allocates memory in |
| /// response to Allocate() calls; |
| /// Chunks stay around for the lifetime of the mempool or until they are passed on to |
| /// another mempool. |
| // |
| /// The caller registers a MemTracker with the pool; chunk allocations are counted |
| /// against that tracker and all of its ancestors. If chunks get moved between pools |
| /// during AcquireData() calls, the respective MemTrackers are updated accordingly. |
| /// Chunks freed up in the d'tor are subtracted from the registered trackers. |
| // |
| /// An Allocate() call will attempt to allocate memory from the chunk that was most |
| /// recently added; if that chunk doesn't have enough memory to |
| /// satisfy the allocation request, the free chunks are searched for one that is |
| /// big enough otherwise a new chunk is added to the list. |
| /// In order to keep allocation overhead low, chunk sizes double with each new one |
| /// added, until they hit a maximum size. But if the required size is greater than the |
| /// next chunk size, then a new chunk with the required size is allocated and the next |
| /// chunk size is set to the min(2*(required size), max chunk size). However if the |
| /// 'enforce_binary_chunk_sizes' flag passed to the c'tor is true, then all chunk sizes |
| /// allocated will be rounded up to the next power of two. |
| /// |
| /// Allocated chunks can be reused for new allocations if Clear() is called to free |
| /// all allocations or ReturnPartialAllocation() is called to return part of the last |
| /// allocation. |
| /// |
| /// All chunks before 'current_chunk_idx_' have allocated memory, while all chunks |
| /// after 'current_chunk_idx_' are free. The chunk at 'current_chunk_idx_' may or may |
| /// not have allocated memory. |
| /// |
| /// Example: |
| /// MemPool* p = new MemPool(); |
| /// for (int i = 0; i < 1024; ++i) { |
| /// returns 8-byte aligned memory (effectively 24 bytes): |
| /// .. = p->Allocate(17); |
| /// } |
| /// at this point, 17K have been handed out in response to Allocate() calls and |
| /// 28K of chunks have been allocated (chunk sizes: 4K, 8K, 16K) |
| /// We track total and peak allocated bytes. At this point they would be the same: |
| /// 28k bytes. A call to Clear will return the allocated memory so |
| /// total_allocated_bytes_ becomes 0. |
| /// p->Clear(); |
| /// the entire 1st chunk is returned: |
| /// .. = p->Allocate(4 * 1024); |
| /// 4K of the 2nd chunk are returned: |
| /// .. = p->Allocate(4 * 1024); |
| /// a new 20K chunk is created |
| /// .. = p->Allocate(20 * 1024); |
| // |
| /// MemPool* p2 = new MemPool(); |
| /// the new mempool receives all chunks containing data from p |
| /// p2->AcquireData(p, false); |
| /// At this point p.total_allocated_bytes_ would be 0. |
| /// The one remaining (empty) chunk is released: |
| /// delete p; |
| // |
| /// This class is not thread-safe. A DFAKE_MUTEX is used to help enforce correct usage. |
| |
| class MemPool { |
| public: |
| /// 'tracker' tracks the amount of memory allocated by this pool. Must not be NULL. |
| /// If 'enforce_binary_chunk_sizes' is set to true then all chunk sizes |
| /// allocated will be rounded up to the next power of two. |
| MemPool(MemTracker* mem_tracker, bool enforce_binary_chunk_sizes = false); |
| |
| /// Frees all chunks of memory and subtracts the total allocated bytes |
| /// from the registered limits. |
| ~MemPool(); |
| |
| /// Allocates a section of memory of 'size' bytes with DEFAULT_ALIGNMENT at the end |
| /// of the the current chunk. Creates a new chunk if there aren't any chunks |
| /// with enough capacity. |
| uint8_t* Allocate(int64_t size) noexcept { |
| DFAKE_SCOPED_LOCK(mutex_); |
| return Allocate<false>(size, DEFAULT_ALIGNMENT); |
| } |
| |
| /// Same as Allocate() except the mem limit is checked before the allocation and |
| /// this call will fail (returns NULL) if it does. |
| /// The caller must handle the NULL case. This should be used for allocations |
| /// where the size can be very big to bound the amount by which we exceed mem limits. |
| uint8_t* TryAllocate(int64_t size) noexcept { |
| DFAKE_SCOPED_LOCK(mutex_); |
| return Allocate<true>(size, DEFAULT_ALIGNMENT); |
| } |
| |
| /// Same as TryAllocate() except a non-default alignment can be specified. It |
| /// should be a power-of-two in [1, alignof(std::max_align_t)]. |
| uint8_t* TryAllocateAligned(int64_t size, int alignment) noexcept { |
| DFAKE_SCOPED_LOCK(mutex_); |
| DCHECK_GE(alignment, 1); |
| DCHECK_LE(alignment, alignof(std::max_align_t)); |
| DCHECK_EQ(BitUtil::RoundUpToPowerOfTwo(alignment), alignment); |
| return Allocate<true>(size, alignment); |
| } |
| |
| /// Same as TryAllocate() except returned memory is not aligned at all. |
| uint8_t* TryAllocateUnaligned(int64_t size) noexcept { |
| DFAKE_SCOPED_LOCK(mutex_); |
| // Call templated implementation directly so that it is inlined here and the |
| // alignment logic can be optimised out. |
| return Allocate<true>(size, 1); |
| } |
| |
| /// Returns 'byte_size' to the current chunk back to the mem pool. This can |
| /// only be used to return either all or part of the previous allocation returned |
| /// by Allocate(). |
| void ReturnPartialAllocation(int64_t byte_size) { |
| DFAKE_SCOPED_LOCK(mutex_); |
| DCHECK_GE(byte_size, 0); |
| DCHECK(current_chunk_idx_ != -1); |
| ChunkInfo& info = chunks_[current_chunk_idx_]; |
| DCHECK_GE(info.allocated_bytes, byte_size); |
| info.allocated_bytes -= byte_size; |
| ASAN_POISON_MEMORY_REGION(info.data + info.allocated_bytes, byte_size); |
| total_allocated_bytes_ -= byte_size; |
| } |
| |
| /// Return a dummy pointer for zero-length allocations. |
| static uint8_t* EmptyAllocPtr() { |
| return reinterpret_cast<uint8_t*>(&zero_length_region_); |
| } |
| |
| /// Makes all allocated chunks available for re-use, but doesn't delete any chunks. |
| void Clear(); |
| |
| /// Deletes all allocated chunks. FreeAll() or AcquireData() must be called for |
| /// each mem pool |
| void FreeAll(); |
| |
| /// Absorb all chunks that hold data from src. If keep_current is true, let src hold on |
| /// to its last allocated chunk that contains data. |
| /// All offsets handed out by calls to GetCurrentOffset() for 'src' become invalid. |
| void AcquireData(MemPool* src, bool keep_current); |
| |
| /// Change the MemTracker, updating consumption on the current and new tracker. |
| void SetMemTracker(MemTracker* new_tracker); |
| |
| std::string DebugString(); |
| |
| int64_t total_allocated_bytes() const { return total_allocated_bytes_; } |
| int64_t total_reserved_bytes() const { return total_reserved_bytes_; } |
| MemTracker* mem_tracker() { return mem_tracker_; } |
| |
| /// Return sum of chunk_sizes_. |
| int64_t GetTotalChunkSizes() const; |
| |
| /// TODO: make a macro for doing this |
| /// For C++/IR interop, we need to be able to look up types by name. |
| static const char* LLVM_CLASS_NAME; |
| |
| static const int DEFAULT_ALIGNMENT = 8; |
| |
| private: |
| friend class MemPoolTest; |
| static const int INITIAL_CHUNK_SIZE = 4 * 1024; |
| |
| /// The maximum size of chunk that should be allocated. Allocations larger than this |
| /// size will get their own individual chunk. Chosen to be small enough that it gets |
| /// a freelist in TCMalloc's central cache. |
| static const int MAX_CHUNK_SIZE = 512 * 1024; |
| |
| struct ChunkInfo { |
| uint8_t* data; // Owned by the ChunkInfo. |
| int64_t size; // in bytes |
| |
| /// bytes allocated via Allocate() in this chunk |
| int64_t allocated_bytes; |
| |
| explicit ChunkInfo(int64_t size, uint8_t* buf); |
| |
| ChunkInfo() |
| : data(NULL), |
| size(0), |
| allocated_bytes(0) {} |
| }; |
| |
| /// A static field used as non-NULL pointer for zero length allocations. NULL is |
| /// reserved for allocation failures. It must be as aligned as max_align_t for |
| /// TryAllocateAligned(). |
| static uint32_t zero_length_region_ alignas(std::max_align_t); |
| |
| /// Ensures a MemPool is not used by two threads concurrently. |
| DFAKE_MUTEX(mutex_); |
| |
| /// chunk from which we served the last Allocate() call; |
| /// always points to the last chunk that contains allocated data; |
| /// chunks 0..current_chunk_idx_ - 1 are guaranteed to contain data |
| /// (chunks_[i].allocated_bytes > 0 for i: 0..current_chunk_idx_ - 1); |
| /// chunks after 'current_chunk_idx_' are "free chunks" that contain no data. |
| /// -1 if no chunks present |
| int current_chunk_idx_; |
| |
| /// The size of the next chunk to allocate. |
| int next_chunk_size_; |
| |
| /// sum of allocated_bytes_ |
| int64_t total_allocated_bytes_; |
| |
| /// sum of all bytes allocated in chunks_ |
| int64_t total_reserved_bytes_; |
| |
| std::vector<ChunkInfo> chunks_; |
| |
| /// The current and peak memory footprint of this pool. This is different from |
| /// total allocated_bytes_ since it includes bytes in chunks that are not used. |
| MemTracker* mem_tracker_; |
| |
| /// If set to true, all chunk sizes allocated will be rounded up to the next power of |
| /// two. |
| const bool enforce_binary_chunk_sizes_; |
| |
| /// Find or allocated a chunk with at least min_size spare capacity and update |
| /// current_chunk_idx_. Also updates chunks_, chunk_sizes_ and allocated_bytes_ |
| /// if a new chunk needs to be created. |
| /// If check_limits is true, this call can fail (returns false) if adding a |
| /// new chunk exceeds the mem limits. |
| bool FindChunk(int64_t min_size, bool check_limits) noexcept; |
| |
| /// Check integrity of the supporting data structures; always returns true but DCHECKs |
| /// all invariants. |
| /// If 'check_current_chunk_empty' is true, checks that the current chunk contains no |
| /// data. Otherwise the current chunk can be either empty or full. |
| bool CheckIntegrity(bool check_current_chunk_empty); |
| |
| /// Return offset to unoccupied space in current chunk. |
| int64_t GetFreeOffset() const { |
| if (current_chunk_idx_ == -1) return 0; |
| return chunks_[current_chunk_idx_].allocated_bytes; |
| } |
| |
| template <bool CHECK_LIMIT_FIRST> |
| uint8_t* ALWAYS_INLINE Allocate(int64_t size, int alignment) noexcept { |
| DCHECK_GE(size, 0); |
| if (UNLIKELY(size == 0)) return reinterpret_cast<uint8_t*>(&zero_length_region_); |
| |
| if (current_chunk_idx_ != -1) { |
| ChunkInfo& info = chunks_[current_chunk_idx_]; |
| int64_t aligned_allocated_bytes = BitUtil::RoundUpToPowerOf2( |
| info.allocated_bytes, alignment); |
| if (aligned_allocated_bytes + size <= info.size) { |
| // Ensure the requested alignment is respected. |
| int64_t padding = aligned_allocated_bytes - info.allocated_bytes; |
| uint8_t* result = info.data + aligned_allocated_bytes; |
| ASAN_UNPOISON_MEMORY_REGION(result, size); |
| DCHECK_LE(info.allocated_bytes + size, info.size); |
| info.allocated_bytes += padding + size; |
| total_allocated_bytes_ += padding + size; |
| DCHECK_LE(current_chunk_idx_, chunks_.size() - 1); |
| return result; |
| } |
| } |
| |
| // If we couldn't allocate a new chunk, return NULL. malloc() guarantees alignment |
| // of alignof(std::max_align_t), so we do not need to do anything additional to |
| // guarantee alignment. |
| static_assert( |
| INITIAL_CHUNK_SIZE >= alignof(std::max_align_t), "Min chunk size too low"); |
| if (UNLIKELY(!FindChunk(size, CHECK_LIMIT_FIRST))) return NULL; |
| |
| ChunkInfo& info = chunks_[current_chunk_idx_]; |
| uint8_t* result = info.data + info.allocated_bytes; |
| ASAN_UNPOISON_MEMORY_REGION(result, size); |
| DCHECK_LE(info.allocated_bytes + size, info.size); |
| info.allocated_bytes += size; |
| total_allocated_bytes_ += size; |
| DCHECK_LE(current_chunk_idx_, chunks_.size() - 1); |
| return result; |
| } |
| }; |
| |
| // Stamp out templated implementations here so they're included in IR module |
| template uint8_t* MemPool::Allocate<false>(int64_t size, int alignment) noexcept; |
| template uint8_t* MemPool::Allocate<true>(int64_t size, int alignment) noexcept; |
| } |
| |
| #endif |