• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #include "tensorflow/core/lib/io/zlib_outputbuffer.h"
17 
18 #include "tensorflow/core/lib/core/errors.h"
19 
20 namespace tensorflow {
21 namespace io {
22 
ZlibOutputBuffer(WritableFile * file,int32 input_buffer_bytes,int32 output_buffer_bytes,const ZlibCompressionOptions & zlib_options)23 ZlibOutputBuffer::ZlibOutputBuffer(
24     WritableFile* file,
25     int32 input_buffer_bytes,  // size of z_stream.next_in buffer
26     int32 output_buffer_bytes,
27     const ZlibCompressionOptions&
28         zlib_options)  // size of z_stream.next_out buffer
29     : file_(file),
30       init_status_(),
31       input_buffer_capacity_(input_buffer_bytes),
32       output_buffer_capacity_(output_buffer_bytes),
33       z_stream_input_(new Bytef[input_buffer_bytes]),
34       z_stream_output_(new Bytef[output_buffer_bytes]),
35       zlib_options_(zlib_options),
36       z_stream_(new z_stream) {}
37 
~ZlibOutputBuffer()38 ZlibOutputBuffer::~ZlibOutputBuffer() {
39   if (z_stream_) {
40     LOG(WARNING) << "ZlibOutputBuffer::Close() not called. Possible data loss";
41   }
42 }
43 
Init()44 Status ZlibOutputBuffer::Init() {
45   // Output buffer size should be greater than 1 because deflation needs at
46   // least one byte for book keeping etc.
47   if (output_buffer_capacity_ <= 1) {
48     return errors::InvalidArgument(
49         "output_buffer_bytes should be greater than "
50         "1");
51   }
52   memset(z_stream_.get(), 0, sizeof(z_stream));
53   z_stream_->zalloc = Z_NULL;
54   z_stream_->zfree = Z_NULL;
55   z_stream_->opaque = Z_NULL;
56   int status =
57       deflateInit2(z_stream_.get(), zlib_options_.compression_level,
58                    zlib_options_.compression_method, zlib_options_.window_bits,
59                    zlib_options_.mem_level, zlib_options_.compression_strategy);
60   if (status != Z_OK) {
61     z_stream_.reset(nullptr);
62     return errors::InvalidArgument("deflateInit failed with status", status);
63   }
64   z_stream_->next_in = z_stream_input_.get();
65   z_stream_->next_out = z_stream_output_.get();
66   z_stream_->avail_in = 0;
67   z_stream_->avail_out = output_buffer_capacity_;
68   return Status::OK();
69 }
70 
AvailableInputSpace() const71 int32 ZlibOutputBuffer::AvailableInputSpace() const {
72   return input_buffer_capacity_ - z_stream_->avail_in;
73 }
74 
AddToInputBuffer(StringPiece data)75 void ZlibOutputBuffer::AddToInputBuffer(StringPiece data) {
76   size_t bytes_to_write = data.size();
77   CHECK_LE(bytes_to_write, AvailableInputSpace());
78 
79   // Input stream ->
80   // [....................input_buffer_capacity_...............]
81   // [<...read_bytes...><...avail_in...>......empty space......]
82   //  ^                 ^
83   //  |                 |
84   //  z_stream_input_   next_in
85   //
86   // Data in the input stream is sharded as show above. z_stream_->next_in could
87   // be pointing to some byte in the buffer with avail_in number of bytes
88   // available to be read.
89   //
90   // In order to avoid shifting the avail_in bytes at next_in to the head of
91   // the buffer we try to fit `data` in the empty space at the tail of the
92   // input stream.
93   // TODO(srbs): This could be avoided if we had a circular buffer.
94   // If it doesn't fit we free the space at the head of the stream and then
95   // append `data` at the end of existing data.
96 
97   int32 read_bytes = z_stream_->next_in - z_stream_input_.get();
98   int32 unread_bytes = z_stream_->avail_in;
99   int32 free_tail_bytes = input_buffer_capacity_ - (read_bytes + unread_bytes);
100 
101   if (static_cast<int32>(bytes_to_write) > free_tail_bytes) {
102     memmove(z_stream_input_.get(), z_stream_->next_in, z_stream_->avail_in);
103     z_stream_->next_in = z_stream_input_.get();
104   }
105   memcpy(z_stream_->next_in + z_stream_->avail_in, data.data(), bytes_to_write);
106   z_stream_->avail_in += bytes_to_write;
107 }
108 
DeflateBuffered(int flush_mode)109 Status ZlibOutputBuffer::DeflateBuffered(int flush_mode) {
110   do {
111     // From zlib manual (http://www.zlib.net/manual.html):
112     //
113     // "In the case of a Z_FULL_FLUSH or Z_SYNC_FLUSH, make sure that
114     // avail_out is greater than six to avoid repeated flush markers due
115     // to avail_out == 0 on return."
116     //
117     // If above condition is met or if output buffer is full we flush contents
118     // to file.
119     if (z_stream_->avail_out == 0 ||
120         (IsSyncOrFullFlush(flush_mode) && z_stream_->avail_out < 6)) {
121       TF_RETURN_IF_ERROR(FlushOutputBufferToFile());
122     }
123     TF_RETURN_IF_ERROR(Deflate(flush_mode));
124   } while (z_stream_->avail_out == 0);
125 
126   DCHECK(z_stream_->avail_in == 0);
127   z_stream_->next_in = z_stream_input_.get();
128   return Status::OK();
129 }
130 
FlushOutputBufferToFile()131 Status ZlibOutputBuffer::FlushOutputBufferToFile() {
132   uint32 bytes_to_write = output_buffer_capacity_ - z_stream_->avail_out;
133   if (bytes_to_write > 0) {
134     Status s = file_->Append(StringPiece(
135         reinterpret_cast<char*>(z_stream_output_.get()), bytes_to_write));
136     if (s.ok()) {
137       z_stream_->next_out = z_stream_output_.get();
138       z_stream_->avail_out = output_buffer_capacity_;
139     }
140     return s;
141   }
142   return Status::OK();
143 }
144 
Append(StringPiece data)145 Status ZlibOutputBuffer::Append(StringPiece data) {
146   // If there is sufficient free space in z_stream_input_ to fit data we
147   // add it there and return.
148   // If there isn't enough space we deflate the existing contents of
149   // z_input_stream_. If data now fits in z_input_stream_ we add it there
150   // else we directly deflate it.
151   //
152   // The deflated output is accumulated in z_stream_output_ and gets written to
153   // file as and when needed.
154 
155   size_t bytes_to_write = data.size();
156 
157   if (static_cast<int32>(bytes_to_write) <= AvailableInputSpace()) {
158     AddToInputBuffer(data);
159     return Status::OK();
160   }
161 
162   TF_RETURN_IF_ERROR(DeflateBuffered(zlib_options_.flush_mode));
163 
164   // At this point input stream should be empty.
165   if (static_cast<int32>(bytes_to_write) <= AvailableInputSpace()) {
166     AddToInputBuffer(data);
167     return Status::OK();
168   }
169 
170   // `data` is too large to fit in input buffer so we deflate it directly.
171   // Note that at this point we have already deflated all existing input so
172   // we do not need to backup next_in and avail_in.
173   z_stream_->next_in = reinterpret_cast<Bytef*>(const_cast<char*>(data.data()));
174   z_stream_->avail_in = bytes_to_write;
175 
176   do {
177     if (z_stream_->avail_out == 0) {
178       // No available output space.
179       // Write output buffer to file.
180       TF_RETURN_IF_ERROR(FlushOutputBufferToFile());
181     }
182     TF_RETURN_IF_ERROR(Deflate(zlib_options_.flush_mode));
183   } while (z_stream_->avail_out == 0);
184 
185   DCHECK(z_stream_->avail_in == 0);  // All input will be used up.
186 
187   // Restore z_stream input pointers.
188   z_stream_->next_in = z_stream_input_.get();
189 
190   return Status::OK();
191 }
192 
193 #if defined(TF_CORD_SUPPORT)
Append(const absl::Cord & cord)194 Status ZlibOutputBuffer::Append(const absl::Cord& cord) {
195   for (absl::string_view fragment : cord.Chunks()) {
196     TF_RETURN_IF_ERROR(Append(fragment));
197   }
198   return Status::OK();
199 }
200 #endif
201 
Flush()202 Status ZlibOutputBuffer::Flush() {
203   TF_RETURN_IF_ERROR(DeflateBuffered(Z_PARTIAL_FLUSH));
204   TF_RETURN_IF_ERROR(FlushOutputBufferToFile());
205   return file_->Flush();
206 }
207 
Name(StringPiece * result) const208 Status ZlibOutputBuffer::Name(StringPiece* result) const {
209   return file_->Name(result);
210 }
211 
Sync()212 Status ZlibOutputBuffer::Sync() {
213   TF_RETURN_IF_ERROR(Flush());
214   return file_->Sync();
215 }
216 
Close()217 Status ZlibOutputBuffer::Close() {
218   if (z_stream_) {
219     TF_RETURN_IF_ERROR(DeflateBuffered(Z_FINISH));
220     TF_RETURN_IF_ERROR(FlushOutputBufferToFile());
221     deflateEnd(z_stream_.get());
222     z_stream_.reset(nullptr);
223   }
224   return Status::OK();
225 }
226 
Deflate(int flush)227 Status ZlibOutputBuffer::Deflate(int flush) {
228   int error = deflate(z_stream_.get(), flush);
229   if (error == Z_OK || error == Z_BUF_ERROR ||
230       (error == Z_STREAM_END && flush == Z_FINISH)) {
231     return Status::OK();
232   }
233   string error_string = strings::StrCat("deflate() failed with error ", error);
234   if (z_stream_->msg != nullptr) {
235     strings::StrAppend(&error_string, ": ", z_stream_->msg);
236   }
237   return errors::DataLoss(error_string);
238 }
239 
Tell(int64 * position)240 Status ZlibOutputBuffer::Tell(int64* position) { return file_->Tell(position); }
241 
242 }  // namespace io
243 }  // namespace tensorflow
244