ClickHouse/src/IO/LimitReadBuffer.cpp at master · StackRadiusForks/ClickHouse · GitHub

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
87
88
89
90
91
92
93
94
95
96
#include <IO/LimitReadBuffer.h>
#include <Common/Exception.h>
#include <Core/Settings.h>


namespace DB
{

namespace ErrorCodes
{
    extern const int LIMIT_EXCEEDED;
    extern const int CANNOT_READ_ALL_DATA;
}


bool LimitReadBuffer::nextImpl()
{
    chassert(position() >= in->position());

    /// Let underlying buffer calculate read bytes in `next()` call.
    in->position() = position();

    if (bytes >= settings.read_no_less)
    {
        if (settings.expect_eof && bytes > settings.read_no_more)
            throw Exception(ErrorCodes::LIMIT_EXCEEDED, "Limit for LimitReadBuffer exceeded: {}", settings.excetion_hint);

        if (bytes >= settings.read_no_more)
            return false;

        //throw Exception(ErrorCodes::CANNOT_READ_ALL_DATA, "Unexpected data, got {} bytes, expected {}", bytes, settings.read_atmost);
    }

    if (!in->next())
    {
        if (bytes < settings.read_no_less)
            throw Exception(ErrorCodes::CANNOT_READ_ALL_DATA, "Unexpected EOF, got {} of {} bytes", bytes, settings.read_no_less);

        /// Clearing the buffer with existing data.
        BufferBase::set(in->position(), 0, 0);

        return false;
    }

    BufferBase::set(in->position(), std::min(in->available(), getEffectiveBufferSize() - bytes), 0);

    return true;
}


LimitReadBuffer::LimitReadBuffer(ReadBuffer & in_, Settings settings_)
    : ReadBuffer(in_.position(), 0)
    , in(&in_)
    , settings(std::move(settings_))
{
    chassert(in);
    chassert(settings.read_no_less <= settings.read_no_more);

    BufferBase::set(in->position(), std::min(in->available(), getEffectiveBufferSize()), 0);
}


LimitReadBuffer::LimitReadBuffer(std::unique_ptr<ReadBuffer> in_, Settings settings_)
    : LimitReadBuffer(*in_, std::move(settings_))
{
    holder = std::move(in_);
}


LimitReadBuffer::~LimitReadBuffer()
{
    /// Update underlying buffer's position in case when limit wasn't reached.
    if (!working_buffer.empty())
        in->position() = position();
}

bool LimitReadBuffer::poll(size_t timeout_microseconds)
{
    if (hasPendingData())
        return true;

    in->position() = position();

    if (bytes >= getEffectiveBufferSize())
        return true;

    return in->poll(timeout_microseconds);
}

size_t LimitReadBuffer::getEffectiveBufferSize() const
{
    if (settings.read_no_less)
        return settings.read_no_less;
    return settings.read_no_more;
}
}