aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/clickhouse/src/IO/BrotliWriteBuffer.cpp
blob: 6ec427049c74be0ed9de24f584cc105f1d0f2929 (plain) (blame)
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
#include "clickhouse_config.h"

#if USE_BROTLI
#    include <IO/BrotliWriteBuffer.h>
#    error #include <brotli/encode.h>

namespace DB
{

namespace ErrorCodes
{
    extern const int BROTLI_WRITE_FAILED;
}


class BrotliWriteBuffer::BrotliStateWrapper
{
public:
    BrotliStateWrapper()
    : state(BrotliEncoderCreateInstance(nullptr, nullptr, nullptr))
    {
    }

    ~BrotliStateWrapper()
    {
        BrotliEncoderDestroyInstance(state);
    }

    BrotliEncoderState * state;
};

BrotliWriteBuffer::BrotliWriteBuffer(std::unique_ptr<WriteBuffer> out_, int compression_level, size_t buf_size, char * existing_memory, size_t alignment)
    : WriteBufferWithOwnMemoryDecorator(std::move(out_), buf_size, existing_memory, alignment)
    , brotli(std::make_unique<BrotliStateWrapper>())
    , in_available(0)
    , in_data(nullptr)
    , out_capacity(0)
    , out_data(nullptr)
{
    BrotliEncoderSetParameter(brotli->state, BROTLI_PARAM_QUALITY, static_cast<uint32_t>(compression_level));
    // Set LZ77 window size. According to brotli sources default value is 24 (c/tools/brotli.c:81)
    BrotliEncoderSetParameter(brotli->state, BROTLI_PARAM_LGWIN, 24);
}

BrotliWriteBuffer::~BrotliWriteBuffer() = default;

void BrotliWriteBuffer::nextImpl()
{
    if (!offset())
    {
        return;
    }

    in_data = reinterpret_cast<unsigned char *>(working_buffer.begin());
    in_available = offset();

    try
    {
        do
        {
            out->nextIfAtEnd();
            out_data = reinterpret_cast<unsigned char *>(out->position());
            out_capacity = out->buffer().end() - out->position();

            int result = BrotliEncoderCompressStream(
                    brotli->state,
                    in_available ? BROTLI_OPERATION_PROCESS : BROTLI_OPERATION_FINISH,
                    &in_available,
                    &in_data,
                    &out_capacity,
                    &out_data,
                    nullptr);

            out->position() = out->buffer().end() - out_capacity;

            if (result == 0)
            {
                throw Exception(ErrorCodes::BROTLI_WRITE_FAILED, "brotli compress failed");
            }
        }
        while (in_available > 0);
    }
    catch (...)
    {
        /// Do not try to write next time after exception.
        out->position() = out->buffer().begin();
        throw;
    }
}

void BrotliWriteBuffer::finalizeBefore()
{
    next();

    while (true)
    {
        out->nextIfAtEnd();
        out_data = reinterpret_cast<unsigned char *>(out->position());
        out_capacity = out->buffer().end() - out->position();

        int result = BrotliEncoderCompressStream(
                brotli->state,
                BROTLI_OPERATION_FINISH,
                &in_available,
                &in_data,
                &out_capacity,
                &out_data,
                nullptr);

        out->position() = out->buffer().end() - out_capacity;

        if (BrotliEncoderIsFinished(brotli->state))
        {
            return;
        }

        if (result == 0)
        {
            throw Exception(ErrorCodes::BROTLI_WRITE_FAILED, "brotli compress failed");
        }
    }
}

}

#endif