1 | #include <Common/config.h> |
2 | #if USE_BROTLI |
3 | |
4 | #include <IO/BrotliWriteBuffer.h> |
5 | #include <brotli/encode.h> |
6 | |
7 | namespace DB |
8 | { |
9 | |
10 | namespace ErrorCodes |
11 | { |
12 | extern const int BROTLI_WRITE_FAILED; |
13 | } |
14 | |
15 | |
16 | class BrotliWriteBuffer::BrotliStateWrapper |
17 | { |
18 | public: |
19 | BrotliStateWrapper() |
20 | : state(BrotliEncoderCreateInstance(nullptr, nullptr, nullptr)) |
21 | { |
22 | } |
23 | |
24 | ~BrotliStateWrapper() |
25 | { |
26 | BrotliEncoderDestroyInstance(state); |
27 | } |
28 | |
29 | public: |
30 | BrotliEncoderState * state; |
31 | }; |
32 | |
33 | BrotliWriteBuffer::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 | |
47 | BrotliWriteBuffer::~BrotliWriteBuffer() |
48 | { |
49 | try |
50 | { |
51 | finish(); |
52 | } |
53 | catch (...) |
54 | { |
55 | tryLogCurrentException(__PRETTY_FUNCTION__); |
56 | } |
57 | } |
58 | |
59 | void 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 | |
94 | void 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 | |