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