source: trunk/Mars/mcore/MemoryManager.h@ 17248

#include <forward_list>

class MemoryStock
{
    friend class MemoryChunk;
    friend class MemoryManager;

    size_t fChunkSize;
    size_t fMaxMemory;

    size_t fInUse;
    size_t fAllocated;

    size_t fMaxInUse;

    std::mutex fMutexMem;
    std::mutex fMutexCond;
    std::condition_variable fCond;

    std::forward_list<std::shared_ptr<char>> fMemoryStock;

public:
    MemoryStock(size_t chunk, size_t max) : fChunkSize(chunk), fMaxMemory(max),
        fInUse(0), fAllocated(0), fMaxInUse(0)
    {
        if (chunk>max)
            throw std::runtime_error("Size mismatch: Size of a single chunk exceeds maximum size of memory");
    }

private:
    std::shared_ptr<char> pop(bool block)
    {
        if (block)
        {
            // No free slot available, next alloc would exceed max memory:
            // block until a slot is available
            std::unique_lock<std::mutex> lock(fMutexCond);
            while (fMemoryStock.empty() && fAllocated+fChunkSize>fMaxMemory)
                fCond.wait(lock);
        }
        else
        {
            // No free slot available, next alloc would exceed max memory
            // return an empty pointer
            if (fMemoryStock.empty() && fAllocated+fChunkSize>fMaxMemory)
                return std::shared_ptr<char>();
        }

        // We will return this amount of memory
        // This is not 100% thread safe, but it is not a super accurate measure anyway
        fInUse += fChunkSize;
        if (fInUse>fMaxInUse)
            fMaxInUse = fInUse;

        if (fMemoryStock.empty())
        {
            // No free slot available, allocate a new one
            fAllocated += fChunkSize;
            return std::shared_ptr<char>(new char[fChunkSize]);
        }

        // Get the next free slot from the stack and return it
        const std::lock_guard<std::mutex> lock(fMutexMem);

        const auto mem = fMemoryStock.front();
        fMemoryStock.pop_front();
        return mem;
    };

    void push(const std::shared_ptr<char> &mem)
    {
        if (!mem)
            return;

        // Decrease the amont of memory in use accordingly
        fInUse -= fChunkSize;

        // If the maximum memory has changed, we might be over the limit.
        // In this case: free a slot
        if (fAllocated>fMaxMemory)
        {
            fAllocated -= fChunkSize;
            return;
        }

        {
            const std::lock_guard<std::mutex> lock(fMutexMem);
            fMemoryStock.emplace_front(mem);
        }

        {
            const std::lock_guard<std::mutex> lock(fMutexCond);
            fCond.notify_one();
        }
    }
};

class MemoryChunk
{
    friend class MemoryManager;

    std::shared_ptr<MemoryStock> fMemoryStock;
    std::shared_ptr<char>   fPointer;
    char*                   fRealPointer;

    MemoryChunk(const std::shared_ptr<MemoryStock> &mem, bool block) : fMemoryStock(mem)
    {
        fPointer = fMemoryStock->pop(block);
        fRealPointer = fPointer.get();
    }

public:
    ~MemoryChunk()
    {
        fMemoryStock->push(fPointer);
    }

public:
    char* get() { return fPointer.get();}
};

class MemoryManager
{
    std::shared_ptr<MemoryStock> fMemoryStock;

public:

    MemoryManager(size_t chunk, size_t max) : fMemoryStock(std::make_shared<MemoryStock>(chunk, max))
    {
    }

    std::shared_ptr<MemoryChunk> malloc(bool block=true)
    {
        MemoryChunk *chunk = new MemoryChunk(fMemoryStock, block);
        //Etienne cannot get the aliasing to work (at least with g++ 4.4.6
        return std::shared_ptr<MemoryChunk>(std::shared_ptr<MemoryChunk>(chunk));//, chunk->fRealPointer);
    }

    size_t getChunkSize() const { return fMemoryStock->fChunkSize; }
    void   setChunkSize(const size_t size)
    {
        if (getInUse() != 0)
            throw std::runtime_error("Cannot change the chunk size while there is memory in use");

        if (getMaxMemory() < size)
        {
            std::ostringstream str;
            str << "Chunk size(" << size << ") larger thank allowed memory(" << getMaxMemory() << "). Cannot allocate a single chunk.";
            throw std::runtime_error(str.str());
        }
        fMemoryStock->fChunkSize = size;
    }
    size_t getMaxMemory() const { return fMemoryStock->fMaxMemory; }
    size_t getInUse() const { return fMemoryStock->fInUse; }
    size_t getAllocated() const { return fMemoryStock->fAllocated; }
    size_t getMaxInUse() const { return fMemoryStock->fMaxInUse; }
};