SpartaSharedPointerAllocator.hpp

#pragma once
 
#include "sparta/utils/SpartaSharedPointer.hpp"
 
namespace sparta
{
 
    template<class PointerT>
    class SpartaSharedPointerAllocator : public BaseAllocator
    {
    public:
 
        using element_type = PointerT;
 
        using WaterMarkWarningCallback = std::function<void (const SpartaSharedPointerAllocator &)>;
 
        using OverAllocationCallback = std::function<void (const SpartaSharedPointerAllocator &)>;
 
        SpartaSharedPointerAllocator(const size_t max_num_blocks,
                                     const size_t water_mark) :
            memory_blocks_(max_num_blocks),
            water_mark_(water_mark),
            watermark_warning_callback_(waterMarkWarningCallback_),
            over_allocation_callback_(overAllocationCallback_)
        {
            sparta_assert(water_mark <= max_num_blocks,
                          "The water_mark on SpartaSharedPointerAllocator should be less than or " <<
                          "equal to the maximum number of blocks. water_mark=" << water_mark <<
                          " max_num_blocks=" << max_num_blocks);
            free_blocks_.resize(max_num_blocks, nullptr);
        }
 
        SpartaSharedPointerAllocator(const SpartaSharedPointerAllocator&) = delete;
        SpartaSharedPointerAllocator(SpartaSharedPointerAllocator &&) = delete;
        SpartaSharedPointerAllocator operator=(const SpartaSharedPointerAllocator&) = delete;
 
        ~SpartaSharedPointerAllocator() {
            if(hasOutstandingObjects())
            {
                std::cerr << "WARNING: Seems that not all of the blocks made it back.  \n'" <<
                    __PRETTY_FUNCTION__ << "'\nAllocated: " << allocated_ <<
                    "\nReturned: " << free_idx_ << std::endl;
            }
        }
 
        size_t getNumFree() const {
            return free_idx_;
        }
 
        size_t getNumAllocated() const {
            return memory_blocks_.size();
        }
 
        bool hasOutstandingObjects() const {
            return (allocated_ != free_idx_);
        }
 
        std::vector<const PointerT*> getOutstandingAllocatedObjects() const
        {
            std::vector<const PointerT*> allocated_objs;
 
            const size_t size = memory_blocks_.size();
            for(uint32_t i = 0; i < size; ++i) {
                if(memory_blocks_[i]->ref_count->count > 0) {
                    allocated_objs.emplace_back(memory_blocks_[i]->object);
                }
            }
 
            return allocated_objs;
        }
 
        void registerCustomWaterMarkCallback(const WaterMarkWarningCallback & callback) {
            watermark_warning_callback_ = callback;
        }
 
        void registerCustomOverAllocationCallback(const OverAllocationCallback & callback) {
            over_allocation_callback_ = callback;
        }
 
    private:
 
        // Let's make friends
        friend class SpartaSharedPointer<PointerT>;
 
        // Make the allocate function a buddy
        template<typename PtrT, typename... Args>
        friend SpartaSharedPointer<PtrT>
        allocate_sparta_shared_pointer(SpartaSharedPointerAllocator<PtrT> &, Args&&...args);
 
        template<typename T>
        struct AlignedStorage
        {
            alignas(T) std::byte buf[sizeof(T)];
        };
 
        // Internal MemoryBlock
        struct MemBlock : public BaseAllocator::MemBlockBase
        {
            using RefCountType = typename SpartaSharedPointer<PointerT>::RefCount;
 
            using RefCountAlignedStorage = AlignedStorage<RefCountType>;
 
            using PointerTAlignedStorage = AlignedStorage<PointerT>;
 
            template<typename ...ObjArgs>
            MemBlock(SpartaSharedPointerAllocator * alloc_in, ObjArgs&&...obj_args) :
                MemBlockBase(alloc_in)
            {
                // Build the new object using the inplacement new.
                object = new (&object_storage) PointerT(std::forward<ObjArgs>(obj_args)...);
 
                // Build the reference count using that object
                ref_count = new (&ref_count_storage) RefCountType(object, (void*)this);
            }
 
            PointerT              * object = nullptr;
            RefCountType          * ref_count = nullptr;
 
            RefCountAlignedStorage  ref_count_storage;
            PointerTAlignedStorage  object_storage;
 
            MemBlock(const MemBlock &) = delete;
            MemBlock(MemBlock &&)      = delete;
            MemBlock & operator=(const MemBlock&) = delete;
            ~MemBlock() {} // Does nothing; implementation for cppcheck
        };
 
        // Class to manager the large block of memory allocated
        class MemBlockVector
        {
            // Align the storage for the MemBlock
            using MemBlockAlignedStorage = AlignedStorage<MemBlock>;
 
