1#include <IO/ZlibDeflatingWriteBuffer.h>
2#include <Common/MemorySanitizer.h>
3
4
5namespace DB
6{
7
8ZlibDeflatingWriteBuffer::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
41ZlibDeflatingWriteBuffer::~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
57void 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
80void 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