1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
|
#include <IO/PeekableWriteBuffer.h>
namespace DB
{
PeekableWriteBuffer::PeekableWriteBuffer(DB::WriteBuffer & sub_buf_) : BufferWithOwnMemory(0), sub_buf(sub_buf_)
{
Buffer & sub_working = sub_buf.buffer();
BufferBase::set(sub_working.begin() + sub_buf.offset(), sub_working.size() - sub_buf.offset(), 0);
}
void PeekableWriteBuffer::nextImpl()
{
if (checkpoint)
{
if (write_to_own_memory)
{
size_t prev_size = position() - memory.data();
size_t new_size = memory.size() * 2;
memory.resize(new_size);
BufferBase::set(memory.data() + prev_size, memory.size() - prev_size, 0);
return;
}
if (memory.size() == 0)
memory.resize(DBMS_DEFAULT_BUFFER_SIZE);
sub_buf.position() = position();
BufferBase::set(memory.data(), memory.size(), 0);
write_to_own_memory = true;
return;
}
sub_buf.position() = position();
sub_buf.next();
BufferBase::set(sub_buf.buffer().begin(), sub_buf.buffer().size(), sub_buf.offset());
}
void PeekableWriteBuffer::dropCheckpoint()
{
assert(checkpoint);
checkpoint = std::nullopt;
/// If we have saved data in own memory, write it to sub-buf.
if (write_to_own_memory)
{
try
{
sub_buf.next();
sub_buf.write(memory.data(), position() - memory.data());
Buffer & sub_working = sub_buf.buffer();
BufferBase::set(sub_working.begin(), sub_working.size(), sub_buf.offset());
write_to_own_memory = false;
}
catch (...)
{
/// If exception happened during writing to sub buffer, we should
/// update buffer to not leave it in invalid state.
Buffer & sub_working = sub_buf.buffer();
BufferBase::set(sub_working.begin(), sub_working.size(), sub_buf.offset());
write_to_own_memory = false;
throw;
}
}
}
void PeekableWriteBuffer::rollbackToCheckpoint(bool drop)
{
assert(checkpoint);
/// Just ignore all data written after checkpoint.
if (write_to_own_memory)
{
Buffer & sub_working = sub_buf.buffer();
BufferBase::set(sub_working.begin(), sub_working.size(), sub_buf.offset());
write_to_own_memory = false;
}
position() = *checkpoint;
if (drop)
checkpoint = std::nullopt;
}
}
|
