1 | #include <IO/ZlibDeflatingWriteBuffer.h> |
2 | #include <Common/MemorySanitizer.h> |
3 | |
4 | |
5 | namespace DB |
6 | { |
7 | |
8 | ZlibDeflatingWriteBuffer::ZlibDeflatingWriteBuffer( |
9 | std::unique_ptr<WriteBuffer> out_, |
10 | CompressionMethod compression_method, |
11 | int compression_level, |
12 | size_t buf_size, |
13 | char * existing_memory, |
14 | size_t alignment) |
15 | : BufferWithOwnMemory<WriteBuffer>(buf_size, existing_memory, alignment) |
16 | , out(std::move(out_)) |
17 | { |
18 | zstr.zalloc = nullptr; |
19 | zstr.zfree = nullptr; |
20 | zstr.opaque = nullptr; |
21 | zstr.next_in = nullptr; |
22 | zstr.avail_in = 0; |
23 | zstr.next_out = nullptr; |
24 | zstr.avail_out = 0; |
25 | |
26 | int window_bits = 15; |
27 | if (compression_method == CompressionMethod::Gzip) |
28 | { |
29 | window_bits += 16; |
30 | } |
31 | |
32 | #pragma GCC diagnostic push |
33 | #pragma GCC diagnostic ignored "-Wold-style-cast" |
34 | int rc = deflateInit2(&zstr, compression_level, Z_DEFLATED, window_bits, 8, Z_DEFAULT_STRATEGY); |
35 | #pragma GCC diagnostic pop |
36 | |
37 | if (rc != Z_OK) |
38 | throw Exception(std::string("deflateInit2 failed: " ) + zError(rc) + "; zlib version: " + ZLIB_VERSION, ErrorCodes::ZLIB_DEFLATE_FAILED); |
39 | } |
40 | |
41 | ZlibDeflatingWriteBuffer::~ZlibDeflatingWriteBuffer() |
42 | { |
43 | try |
44 | { |
45 | finish(); |
46 | |
47 | int rc = deflateEnd(&zstr); |
48 | if (rc != Z_OK) |
49 | throw Exception(std::string("deflateEnd failed: " ) + zError(rc), ErrorCodes::ZLIB_DEFLATE_FAILED); |
50 | } |
51 | catch (...) |
52 | { |
53 | tryLogCurrentException(__PRETTY_FUNCTION__); |
54 | } |
55 | } |
56 | |
57 | void ZlibDeflatingWriteBuffer::nextImpl() |
58 | { |
59 | if (!offset()) |
60 | return; |
61 | |
62 | zstr.next_in = reinterpret_cast<unsigned char *>(working_buffer.begin()); |
63 | zstr.avail_in = offset(); |
64 | |
65 | do |
66 | { |
67 | out->nextIfAtEnd(); |
68 | zstr.next_out = reinterpret_cast<unsigned char *>(out->position()); |
69 | zstr.avail_out = out->buffer().end() - out->position(); |
70 | |
71 | int rc = deflate(&zstr, Z_NO_FLUSH); |
72 | out->position() = out->buffer().end() - zstr.avail_out; |
73 | |
74 | if (rc != Z_OK) |
75 | throw Exception(std::string("deflate failed: " ) + zError(rc), ErrorCodes::ZLIB_DEFLATE_FAILED); |
76 | } |
77 | while (zstr.avail_in > 0 || zstr.avail_out == 0); |
78 | } |
79 | |
80 | void ZlibDeflatingWriteBuffer::finish() |
81 | { |
82 | if (finished) |
83 | return; |
84 | |
85 | next(); |
86 | |
87 | while (true) |
88 | { |
89 | out->nextIfAtEnd(); |
90 | zstr.next_out = reinterpret_cast<unsigned char *>(out->position()); |
91 | zstr.avail_out = out->buffer().end() - out->position(); |
92 | |
93 | int rc = deflate(&zstr, Z_FINISH); |
94 | out->position() = out->buffer().end() - zstr.avail_out; |
95 | |
96 | if (rc == Z_STREAM_END) |
97 | { |
98 | finished = true; |
99 | return; |
100 | } |
101 | |
102 | if (rc != Z_OK) |
103 | throw Exception(std::string("deflate finish failed: " ) + zError(rc), ErrorCodes::ZLIB_DEFLATE_FAILED); |
104 | } |
105 | } |
106 | |
107 | } |
108 | |