You've already forked mariadb-columnstore-engine
mirror of
https://github.com/mariadb-corporation/mariadb-columnstore-engine.git
synced 2025-08-01 06:46:55 +03:00
feat(PP,ByteStream): new counting memory allocator
This commit is contained in:
91
utils/common/countingallocator.h
Normal file
91
utils/common/countingallocator.h
Normal file
@ -0,0 +1,91 @@
|
||||
/* Copyright (C) 2024 MariaDB Corporation
|
||||
|
||||
This program is free software; you can redistribute it and/or
|
||||
modify it under the terms of the GNU General Public License
|
||||
as published by the Free Software Foundation; version 2 of
|
||||
the License.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program; if not, write to the Free Software
|
||||
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
|
||||
MA 02110-1301, USA. */
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
#include <memory>
|
||||
#include <atomic>
|
||||
#include <cstddef>
|
||||
#include <iostream>
|
||||
#include <utility>
|
||||
|
||||
namespace allocators
|
||||
{
|
||||
|
||||
// const constexpr std::uint64_t CounterUpdateUnitSize = 4 * 1024 * 1024;
|
||||
const constexpr std::int64_t MemoryLimitLowerBound = 100 * 1024 * 1024; // WIP
|
||||
|
||||
// Custom Allocator that tracks allocated memory using an atomic counter
|
||||
template <typename T>
|
||||
class CountingAllocator {
|
||||
public:
|
||||
using value_type = T;
|
||||
|
||||
// Constructor accepting a reference to an atomic counter
|
||||
explicit CountingAllocator(std::atomic<int64_t>& memoryLimit, const uint64_t lowerBound = MemoryLimitLowerBound) noexcept
|
||||
: memoryLimitRef_(memoryLimit), memoryLimitLowerBound(lowerBound) {}
|
||||
|
||||
// Copy constructor (template to allow conversion between different types)
|
||||
template <typename U>
|
||||
CountingAllocator(const CountingAllocator<U>& other) noexcept
|
||||
: memoryLimitRef_(other.memoryLimitRef_) {}
|
||||
|
||||
// Allocate memory for n objects of type T
|
||||
T* allocate(std::size_t n) {
|
||||
auto memCounted = memoryLimitRef_.fetch_sub(n * sizeof(T), std::memory_order_relaxed);
|
||||
if (memCounted < memoryLimitLowerBound) {
|
||||
memoryLimitRef_.fetch_add(n * sizeof(T), std::memory_order_relaxed);
|
||||
throw std::bad_alloc();
|
||||
}
|
||||
|
||||
T* ptr = static_cast<T*>(::operator new(n * sizeof(T)));
|
||||
// std::cout << "[Allocate] " << n * sizeof(T) << " bytes at " << static_cast<void*>(ptr)
|
||||
// << ". current timit: " << std::dec << memoryLimitRef_.load() << std::hex << " bytes.\n";
|
||||
return ptr;
|
||||
}
|
||||
|
||||
// Deallocate memory for n objects of type T
|
||||
void deallocate(T* ptr, std::size_t n) noexcept {
|
||||
::operator delete(ptr);
|
||||
memoryLimitRef_.fetch_add(n * sizeof(T), std::memory_order_relaxed);
|
||||
// std::cout << "[Deallocate] " << n * sizeof(T) << " bytes from " << static_cast<void*>(ptr)
|
||||
// << ". current timit: " << std::dec << memoryLimitRef_.load() << std::hex << " bytes.\n";
|
||||
}
|
||||
|
||||
// Equality operators (allocators are equal if they share the same counter)
|
||||
template <typename U>
|
||||
bool operator==(const CountingAllocator<U>& other) const noexcept {
|
||||
return &memoryLimitRef_ == &other.memoryLimitRef_;
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
bool operator!=(const CountingAllocator<U>& other) const noexcept {
|
||||
return !(*this == other);
|
||||
}
|
||||
|
||||
private:
|
||||
std::atomic<int64_t>& memoryLimitRef_;
|
||||
int64_t memoryLimitLowerBound = 0;
|
||||
|
||||
// Grant access to other instances of CountingAllocator with different types
|
||||
template <typename U>
|
||||
friend class CountingAllocator;
|
||||
};
|
||||
|
||||
} // namespace allocators
|
@ -50,8 +50,8 @@ void ByteStream::doCopy(const ByteStream& rhs)
|
||||
|
||||
if (fMaxLen < rlen)
|
||||
{
|
||||
delete[] fBuf;
|
||||
fBuf = new uint8_t[rlen + ISSOverhead];
|
||||
deallocate(fBuf);
|
||||
fBuf = allocate(rlen + ISSOverhead);
|
||||
fMaxLen = rlen;
|
||||
}
|
||||
|
||||
@ -83,7 +83,7 @@ ByteStream& ByteStream::operator=(const ByteStream& rhs)
|
||||
doCopy(rhs);
|
||||
else
|
||||
{
|
||||
delete[] fBuf;
|
||||
deallocate(fBuf);
|
||||
fBuf = fCurInPtr = fCurOutPtr = 0;
|
||||
fMaxLen = 0;
|
||||
// Clear `longStrings`.
|
||||
@ -100,6 +100,13 @@ ByteStream::ByteStream(BSSizeType initSize) : fBuf(0), fCurInPtr(0), fCurOutPtr(
|
||||
growBuf(initSize);
|
||||
}
|
||||
|
||||
ByteStream::ByteStream(allocators::CountingAllocator<uint8_t>* allocator, uint32_t initSize)
|
||||
: fBuf(0), fCurInPtr(0), fCurOutPtr(0), fMaxLen(0), allocator(allocator)
|
||||
{
|
||||
if (initSize > 0)
|
||||
growBuf(initSize);
|
||||
}
|
||||
|
||||
void ByteStream::add(const uint8_t b)
|
||||
{
|
||||
if (fBuf == 0 || (static_cast<BSSizeType>(fCurInPtr - fBuf) == fMaxLen + ISSOverhead))
|
||||
@ -108,6 +115,26 @@ void ByteStream::add(const uint8_t b)
|
||||
*fCurInPtr++ = b;
|
||||
}
|
||||
|
||||
BSBufType* ByteStream::allocate(const size_t size)
|
||||
{
|
||||
if (allocator)
|
||||
{
|
||||
auto* mem = allocator->allocate(size);
|
||||
return new (mem) BSBufType[size];
|
||||
}
|
||||
return new BSBufType[size];
|
||||
}
|
||||
|
||||
void ByteStream::deallocate(BSBufType* ptr)
|
||||
{
|
||||
if (allocator)
|
||||
{
|
||||
size_t count = (fMaxLen) ? fMaxLen + ISSOverhead : 0;
|
||||
return allocator->deallocate(ptr, count);
|
||||
}
|
||||
return delete[] fBuf;
|
||||
}
|
||||
|
||||
void ByteStream::growBuf(BSSizeType toSize)
|
||||
{
|
||||
if (fBuf == 0)
|
||||
@ -117,7 +144,7 @@ void ByteStream::growBuf(BSSizeType toSize)
|
||||
else
|
||||
toSize = ((toSize + BlockSize - 1) / BlockSize) * BlockSize;
|
||||
|
||||
fBuf = new uint8_t[toSize + ISSOverhead];
|
||||
fBuf = allocate(toSize + ISSOverhead);
|
||||
#ifdef ZERO_ON_NEW
|
||||
memset(fBuf, 0, (toSize + ISSOverhead));
|
||||
#endif
|
||||
@ -137,14 +164,14 @@ void ByteStream::growBuf(BSSizeType toSize)
|
||||
// Make sure we at least double the allocation
|
||||
toSize = std::max(toSize, fMaxLen * 2);
|
||||
|
||||
uint8_t* t = new uint8_t[toSize + ISSOverhead];
|
||||
BSBufType* t = allocate(toSize + ISSOverhead);
|
||||
BSSizeType curOutOff = fCurOutPtr - fBuf;
|
||||
BSSizeType curInOff = fCurInPtr - fBuf;
|
||||
memcpy(t, fBuf, fCurInPtr - fBuf);
|
||||
#ifdef ZERO_ON_NEW
|
||||
memset(t + (fCurInPtr - fBuf), 0, (toSize + ISSOverhead) - (fCurInPtr - fBuf));
|
||||
#endif
|
||||
delete[] fBuf;
|
||||
deallocate(fBuf);
|
||||
fBuf = t;
|
||||
fMaxLen = toSize;
|
||||
fCurInPtr = fBuf + curInOff;
|
||||
@ -541,8 +568,8 @@ void ByteStream::load(const uint8_t* bp, BSSizeType len)
|
||||
|
||||
if (len > fMaxLen)
|
||||
{
|
||||
delete[] fBuf;
|
||||
fBuf = new uint8_t[newMaxLen + ISSOverhead];
|
||||
deallocate(fBuf);
|
||||
fBuf = allocate(newMaxLen + ISSOverhead);
|
||||
fMaxLen = newMaxLen;
|
||||
}
|
||||
|
||||
@ -575,8 +602,10 @@ void ByteStream::swap(ByteStream& rhs)
|
||||
std::swap(fCurOutPtr, rhs.fCurOutPtr);
|
||||
std::swap(fMaxLen, rhs.fMaxLen);
|
||||
std::swap(longStrings, rhs.longStrings);
|
||||
std::swap(allocator, rhs.allocator);
|
||||
}
|
||||
|
||||
// WIP use allocator
|
||||
ifstream& operator>>(ifstream& ifs, ByteStream& bs)
|
||||
{
|
||||
int ifs_len;
|
||||
@ -653,7 +682,6 @@ void ByteStream::needAtLeast(BSSizeType amount)
|
||||
growBuf(fMaxLen + amount);
|
||||
}
|
||||
|
||||
|
||||
ByteStream& ByteStream::operator<<(const ByteStream& bs)
|
||||
{
|
||||
BSSizeType len = bs.length();
|
||||
|
@ -37,6 +37,7 @@
|
||||
#include "serializeable.h"
|
||||
#include "any.hpp"
|
||||
#include "nullstring.h"
|
||||
#include "countingallocator.h"
|
||||
|
||||
class ByteStreamTestSuite;
|
||||
|
||||
@ -46,7 +47,7 @@ namespace messageqcpp
|
||||
{
|
||||
typedef boost::shared_ptr<ByteStream> SBS;
|
||||
using BSSizeType = uint64_t;
|
||||
|
||||
using BSBufType = uint8_t;
|
||||
/**
|
||||
* @brief A class to marshall bytes as a stream
|
||||
*
|
||||
@ -78,6 +79,7 @@ class ByteStream : public Serializeable
|
||||
* default ctor
|
||||
*/
|
||||
EXPORT explicit ByteStream(BSSizeType initSize = 8192); // multiples of pagesize are best
|
||||
explicit ByteStream(allocators::CountingAllocator<BSBufType>* alloc, uint32_t initSize = 8192);
|
||||
/**
|
||||
* ctor with a uint8_t array and len initializer
|
||||
*/
|
||||
@ -466,6 +468,9 @@ class ByteStream : public Serializeable
|
||||
void doCopy(const ByteStream& rhs);
|
||||
|
||||
private:
|
||||
BSBufType* allocate(const size_t size);
|
||||
void deallocate(BSBufType* ptr);
|
||||
|
||||
// Put struct `MemChunk` declaration here, to avoid circular dependency.
|
||||
struct MemChunk
|
||||
{
|
||||
@ -474,11 +479,13 @@ class ByteStream : public Serializeable
|
||||
uint8_t data[];
|
||||
};
|
||||
|
||||
uint8_t* fBuf; /// the start of the allocated buffer
|
||||
uint8_t* fCurInPtr; // the point in fBuf where data is inserted next
|
||||
uint8_t* fCurOutPtr; // the point in fBuf where data is extracted from next
|
||||
BSSizeType fMaxLen; // how big fBuf is currently
|
||||
std::vector<std::shared_ptr<uint8_t[]>> longStrings; // Stores `long strings`.
|
||||
BSBufType* fBuf; /// the start of the allocated buffer
|
||||
BSBufType* fCurInPtr; // the point in fBuf where data is inserted next
|
||||
BSBufType* fCurOutPtr; // the point in fBuf where data is extracted from next
|
||||
BSSizeType fMaxLen; // how big fBuf is currently
|
||||
// Stores `long strings`.
|
||||
std::vector<std::shared_ptr<uint8_t[]>> longStrings;
|
||||
allocators::CountingAllocator<BSBufType>* allocator = nullptr;
|
||||
};
|
||||
|
||||
template <int W, typename T = void>
|
||||
@ -533,7 +540,7 @@ inline ByteStream::ByteStream(const uint8_t* bp, BSSizeType len) : fBuf(nullptr)
|
||||
}
|
||||
inline ByteStream::~ByteStream()
|
||||
{
|
||||
delete[] fBuf;
|
||||
deallocate(fBuf);
|
||||
}
|
||||
|
||||
inline const uint8_t* ByteStream::buf() const
|
||||
@ -558,7 +565,7 @@ inline BSSizeType ByteStream::lengthWithHdrOverhead() const
|
||||
}
|
||||
inline void ByteStream::reset()
|
||||
{
|
||||
delete[] fBuf;
|
||||
deallocate(fBuf);
|
||||
fMaxLen = 0;
|
||||
fCurInPtr = fCurOutPtr = fBuf = nullptr;
|
||||
}
|
||||
|
@ -103,7 +103,7 @@ const SBS CompressedInetStreamSocket::read(const struct timespec* timeout, bool*
|
||||
uint32_t storedLen = *(uint32_t*)readBS->buf();
|
||||
|
||||
if (!storedLen)
|
||||
return SBS(new ByteStream(0));
|
||||
return SBS(new ByteStream(0U));
|
||||
|
||||
uncompressedSize = storedLen;
|
||||
ret.reset(new ByteStream(uncompressedSize));
|
||||
|
@ -487,25 +487,25 @@ const SBS InetStreamSocket::read(const struct ::timespec* timeout, bool* isTimeO
|
||||
// {
|
||||
// logIoError("InetStreamSocket::read: timeout during readToMagic", 0);
|
||||
// }
|
||||
return SBS(new ByteStream(0));
|
||||
return SBS(new ByteStream(0U));
|
||||
}
|
||||
|
||||
// we need to read the 4-byte message length first.
|
||||
uint32_t msglen;
|
||||
if (!readFixedSizeData(pfd, reinterpret_cast<uint8_t*>(&msglen), sizeof(msglen), timeout, isTimeOut, stats,
|
||||
msecs))
|
||||
return SBS(new ByteStream(0));
|
||||
return SBS(new ByteStream(0U));
|
||||
|
||||
// Read the number of the `long strings`.
|
||||
uint32_t longStringSize;
|
||||
if (!readFixedSizeData(pfd, reinterpret_cast<uint8_t*>(&longStringSize), sizeof(longStringSize), timeout,
|
||||
isTimeOut, stats, msecs))
|
||||
return SBS(new ByteStream(0));
|
||||
return SBS(new ByteStream(0U));
|
||||
|
||||
// Read the actual data of the `ByteStream`.
|
||||
SBS res(new ByteStream(msglen));
|
||||
if (!readFixedSizeData(pfd, res->getInputPtr(), msglen, timeout, isTimeOut, stats, msecs))
|
||||
return SBS(new ByteStream(0));
|
||||
return SBS(new ByteStream(0U));
|
||||
res->advanceInputPtr(msglen);
|
||||
|
||||
std::vector<std::shared_ptr<uint8_t[]>> longStrings;
|
||||
@ -517,7 +517,7 @@ const SBS InetStreamSocket::read(const struct ::timespec* timeout, bool* isTimeO
|
||||
rowgroup::StringStore::MemChunk memChunk;
|
||||
if (!readFixedSizeData(pfd, reinterpret_cast<uint8_t*>(&memChunk),
|
||||
sizeof(rowgroup::StringStore::MemChunk), timeout, isTimeOut, stats, msecs))
|
||||
return SBS(new ByteStream(0));
|
||||
return SBS(new ByteStream(0U));
|
||||
|
||||
// Allocate new memory for the `long string`.
|
||||
std::shared_ptr<uint8_t[]> longString(
|
||||
@ -532,7 +532,7 @@ const SBS InetStreamSocket::read(const struct ::timespec* timeout, bool* isTimeO
|
||||
// Read the `long string`.
|
||||
if (!readFixedSizeData(pfd, memChunkPointer->data, memChunkPointer->currentSize, timeout, isTimeOut,
|
||||
stats, msecs))
|
||||
return SBS(new ByteStream(0));
|
||||
return SBS(new ByteStream(0U));
|
||||
|
||||
longStrings.push_back(longString);
|
||||
}
|
||||
@ -540,7 +540,7 @@ const SBS InetStreamSocket::read(const struct ::timespec* timeout, bool* isTimeO
|
||||
catch (std::bad_alloc& exception)
|
||||
{
|
||||
logIoError("InetStreamSocket::read: error during read for 'long strings' - 'bad_alloc'", 0);
|
||||
return SBS(new ByteStream(0));
|
||||
return SBS(new ByteStream(0U));
|
||||
}
|
||||
catch (std::exception& exception)
|
||||
{
|
||||
|
Reference in New Issue
Block a user