1#include <Common/config.h>
2#if USE_BROTLI
3
4#include <IO/BrotliWriteBuffer.h>
5#include <brotli/encode.h>
6
7namespace DB
8{
9
10namespace ErrorCodes
11{
12 extern const int BROTLI_WRITE_FAILED;
13}
14
15
16class BrotliWriteBuffer::BrotliStateWrapper
17{
18public:
19 BrotliStateWrapper()
20 : state(BrotliEncoderCreateInstance(nullptr, nullptr, nullptr))
21 {
22 }
23
24 ~BrotliStateWrapper()
25 {
26 BrotliEncoderDestroyInstance(state);
27 }
28
29public:
30 BrotliEncoderState * state;
31};
32
33BrotliWriteBuffer::BrotliWriteBuffer(WriteBuffer & out_, int compression_level, size_t buf_size, char * existing_memory, size_t alignment)
34 : BufferWithOwnMemory<WriteBuffer>(buf_size, existing_memory, alignment)
35 , brotli(std::make_unique<BrotliStateWrapper>())
36 , in_available(0)
37 , in_data(nullptr)
38 , out_capacity(0)
39 , out_data(nullptr)
40 , out(out_)
41{
42 BrotliEncoderSetParameter(brotli->state, BROTLI_PARAM_QUALITY, static_cast<uint32_t>(compression_level));
43 // Set LZ77 window size. According to brotli sources default value is 24 (c/tools/brotli.c:81)
44 BrotliEncoderSetParameter(brotli->state, BROTLI_PARAM_LGWIN, 24);
45}
46
47BrotliWriteBuffer::~BrotliWriteBuffer()
48{
49 try
50 {
51 finish();
52 }
53 catch (...)
54 {
55 tryLogCurrentException(__PRETTY_FUNCTION__);
56 }
57}
58
59void BrotliWriteBuffer::nextImpl()
60{
61 if (!offset())
62 {
63 return;
64 }
65
66 in_data = reinterpret_cast<unsigned char *>(working_buffer.begin());
67 in_available = offset();
68
69 do
70 {
71 out.nextIfAtEnd();
72 out_data = reinterpret_cast<unsigned char *>(out.position());
73 out_capacity = out.buffer().end() - out.position();
74
75 int result = BrotliEncoderCompressStream(
76 brotli->state,
77 in_available ? BROTLI_OPERATION_PROCESS : BROTLI_OPERATION_FINISH,
78 &in_available,
79 &in_data,
80 &out_capacity,
81 &out_data,
82 nullptr);
83
84 out.position() = out.buffer().end() - out_capacity;
85
86 if (result == 0)
87 {
88 throw Exception("brotli compress failed", ErrorCodes::BROTLI_WRITE_FAILED);
89 }
90 }
91 while (in_available > 0 || out_capacity == 0);
92}
93
94void BrotliWriteBuffer::finish()
95{
96 if (finished)
97 return;
98
99 next();
100
101 while (true)
102 {
103 out.nextIfAtEnd();
104 out_data = reinterpret_cast<unsigned char *>(out.position());
105 out_capacity = out.buffer().end() - out.position();
106
107 int result = BrotliEncoderCompressStream(
108 brotli->state,
109 BROTLI_OPERATION_FINISH,
110 &in_available,
111 &in_data,
112 &out_capacity,
113 &out_data,
114 nullptr);
115
116 out.position() = out.buffer().end() - out_capacity;
117
118 if (BrotliEncoderIsFinished(brotli->state))
119 {
120 finished = true;
121 return;
122 }
123
124 if (result == 0)
125 {
126 throw Exception("brotli compress failed", ErrorCodes::BROTLI_WRITE_FAILED);
127 }
128 }
129}
130
131}
132
133#endif
134