            std::vector<MemBlockAlignedStorage> data_;
            std::size_t size_ = 0;
 
        public:
 
            // Create num_blocks amount of MemBlocks to use
            explicit MemBlockVector(const size_t num_blocks) :
                data_(num_blocks)
            {}
 
            // Allocate an object in aligned storage
            template<typename ...Args>
            MemBlock * allocate(SpartaSharedPointerAllocator * alloc, Args&&... args)
            {
                sparta_assert(size_ < data_.size(), "Out of memory");
 
                // construct value in memory of aligned storage
                // using inplace operator new
                auto block = new(&data_[size_]) MemBlock(alloc, std::forward<Args>(args)...);
                ++size_;
 
                return block;
            }
 
            // Access an object in aligned storage
            const MemBlock* operator[](std::size_t idx) const
            {
                sparta_assert(idx < size_);
                return reinterpret_cast<const MemBlock*>(&data_[idx]);
            }
 
            // The capacity or num_blocks
            size_t capacity() const {
                return data_.size();
            }
 
            // The number of blocks requested so far
            size_t size() const {
                return size_;
            }
 
            // Start over.
            void clear() {
                size_ = 0;
            }
        };
 
        template<typename ...PointerTArgs>
        typename SpartaSharedPointer<PointerT>::RefCount * allocate_(PointerTArgs&&... args)
        {
            // Return memory allocated here.
            MemBlock * block = nullptr;
 
            // Check for previously freed blocks and reuse them.
            if(free_idx_ > 0) {
                --free_idx_;
                block = free_blocks_[free_idx_];
                sparta_assert(block->ref_count->p != nullptr);
                block->ref_count->mem_block = (void*)block;
                block->ref_count->count = 1;
                // perform an in place new on the reused pointer
                new (block->ref_count->p) PointerT(std::forward<PointerTArgs>(args)...);
            }
            else {
                if(SPARTA_EXPECT_FALSE(allocated_ > water_mark_)) {
                    if(SPARTA_EXPECT_FALSE(!water_mark_warning_)) {
                        watermark_warning_callback_(*this);
                        water_mark_warning_ = true;
                    }
 
                    // Only need to check for overallocation if we've passed the watermark
                    if(SPARTA_EXPECT_FALSE(allocated_ >= memory_blocks_.capacity())) {
                        over_allocation_callback_(*this);
                        SpartaException ex;
                        ex << "This allocator has run out of memory: \n\n\t"
                           << __PRETTY_FUNCTION__
                           << "\n\n"
                           << "\t\tNumber blocks preallocated: " << memory_blocks_.capacity()
                           << "\n\t\tWatermark                 : " << water_mark_;
                        throw ex;
                    }
                }
                block = memory_blocks_.allocate(this, std::forward<PointerTArgs>(args)...);
                ++allocated_;
            }
            return block->ref_count;
        }
 
        // Default watermark warning callback
        static void waterMarkWarningCallback_(const SpartaSharedPointerAllocator & allocator) {
            std::cerr << "WARNING: The watermark for this allocator has been surpassed: \n\n\t"
                      << __PRETTY_FUNCTION__
                      << "\n\n"
                      << "\t\tNumber blocks preallocated: " << allocator.memory_blocks_.capacity()
                      << "\n\t\tWatermark                 : " << allocator.water_mark_ << std::endl;
        }
 
        // Default over allocation callback
        static void overAllocationCallback_(const SpartaSharedPointerAllocator & allocator) {
        }
 
        void releaseObject_(void * block) const noexcept override {
            static_cast<MemBlock *>(block)->ref_count->p->~PointerT();
        }
 
        void releaseBlock_(void * block) override {
            sparta_assert(free_idx_ < free_blocks_.capacity());
            free_blocks_[free_idx_] = static_cast<MemBlock *>(block);
            ++free_idx_;
        }
 
        MemBlockVector           memory_blocks_;
        std::vector<MemBlock*>   free_blocks_;
        size_t                   free_idx_ = 0;
        size_t                   allocated_  = 0;
        const size_t             water_mark_;
        bool                     water_mark_warning_ = false;
        WaterMarkWarningCallback watermark_warning_callback_;
        OverAllocationCallback   over_allocation_callback_;
    };
 
    template<typename PointerT, typename... Args>
    SpartaSharedPointer<PointerT>
    allocate_sparta_shared_pointer(SpartaSharedPointerAllocator<PointerT> & alloc,
                                   Args&&...args)
    {
        static_assert(std::is_constructible<PointerT, Args...>::value,
                      "Can't construct object in allocate_sparta_shared_pointer with the arguments given");
 
        SpartaSharedPointer<PointerT> ptr(alloc.allocate_(std::forward<Args>(args)...));
        return ptr;
    }
 
}