• 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 #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 #if defined(TF_CORD_SUPPORT)
69   Status Append(const absl::Cord& cord) override;
70 #endif
71 
72   // Deflates any cached input and writes all output to file.
73   Status Flush() override;
74 
75   // Compresses any cached input and writes all output to file. This must be
76   // called before the destructor to avoid any data loss.
77   //
78   // Contrary to `Flush()` this informs zlib that it should not expect any
79   // further input by using Z_FINISH flush mode. Also cleans up z_stream.
80   //
81   // After calling this, any further calls to `Write()`, `Flush()` or `Close()`
82   // will fail.
83   Status Close() override;
84 
85   // Returns the name of the underlying file.
86   Status Name(StringPiece* result) const override;
87 
88   // Deflates any cached input, writes all output to file and syncs it.
89   Status Sync() override;
90 
91   // Returns the write position in the underlying file. The position does not
92   // reflect buffered, un-flushed data.
93   Status Tell(int64* position) override;
94 
95  private:
96   WritableFile* file_;  // Not owned
97   Status init_status_;
98   size_t input_buffer_capacity_;
99   size_t output_buffer_capacity_;
100 
101   // Buffer for storing contents read from input `file_`.
102   // TODO(srbs): Consider using circular buffers. That would greatly simplify
103   // the implementation.
104   std::unique_ptr<Bytef[]> z_stream_input_;
105 
106   // Buffer for storing deflated contents of `file_`.
107   std::unique_ptr<Bytef[]> z_stream_output_;
108 
109   ZlibCompressionOptions const zlib_options_;
110 
111   // Configuration passed to `deflate`.
112   //
113   // z_stream_->next_in:
114   //   Next byte to compress. Points to some byte in z_stream_input_ buffer.
115   // z_stream_->avail_in:
116   //   Number of bytes available to be compressed at this time.
117   // z_stream_->next_out:
118   //   Next byte to write compressed data to. Points to some byte in
119   //   z_stream_output_ buffer.
120   // z_stream_->avail_out:
121   //   Number of free bytes available at write location.
122   std::unique_ptr<z_stream> z_stream_;
123 
124   // Adds `data` to `z_stream_input_`.
125   // Throws if `data.size()` > AvailableInputSpace().
126   void AddToInputBuffer(StringPiece data);
127 
128   // Returns the total space available in z_input_stream_ buffer.
129   int32 AvailableInputSpace() const;
130 
131   // Deflate contents in z_stream_input_ and store results in z_stream_output_.
132   // The contents of output stream are written to file if more space is needed.
133   // On successful termination it is assured that:
134   // - z_stream_->avail_in == 0
135   // - z_stream_->avail_out > 0
136   //
137   // Note: This method does not flush contents to file.
138   // Returns non-ok status if writing contents to file fails.
139   Status DeflateBuffered(int flush_mode);
140 
141   // Appends contents of `z_stream_output_` to `file_`.
142   // Returns non-OK status if writing to file fails.
143   Status FlushOutputBufferToFile();
144 
145   // Calls `deflate()` and returns DataLoss Status if it failed.
146   Status Deflate(int flush);
147 
IsSyncOrFullFlush(uint8 flush_mode)148   static bool IsSyncOrFullFlush(uint8 flush_mode) {
149     return flush_mode == Z_SYNC_FLUSH || flush_mode == Z_FULL_FLUSH;
150   }
151 
152   TF_DISALLOW_COPY_AND_ASSIGN(ZlibOutputBuffer);
153 };
154 
155 }  // namespace io
156 }  // namespace tensorflow
157 
158 #endif  // TENSORFLOW_CORE_LIB_IO_COMPRESSED_OUTPUTBUFFER_H_
159