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 #ifndef TENSORFLOW_CORE_LIB_IO_COMPRESSED_OUTPUTBUFFER_H_ 17 #define TENSORFLOW_CORE_LIB_IO_COMPRESSED_OUTPUTBUFFER_H_ 18 19 #include <zlib.h> 20 21 #include <string> 22 23 #include "tensorflow/core/lib/core/status.h" 24 #include "tensorflow/core/lib/io/zlib_compression_options.h" 25 #include "tensorflow/core/platform/env.h" 26 #include "tensorflow/core/platform/file_system.h" 27 #include "tensorflow/core/platform/macros.h" 28 #include "tensorflow/core/platform/types.h" 29 30 namespace tensorflow { 31 namespace io { 32 33 // Provides support for writing compressed output to file using zlib 34 // (http://www.zlib.net/). 35 // A given instance of an ZlibOutputBuffer is NOT safe for concurrent use 36 // by multiple threads 37 class ZlibOutputBuffer : public WritableFile { 38 public: 39 // Create an ZlibOutputBuffer for `file` with two buffers that cache the 40 // 1. input data to be deflated 41 // 2. the deflated output 42 // with sizes `input_buffer_bytes` and `output_buffer_bytes` respectively. 43 // Does not take ownership of `file`. 44 // output_buffer_bytes should be greater than 1. 45 ZlibOutputBuffer( 46 WritableFile* file, 47 int32 input_buffer_bytes, // size of z_stream.next_in buffer 48 int32 output_buffer_bytes, // size of z_stream.next_out buffer 49 const ZlibCompressionOptions& zlib_options); 50 51 ~ZlibOutputBuffer(); 52 53 // Initializes some state necessary for the output buffer. This call is 54 // required before any other operation on the buffer. 55 Status Init(); 56 57 // Adds `data` to the compression pipeline. 58 // 59 // The input data is buffered in `z_stream_input_` and is compressed in bulk 60 // when the buffer gets full. The compressed output is not immediately 61 // written to file but rather buffered in `z_stream_output_` and gets written 62 // to file when the buffer is full. 63 // 64 // To immediately write contents to file call `Flush()`. 65 Status Append(const StringPiece& data) override; 66 67 // Deflates any cached input and writes all output to file. 68 Status Flush() override; 69 70 // Compresses any cached input and writes all output to file. This must be 71 // called before the destructor to avoid any data loss. 72 // 73 // Contrary to `Flush()` this informs zlib that it should not expect any 74 // further input by using Z_FINISH flush mode. Also cleans up z_stream. 75 // 76 // After calling this, any further calls to `Write()`, `Flush()` or `Close()` 77 // will fail. 78 Status Close() override; 79 80 // Deflates any cached input, writes all output to file and syncs it. 81 Status Sync() override; 82 83 private: 84 WritableFile* file_; // Not owned 85 Status init_status_; 86 size_t input_buffer_capacity_; 87 size_t output_buffer_capacity_; 88 89 // Buffer for storing contents read from input `file_`. 90 // TODO(srbs): Consider using circular buffers. That would greatly simplify 91 // the implementation. 92 std::unique_ptr<Bytef[]> z_stream_input_; 93 94 // Buffer for storing deflated contents of `file_`. 95 std::unique_ptr<Bytef[]> z_stream_output_; 96 97 ZlibCompressionOptions const zlib_options_; 98 99 // Configuration passed to `deflate`. 100 // 101 // z_stream_->next_in: 102 // Next byte to compress. Points to some byte in z_stream_input_ buffer. 103 // z_stream_->avail_in: 104 // Number of bytes available to be compressed at this time. 105 // z_stream_->next_out: 106 // Next byte to write compressed data to. Points to some byte in 107 // z_stream_output_ buffer. 108 // z_stream_->avail_out: 109 // Number of free bytes available at write location. 110 std::unique_ptr<z_stream> z_stream_; 111 112 // Adds `data` to `z_stream_input_`. 113 // Throws if `data.size()` > AvailableInputSpace(). 114 void AddToInputBuffer(StringPiece data); 115 116 // Returns the total space available in z_input_stream_ buffer. 117 int32 AvailableInputSpace() const; 118 119 // Deflate contents in z_stream_input_ and store results in z_stream_output_. 120 // The contents of output stream are written to file if more space is needed. 121 // On successful termination it is assured that: 122 // - z_stream_->avail_in == 0 123 // - z_stream_->avail_out > 0 124 // 125 // Note: This method does not flush contents to file. 126 // Returns non-ok status if writing contents to file fails. 127 Status DeflateBuffered(bool last = false); 128 129 // Appends contents of `z_stream_output_` to `file_`. 130 // Returns non-OK status if writing to file fails. 131 Status FlushOutputBufferToFile(); 132 133 // Calls `deflate()` and returns DataLoss Status if it failed. 134 Status Deflate(int flush); 135 136 static bool IsSyncOrFullFlush(uint8 flush_mode) { 137 return flush_mode == Z_SYNC_FLUSH || flush_mode == Z_FULL_FLUSH; 138 } 139 140 TF_DISALLOW_COPY_AND_ASSIGN(ZlibOutputBuffer); 141 }; 142 143 } // namespace io 144 } // namespace tensorflow 145 146 #endif // TENSORFLOW_CORE_LIB_IO_COMPRESSED_OUTPUTBUFFER_H_ 147
