tdlight/tdutils/td/utils/buffer.h
levlam 4455b66e8b Fix BufferSlice size statistics.
GitOrigin-RevId: 2dcfeb4981435b93106df5817c849df8a94ba14f
2020-08-05 03:21:10 +03:00

842 lines
22 KiB
C++

//
// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2020
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
#pragma once
#include "td/utils/common.h"
#include "td/utils/port/thread_local.h"
#include "td/utils/Slice.h"
#include <atomic>
#include <limits>
#include <memory>
namespace td {
struct BufferRaw {
explicit BufferRaw(size_t size) : data_size_(size) {
}
size_t data_size_;
// Constant after first reader is created.
// May be change by writer before it.
// So writer may do prepends till there is no reader created.
size_t begin_ = 0;
// Write by writer.
// Read by reader.
std::atomic<size_t> end_{0};
mutable std::atomic<int32> ref_cnt_{1};
std::atomic<bool> has_writer_{true};
bool was_reader_{false};
alignas(4) unsigned char data_[1];
};
class BufferAllocator {
public:
class DeleteWriterPtr {
public:
void operator()(BufferRaw *ptr) {
ptr->has_writer_.store(false, std::memory_order_release);
dec_ref_cnt(ptr);
}
};
class DeleteReaderPtr {
public:
void operator()(BufferRaw *ptr) {
dec_ref_cnt(ptr);
}
};
using WriterPtr = std::unique_ptr<BufferRaw, DeleteWriterPtr>;
using ReaderPtr = std::unique_ptr<BufferRaw, DeleteReaderPtr>;
static WriterPtr create_writer(size_t size);
static WriterPtr create_writer(size_t size, size_t prepend, size_t append);
static ReaderPtr create_reader(size_t size);
static ReaderPtr create_reader(const WriterPtr &raw);
static ReaderPtr create_reader(const ReaderPtr &raw);
static size_t get_buffer_mem();
static int64 get_buffer_slice_size();
static void clear_thread_local();
private:
friend class BufferSlice;
static void track_buffer_slice(int64 size);
static ReaderPtr create_reader_fast(size_t size);
static WriterPtr create_writer_exact(size_t size);
struct BufferRawDeleter {
void operator()(BufferRaw *ptr) {
dec_ref_cnt(ptr);
}
};
struct BufferRawTls {
std::unique_ptr<BufferRaw, BufferRawDeleter> buffer_raw;
};
static TD_THREAD_LOCAL BufferRawTls *buffer_raw_tls;
static void dec_ref_cnt(BufferRaw *ptr);
static BufferRaw *create_buffer_raw(size_t size);
static std::atomic<size_t> buffer_mem;
};
using BufferWriterPtr = BufferAllocator::WriterPtr;
using BufferReaderPtr = BufferAllocator::ReaderPtr;
class BufferSlice {
public:
BufferSlice() = default;
explicit BufferSlice(BufferReaderPtr buffer_ptr) : buffer_(std::move(buffer_ptr)) {
if (is_null()) {
return;
}
begin_ = buffer_->begin_;
end_ = begin_;
sync_with_writer();
}
BufferSlice(BufferReaderPtr buffer_ptr, size_t begin, size_t end)
: buffer_(std::move(buffer_ptr)), begin_(begin), end_(end) {
debug_track();
}
BufferSlice(const BufferSlice &other) = delete;
BufferSlice &operator=(const BufferSlice &other) = delete;
BufferSlice(BufferSlice &&other) : BufferSlice(std::move(other.buffer_), other.begin_, other.end_) {
debug_untrack(); // yes, debug_untrack
}
BufferSlice &operator=(BufferSlice &&other) {
debug_untrack();
buffer_ = std::move(other.buffer_);
begin_ = other.begin_;
end_ = other.end_;
return *this;
}
explicit BufferSlice(size_t size) : buffer_(BufferAllocator::create_reader(size)) {
end_ = buffer_->end_.load(std::memory_order_relaxed);
begin_ = end_ - ((size + 7) & -8);
end_ = begin_ + size;
debug_track();
}
explicit BufferSlice(Slice slice) : BufferSlice(slice.size()) {
as_slice().copy_from(slice);
}
BufferSlice(const char *ptr, size_t size) : BufferSlice(Slice(ptr, size)) {
}
~BufferSlice() {
debug_untrack();
}
void debug_track() const {
BufferAllocator::track_buffer_slice(static_cast<int64>(size()));
}
void debug_untrack() const {
BufferAllocator::track_buffer_slice(-static_cast<int64>(size()));
}
BufferSlice clone() const {
if (is_null()) {
return BufferSlice(BufferReaderPtr(), begin_, end_);
}
return BufferSlice(BufferAllocator::create_reader(buffer_), begin_, end_);
}
BufferSlice copy() const {
if (is_null()) {
return BufferSlice(BufferReaderPtr(), begin_, end_);
}
return BufferSlice(as_slice());
}
Slice as_slice() const {
if (is_null()) {
return Slice();
}
return Slice(buffer_->data_ + begin_, size());
}
operator Slice() const {
return as_slice();
}
MutableSlice as_slice() {
if (is_null()) {
return MutableSlice();
}
return MutableSlice(buffer_->data_ + begin_, size());
}
Slice prepare_read() const {
return as_slice();
}
Slice after(size_t offset) const {
auto full = as_slice();
full.remove_prefix(offset);
return full;
}
bool confirm_read(size_t size) {
debug_untrack();
begin_ += size;
CHECK(begin_ <= end_);
debug_track();
return begin_ == end_;
}
void truncate(size_t limit) {
if (size() > limit) {
debug_untrack();
end_ = begin_ + limit;
debug_track();
}
}
BufferSlice from_slice(Slice slice) const {
auto res = BufferSlice(BufferAllocator::create_reader(buffer_));
res.debug_untrack();
res.begin_ = static_cast<size_t>(slice.ubegin() - buffer_->data_);
res.end_ = static_cast<size_t>(slice.uend() - buffer_->data_);
res.debug_track();
CHECK(buffer_->begin_ <= res.begin_);
CHECK(res.begin_ <= res.end_);
CHECK(res.end_ <= buffer_->end_.load(std::memory_order_relaxed));
return res;
}
// like in std::string
char *data() {
return as_slice().data();
}
const char *data() const {
return as_slice().data();
}
char operator[](size_t at) const {
return as_slice()[at];
}
bool empty() const {
return size() == 0;
}
bool is_null() const {
return !buffer_;
}
size_t size() const {
if (is_null()) {
return 0;
}
return end_ - begin_;
}
// like in std::string
size_t length() const {
return size();
}
// set end_ into writer's end_
size_t sync_with_writer() {
debug_untrack();
CHECK(!is_null());
auto old_end = end_;
end_ = buffer_->end_.load(std::memory_order_acquire);
debug_track();
return end_ - old_end;
}
bool is_writer_alive() const {
CHECK(!is_null());
return buffer_->has_writer_.load(std::memory_order_acquire);
}
void clear() {
debug_untrack();
begin_ = 0;
end_ = 0;
buffer_ = nullptr;
}
private:
BufferReaderPtr buffer_;
size_t begin_ = 0;
size_t end_ = 0;
};
template <class StorerT>
void store(const BufferSlice &buffer_slice, StorerT &storer) {
storer.store_string(buffer_slice);
}
template <class ParserT>
void parse(BufferSlice &buffer_slice, ParserT &parser) {
buffer_slice = parser.template fetch_string<BufferSlice>();
}
class BufferWriter {
public:
BufferWriter() = default;
explicit BufferWriter(size_t size) : BufferWriter(BufferAllocator::create_writer(size)) {
}
BufferWriter(size_t size, size_t prepend, size_t append)
: BufferWriter(BufferAllocator::create_writer(size, prepend, append)) {
}
BufferWriter(Slice slice, size_t prepend, size_t append)
: BufferWriter(BufferAllocator::create_writer(slice.size(), prepend, append)) {
as_slice().copy_from(slice);
}
explicit BufferWriter(BufferWriterPtr buffer_ptr) : buffer_(std::move(buffer_ptr)) {
}
BufferSlice as_buffer_slice() const {
return BufferSlice(BufferAllocator::create_reader(buffer_));
}
bool is_null() const {
return !buffer_;
}
bool empty() const {
return size() == 0;
}
size_t size() const {
if (is_null()) {
return 0;
}
return buffer_->end_.load(std::memory_order_relaxed) - buffer_->begin_;
}
MutableSlice as_slice() {
auto end = buffer_->end_.load(std::memory_order_relaxed);
return MutableSlice(buffer_->data_ + buffer_->begin_, buffer_->data_ + end);
}
Slice as_slice() const {
auto end = buffer_->end_.load(std::memory_order_relaxed);
return Slice(buffer_->data_ + buffer_->begin_, buffer_->data_ + end);
}
MutableSlice prepare_prepend() {
if (is_null()) {
return MutableSlice();
}
CHECK(!buffer_->was_reader_);
return MutableSlice(buffer_->data_, buffer_->begin_);
}
MutableSlice prepare_append() {
if (is_null()) {
return MutableSlice();
}
auto end = buffer_->end_.load(std::memory_order_relaxed);
return MutableSlice(buffer_->data_ + end, buffer_->data_size_ - end);
}
void confirm_append(size_t size) {
if (is_null()) {
CHECK(size == 0);
return;
}
auto new_end = buffer_->end_.load(std::memory_order_relaxed) + size;
CHECK(new_end <= buffer_->data_size_);
buffer_->end_.store(new_end, std::memory_order_release);
}
void confirm_prepend(size_t size) {
if (is_null()) {
CHECK(size == 0);
return;
}
CHECK(buffer_->begin_ >= size);
buffer_->begin_ -= size;
}
private:
BufferWriterPtr buffer_;
};
struct ChainBufferNode {
friend struct DeleteWriterPtr;
struct DeleteWriterPtr {
void operator()(ChainBufferNode *ptr) {
ptr->has_writer_.store(false, std::memory_order_release);
dec_ref_cnt(ptr);
}
};
friend struct DeleteReaderPtr;
struct DeleteReaderPtr {
void operator()(ChainBufferNode *ptr) {
dec_ref_cnt(ptr);
}
};
using WriterPtr = std::unique_ptr<ChainBufferNode, DeleteWriterPtr>;
using ReaderPtr = std::unique_ptr<ChainBufferNode, DeleteReaderPtr>;
static WriterPtr make_writer_ptr(ChainBufferNode *ptr) {
ptr->ref_cnt_.store(1, std::memory_order_relaxed);
ptr->has_writer_.store(true, std::memory_order_relaxed);
return WriterPtr(ptr);
}
static ReaderPtr make_reader_ptr(ChainBufferNode *ptr) {
ptr->ref_cnt_.fetch_add(1, std::memory_order_acq_rel);
return ReaderPtr(ptr);
}
bool has_writer() {
return has_writer_.load(std::memory_order_acquire);
}
bool unique() {
return ref_cnt_.load(std::memory_order_acquire) == 1;
}
ChainBufferNode(BufferSlice slice, bool sync_flag) : slice_(std::move(slice)), sync_flag_(sync_flag) {
}
// reader
// There are two options
// 1. Fixed slice of Buffer
// 2. Slice with non-fixed right end
// In each case slice_ is const. Reader should read it and use sync_with_writer on its own copy.
const BufferSlice slice_;
const bool sync_flag_{false}; // should we call slice_.sync_with_writer or not.
// writer
ReaderPtr next_{nullptr};
private:
std::atomic<int> ref_cnt_{0};
std::atomic<bool> has_writer_{false};
static void clear_nonrecursive(ReaderPtr ptr) {
while (ptr && ptr->unique()) {
ptr = std::move(ptr->next_);
}
}
static void dec_ref_cnt(ChainBufferNode *ptr) {
int left = --ptr->ref_cnt_;
if (left == 0) {
clear_nonrecursive(std::move(ptr->next_));
// TODO(refact): move memory management into allocator (?)
delete ptr;
}
}
};
using ChainBufferNodeWriterPtr = ChainBufferNode::WriterPtr;
using ChainBufferNodeReaderPtr = ChainBufferNode::ReaderPtr;
class ChainBufferNodeAllocator {
public:
static ChainBufferNodeWriterPtr create(BufferSlice slice, bool sync_flag) {
auto *ptr = new ChainBufferNode(std::move(slice), sync_flag);
return ChainBufferNode::make_writer_ptr(ptr);
}
static ChainBufferNodeReaderPtr clone(const ChainBufferNodeReaderPtr &ptr) {
if (!ptr) {
return ChainBufferNodeReaderPtr();
}
return ChainBufferNode::make_reader_ptr(ptr.get());
}
static ChainBufferNodeReaderPtr clone(ChainBufferNodeWriterPtr &ptr) {
if (!ptr) {
return ChainBufferNodeReaderPtr();
}
return ChainBufferNode::make_reader_ptr(ptr.get());
}
};
class ChainBufferIterator {
public:
ChainBufferIterator() = default;
explicit ChainBufferIterator(ChainBufferNodeReaderPtr head) : head_(std::move(head)) {
load_head();
}
ChainBufferIterator clone() const {
return ChainBufferIterator(ChainBufferNodeAllocator::clone(head_), reader_.clone(), need_sync_, offset_);
}
size_t offset() const {
return offset_;
}
void clear() {
*this = ChainBufferIterator();
}
Slice prepare_read() {
if (!head_) {
return Slice();
}
while (true) {
auto res = reader_.prepare_read();
if (!res.empty()) {
return res;
}
auto has_writer = head_->has_writer();
if (need_sync_) {
reader_.sync_with_writer();
res = reader_.prepare_read();
if (!res.empty()) {
return res;
}
}
if (has_writer) {
return Slice();
}
head_ = ChainBufferNodeAllocator::clone(head_->next_);
if (!head_) {
return Slice();
}
load_head();
}
}
// returns only head
BufferSlice read_as_buffer_slice(size_t limit) {
prepare_read();
auto res = reader_.clone();
res.truncate(limit);
confirm_read(res.size());
return res;
}
const BufferSlice &head() const {
return reader_;
}
void confirm_read(size_t size) {
offset_ += size;
reader_.confirm_read(size);
}
void advance_till_end() {
while (true) {
auto ready = prepare_read();
if (ready.empty()) {
break;
}
confirm_read(ready.size());
}
}
size_t advance(size_t offset, MutableSlice dest = MutableSlice()) {
size_t skipped = 0;
while (offset != 0) {
auto ready = prepare_read();
if (ready.empty()) {
break;
}
// read no more than offset
ready.truncate(offset);
offset -= ready.size();
skipped += ready.size();
// copy to dest if possible
auto to_dest_size = min(ready.size(), dest.size());
if (to_dest_size != 0) {
dest.copy_from(ready.substr(0, to_dest_size));
dest.remove_prefix(to_dest_size);
}
confirm_read(ready.size());
}
return skipped;
}
private:
ChainBufferNodeReaderPtr head_;
BufferSlice reader_; // copy of head_->slice_
bool need_sync_ = false; // copy of head_->sync_flag_
size_t offset_ = 0; // position in the union of all nodes
ChainBufferIterator(ChainBufferNodeReaderPtr head, BufferSlice reader, bool need_sync, size_t offset)
: head_(std::move(head)), reader_(std::move(reader)), need_sync_(need_sync), offset_(offset) {
}
void load_head() {
reader_ = head_->slice_.clone();
need_sync_ = head_->sync_flag_;
}
};
class ChainBufferReader {
public:
ChainBufferReader() = default;
explicit ChainBufferReader(ChainBufferNodeReaderPtr head)
: begin_(ChainBufferNodeAllocator::clone(head)), end_(std::move(head)) {
end_.advance_till_end();
}
ChainBufferReader(ChainBufferIterator begin, ChainBufferIterator end, bool sync_flag)
: begin_(std::move(begin)), end_(std::move(end)), sync_flag_(sync_flag) {
}
ChainBufferReader(ChainBufferNodeReaderPtr head, size_t size)
: begin_(ChainBufferNodeAllocator::clone(head)), end_(std::move(head)) {
auto advanced = end_.advance(size);
CHECK(advanced == size);
}
ChainBufferReader(ChainBufferReader &&) = default;
ChainBufferReader &operator=(ChainBufferReader &&) = default;
ChainBufferReader(const ChainBufferReader &) = delete;
ChainBufferReader &operator=(const ChainBufferReader &) = delete;
~ChainBufferReader() = default;
ChainBufferReader clone() {
return ChainBufferReader(begin_.clone(), end_.clone(), sync_flag_);
}
Slice prepare_read() {
auto res = begin_.prepare_read();
res.truncate(size());
return res;
}
void confirm_read(size_t size) {
CHECK(size <= this->size());
begin_.confirm_read(size);
}
size_t advance(size_t offset, MutableSlice dest = MutableSlice()) {
CHECK(offset <= size());
return begin_.advance(offset, dest);
}
size_t size() const {
return end_.offset() - begin_.offset();
}
bool empty() const {
return size() == 0;
}
void sync_with_writer() {
if (sync_flag_) {
end_.advance_till_end();
}
}
void advance_end(size_t size) {
end_.advance(size);
}
const ChainBufferIterator &begin() {
return begin_;
}
const ChainBufferIterator &end() {
return end_;
}
// Return [begin_, tail.begin_)
// *this = tail
ChainBufferReader cut_head(ChainBufferIterator pos) TD_WARN_UNUSED_RESULT {
auto tmp = begin_.clone();
begin_ = pos.clone();
return ChainBufferReader(std::move(tmp), std::move(pos), false);
}
ChainBufferReader cut_head(size_t offset) TD_WARN_UNUSED_RESULT {
CHECK(offset <= size());
auto it = begin_.clone();
it.advance(offset);
return cut_head(std::move(it));
}
BufferSlice move_as_buffer_slice() {
BufferSlice res;
if (begin_.head().size() >= size()) {
res = begin_.read_as_buffer_slice(size());
} else {
auto save_size = size();
res = BufferSlice{save_size};
advance(save_size, res.as_slice());
}
*this = ChainBufferReader();
return res;
}
BufferSlice read_as_buffer_slice(size_t limit = std::numeric_limits<size_t>::max()) {
return begin_.read_as_buffer_slice(min(limit, size()));
}
private:
ChainBufferIterator begin_; // use it for prepare_read. Fix result with size()
ChainBufferIterator end_; // keep end as far as we can. use it for size()
bool sync_flag_ = true; // auto sync of end_
// 1. We have fixed size. Than end_ is useless.
// 2. No fixed size. One has to sync end_ with end_.advance_till_end() in order to calculate size.
};
class ChainBufferWriter {
public:
ChainBufferWriter() {
init();
}
void init(size_t size = 0) {
writer_ = BufferWriter(size);
tail_ = ChainBufferNodeAllocator::create(writer_.as_buffer_slice(), true);
head_ = ChainBufferNodeAllocator::clone(tail_);
}
MutableSlice prepare_append(size_t hint = 0) {
CHECK(!empty());
auto res = prepare_append_inplace();
if (res.empty()) {
return prepare_append_alloc(hint);
}
return res;
}
MutableSlice prepare_append_at_least(size_t size) {
CHECK(!empty());
auto res = prepare_append_inplace();
if (res.size() < size) {
return prepare_append_alloc(size);
}
return res;
}
MutableSlice prepare_append_inplace() {
CHECK(!empty());
return writer_.prepare_append();
}
MutableSlice prepare_append_alloc(size_t hint = 0) {
CHECK(!empty());
if (hint < (1 << 10)) {
hint = 1 << 12;
}
BufferWriter new_writer(hint);
auto new_tail = ChainBufferNodeAllocator::create(new_writer.as_buffer_slice(), true);
tail_->next_ = ChainBufferNodeAllocator::clone(new_tail);
writer_ = std::move(new_writer);
tail_ = std::move(new_tail); // release tail_
return writer_.prepare_append();
}
void confirm_append(size_t size) {
CHECK(!empty());
writer_.confirm_append(size);
}
void append(Slice slice, size_t hint = 0) {
while (!slice.empty()) {
auto ready = prepare_append(td::max(slice.size(), hint));
auto shift = min(ready.size(), slice.size());
ready.copy_from(slice.substr(0, shift));
confirm_append(shift);
slice.remove_prefix(shift);
}
}
void append(BufferSlice slice) {
auto ready = prepare_append_inplace();
// TODO(perf): we have to store some stats in ChainBufferWriter
// for better append logic
if (slice.size() < (1 << 8) || ready.size() >= slice.size()) {
return append(slice.as_slice());
}
auto new_tail = ChainBufferNodeAllocator::create(std::move(slice), false);
tail_->next_ = ChainBufferNodeAllocator::clone(new_tail);
writer_ = BufferWriter();
tail_ = std::move(new_tail); // release tail_
}
void append(ChainBufferReader &&reader) {
while (!reader.empty()) {
append(reader.read_as_buffer_slice());
}
}
void append(ChainBufferReader &reader) {
while (!reader.empty()) {
append(reader.read_as_buffer_slice());
}
}
ChainBufferReader extract_reader() {
CHECK(head_);
return ChainBufferReader(std::move(head_));
}
private:
bool empty() const {
return !tail_;
}
ChainBufferNodeReaderPtr head_;
ChainBufferNodeWriterPtr tail_;
BufferWriter writer_;
};
class BufferBuilder {
public:
BufferBuilder() = default;
BufferBuilder(Slice slice, size_t prepend_size, size_t append_size)
: buffer_writer_(slice, prepend_size, append_size) {
}
explicit BufferBuilder(BufferWriter &&buffer_writer) : buffer_writer_(std::move(buffer_writer)) {
}
void append(BufferSlice slice);
void append(Slice slice);
void prepend(BufferSlice slice);
void prepend(Slice slice);
template <class F>
void for_each(F &&f) const & {
for (auto i = to_prepend_.size(); i > 0; i--) {
f(to_prepend_[i - 1].as_slice());
}
if (!buffer_writer_.empty()) {
f(buffer_writer_.as_slice());
}
for (auto &slice : to_append_) {
f(slice.as_slice());
}
}
template <class F>
void for_each(F &&f) && {
for (auto i = to_prepend_.size(); i > 0; i--) {
f(std::move(to_prepend_[i - 1]));
}
if (!buffer_writer_.empty()) {
f(buffer_writer_.as_buffer_slice());
}
for (auto &slice : to_append_) {
f(std::move(slice));
}
}
size_t size() const;
BufferSlice extract();
private:
BufferWriter buffer_writer_;
std::vector<BufferSlice> to_append_;
std::vector<BufferSlice> to_prepend_;
bool append_inplace(Slice slice);
void append_slow(BufferSlice slice);
bool prepend_inplace(Slice slice);
void prepend_slow(BufferSlice slice);
};
inline Slice as_slice(const BufferSlice &value) {
return value.as_slice();
}
inline MutableSlice as_slice(BufferSlice &value) {
return value.as_slice();
}
inline MutableSlice as_mutable_slice(BufferSlice &value) {
return value.as_slice();
}
} // namespace td