ClickHouse/dbms/src/IO/MemoryReadWriteBuffer.cpp

#include <DB/IO/MemoryReadWriteBuffer.h>
#include <DB/Common/Allocator.h>
#include <common/likely.h>

namespace DB
{

namespace ErrorCodes
{
	extern const int CURRENT_WRITE_BUFFER_IS_EXHAUSTED;
}


class ReadBufferFromMemoryWriteBuffer : public ReadBuffer
{
public:

	ReadBufferFromMemoryWriteBuffer(MemoryWriteBuffer && origin)
	:
	ReadBuffer(nullptr, 0),
	chunk_list(std::move(origin.chunk_list)),
	end_pos(origin.position())
	{
		chunk_head = chunk_list.begin();
		setChunk();
	}

	bool nextImpl() override
	{
		if (chunk_head == chunk_list.end())
			return false;

		++chunk_head;
		return setChunk();
	}

	~ReadBufferFromMemoryWriteBuffer()
	{
		for (const auto & range : chunk_list)
			Allocator<false>().free(range.begin(), range.size());
	}

private:

	/// update buffers and position according to chunk_head pointer
	bool setChunk()
	{
		if (chunk_head != chunk_list.end())
		{
			internalBuffer() = *chunk_head;

			auto next_chunk = chunk_head;
			++next_chunk;

			/// It is last chunk, it should be truncated
			if (next_chunk != chunk_list.end())
				buffer() = internalBuffer();
			else
				buffer() = Buffer(internalBuffer().begin(), end_pos);

			position() = buffer().begin();
		}
		else
		{
			buffer() = internalBuffer() = Buffer(nullptr, nullptr);
			position() = nullptr;
		}

		return buffer().size() != 0;
	}

	using Container = std::forward_list<BufferBase::Buffer>;

	Container chunk_list;
	Container::iterator chunk_head;
	Position end_pos;
};

MemoryWriteBuffer::MemoryWriteBuffer(size_t max_total_size_, size_t initial_chunk_size_, double growth_rate_)
: WriteBuffer(nullptr, 0), max_total_size(max_total_size_), initial_chunk_size(initial_chunk_size_), growth_rate(growth_rate_)
{
	addChunk();
}

void MemoryWriteBuffer::nextImpl()
{
	if (unlikely(hasPendingData()))
	{
		/// ignore flush
		buffer() = Buffer(pos, buffer().end());
		return;
	}

	addChunk();
}

void MemoryWriteBuffer::addChunk()
{
	size_t next_chunk_size;
	if (chunk_list.empty())
	{
		chunk_tail = chunk_list.before_begin();
		next_chunk_size = initial_chunk_size;
	}
	else
	{
		next_chunk_size = std::max(1ul, static_cast<size_t>(chunk_tail->size() * growth_rate));
	}

	if (max_total_size)
	{
		if (total_chunks_size + next_chunk_size > max_total_size)
			next_chunk_size = max_total_size - total_chunks_size;

		if (0 == next_chunk_size)
			throw Exception("MemoryWriteBuffer limit is exhausted", ErrorCodes::CURRENT_WRITE_BUFFER_IS_EXHAUSTED);
	}

	Position begin = reinterpret_cast<Position>(Allocator<false>().alloc(next_chunk_size));
	chunk_tail = chunk_list.emplace_after(chunk_tail, begin, begin + next_chunk_size);
	total_chunks_size += next_chunk_size;

	set(chunk_tail->begin(), chunk_tail->size());
}

std::shared_ptr<ReadBuffer> MemoryWriteBuffer::getReadBuffer()
{
	auto res = std::make_shared<ReadBufferFromMemoryWriteBuffer>(std::move(*this));

	/// invalidate members
	chunk_list.clear();
	chunk_tail = chunk_list.begin();

	return res;
}

MemoryWriteBuffer::~MemoryWriteBuffer()
{
	for (const auto & range : chunk_list)
		Allocator<false>().free(range.begin(), range.size());
}

}