folly/io/Compression.h

   1 /*
   2  * Copyright 2017 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #pragma once
  18
  19 #include <cstdint>
  20 #include <limits>
  21 #include <memory>
  22 #include <string>
  23 #include <vector>
  24
  25 #include <folly/Optional.h>
  26 #include <folly/Range.h>
  27 #include <folly/io/IOBuf.h>
  28
  29 /**
  30  * Compression / decompression over IOBufs
  31  */
  32
  33 namespace folly {
  34 namespace io {
  35
  36 enum class CodecType {
  37   /**
  38    * This codec type is not defined; getCodec() will throw an exception
  39    * if used. Useful if deriving your own classes from Codec without
  40    * going through the getCodec() interface.
  41    */
  42   USER_DEFINED = 0,
  43
  44   /**
  45    * Use no compression.
  46    * Levels supported: 0
  47    */
  48   NO_COMPRESSION = 1,
  49
  50   /**
  51    * Use LZ4 compression.
  52    * Levels supported: 1 = fast, 2 = best; default = 1
  53    */
  54   LZ4 = 2,
  55
  56   /**
  57    * Use Snappy compression.
  58    * Levels supported: 1
  59    */
  60   SNAPPY = 3,
  61
  62   /**
  63    * Use zlib compression.
  64    * Levels supported: 0 = no compression, 1 = fast, ..., 9 = best; default = 6
  65    */
  66   ZLIB = 4,
  67
  68   /**
  69    * Use LZ4 compression, prefixed with size (as Varint).
  70    */
  71   LZ4_VARINT_SIZE = 5,
  72
  73   /**
  74    * Use LZMA2 compression.
  75    * Levels supported: 0 = no compression, 1 = fast, ..., 9 = best; default = 6
  76    */
  77   LZMA2 = 6,
  78   LZMA2_VARINT_SIZE = 7,
  79
  80   /**
  81    * Use ZSTD compression.
  82    */
  83   ZSTD = 8,
  84
  85   /**
  86    * Use gzip compression.  This is the same compression algorithm as ZLIB but
  87    * gzip-compressed files tend to be easier to work with from the command line.
  88    * Levels supported: 0 = no compression, 1 = fast, ..., 9 = best; default = 6
  89    */
  90   GZIP = 9,
  91
  92   /**
  93    * Use LZ4 frame compression.
  94    * Levels supported: 0 = fast, 16 = best; default = 0
  95    */
  96   LZ4_FRAME = 10,
  97
  98   /**
  99    * Use bzip2 compression.
 100    * Levels supported: 1 = fast, 9 = best; default = 9
 101    */
 102   BZIP2 = 11,
 103
 104   NUM_CODEC_TYPES = 12,
 105 };
 106
 107 class Codec {
 108  public:
 109   virtual ~Codec() { }
 110
 111   static constexpr uint64_t UNLIMITED_UNCOMPRESSED_LENGTH = uint64_t(-1);
 112   /**
 113    * Return the maximum length of data that may be compressed with this codec.
 114    * NO_COMPRESSION and ZLIB support arbitrary lengths;
 115    * LZ4 supports up to 1.9GiB; SNAPPY supports up to 4GiB.
 116    * May return UNLIMITED_UNCOMPRESSED_LENGTH if unlimited.
 117    */
 118   uint64_t maxUncompressedLength() const;
 119
 120   /**
 121    * Return the codec's type.
 122    */
 123   CodecType type() const { return type_; }
 124
 125   /**
 126    * Does this codec need the exact uncompressed length on decompression?
 127    */
 128   bool needsUncompressedLength() const;
 129
 130   /**
 131    * Compress data, returning an IOBuf (which may share storage with data).
 132    * Throws std::invalid_argument if data is larger than
 133    * maxUncompressedLength().
 134    *
 135    * Regardless of the behavior of the underlying compressor, compressing
 136    * an empty IOBuf chain will return an empty IOBuf chain.
 137    */
 138   std::unique_ptr<IOBuf> compress(const folly::IOBuf* data);
 139
 140   /**
 141    * Compresses data. May involve additional copies compared to the overload
 142    * that takes and returns IOBufs. Has the same error semantics as the IOBuf
 143    * version.
 144    */
 145   std::string compress(StringPiece data);
 146
 147   /**
 148    * Uncompress data. Throws std::runtime_error on decompression error.
 149    *
 150    * Some codecs (LZ4) require the exact uncompressed length; this is indicated
 151    * by needsUncompressedLength().
 152    *
 153    * For other codes (zlib), knowing the exact uncompressed length ahead of
 154    * time might be faster.
 155    *
 156    * Regardless of the behavior of the underlying compressor, uncompressing
 157    * an empty IOBuf chain will return an empty IOBuf chain.
 158    */
 159   std::unique_ptr<IOBuf> uncompress(
 160       const IOBuf* data,
 161       folly::Optional<uint64_t> uncompressedLength = folly::none);
 162
 163   /**
 164    * Uncompresses data. May involve additional copies compared to the overload
 165    * that takes and returns IOBufs. Has the same error semantics as the IOBuf
 166    * version.
 167    */
 168   std::string uncompress(
 169       StringPiece data,
 170       folly::Optional<uint64_t> uncompressedLength = folly::none);
 171
 172   /**
 173    * Returns a bound on the maximum compressed length when compressing data with
 174    * the given uncompressed length.
 175    */
 176   uint64_t maxCompressedLength(uint64_t uncompressedLength) const;
 177
 178   /**
 179    * Extracts the uncompressed length from the compressed data if possible.
 180    * If the codec doesn't store the uncompressed length, or the data is
 181    * corrupted it returns the given uncompressedLength.
 182    * If the uncompressed length is stored in the compressed data and
 183    * uncompressedLength is not none and they do not match a std::runtime_error
 184    * is thrown.
 185    */
 186   folly::Optional<uint64_t> getUncompressedLength(
 187       const folly::IOBuf* data,
 188       folly::Optional<uint64_t> uncompressedLength = folly::none) const;
 189
 190  protected:
 191   explicit Codec(CodecType type);
 192
 193  public:
 194   /**
 195    * Returns a superset of the set of prefixes for which canUncompress() will
 196    * return true. A superset is allowed for optimizations in canUncompress()
 197    * based on other knowledge such as length. None of the prefixes may be empty.
 198    * default: No prefixes.
 199    */
 200   virtual std::vector<std::string> validPrefixes() const;
 201
 202   /**
 203    * Returns true if the codec thinks it can uncompress the data.
 204    * If a codec doesn't have magic bytes at the beginning, like LZ4 and Snappy,
 205    * it can always return false.
 206    * default: Returns false.
 207    */
 208   virtual bool canUncompress(
 209       const folly::IOBuf* data,
 210       folly::Optional<uint64_t> uncompressedLength = folly::none) const;
 211
 212  private:
 213   // default: no limits (save for special value UNKNOWN_UNCOMPRESSED_LENGTH)
 214   virtual uint64_t doMaxUncompressedLength() const;
 215   // default: doesn't need uncompressed length
 216   virtual bool doNeedsUncompressedLength() const;
 217   virtual std::unique_ptr<IOBuf> doCompress(const folly::IOBuf* data) = 0;
 218   virtual std::unique_ptr<IOBuf> doUncompress(
 219       const folly::IOBuf* data,
 220       folly::Optional<uint64_t> uncompressedLength) = 0;
 221   // default: an implementation is provided by default to wrap the strings into
 222   // IOBufs and delegate to the IOBuf methods. This incurs a copy of the output
 223   // from IOBuf to string. Implementers, at their discretion, can override
 224   // these methods to avoid the copy.
 225   virtual std::string doCompressString(StringPiece data);
 226   virtual std::string doUncompressString(
 227       StringPiece data,
 228       folly::Optional<uint64_t> uncompressedLength);
 229
 230   virtual uint64_t doMaxCompressedLength(uint64_t uncompressedLength) const = 0;
 231   // default: returns the passed uncompressedLength.
 232   virtual folly::Optional<uint64_t> doGetUncompressedLength(
 233       const folly::IOBuf* data,
 234       folly::Optional<uint64_t> uncompressedLength) const;
 235
 236   CodecType type_;
 237 };
 238
 239 class StreamCodec : public Codec {
 240  public:
 241   ~StreamCodec() override {}
 242
 243   /**
 244    * Does the codec need the data length before compression streaming?
 245    */
 246   bool needsDataLength() const;
 247
 248   /*****************************************************************************
 249    * Streaming API
 250    *****************************************************************************
 251    * A low-level stateful streaming API.
 252    * Streaming operations can be started in two ways:
 253    *   1. From a clean Codec on which no non-const methods have been called.
 254    *   2. A call to resetStream(), which will reset any codec to a clean state.
 255    * After a streaming operation has begun, either compressStream() or
 256    * uncompressStream() must be called until the streaming operation ends.
 257    * compressStream() ends when it returns true with flushOp END.
 258    * uncompressStream() ends when it returns true. At this point the codec
 259    * may be reused by calling resetStream().
 260    *
 261    * compress() and uncompress() can be called at any time, but they interrupt
 262    * any ongoing streaming operations (state is lost and resetStream() must be
 263    * called before another streaming operation).
 264    */
 265
 266   /**
 267    * Reset the state of the codec, and set the uncompressed length for the next
 268    * streaming operation. If uncompressedLength is not none it must be exactly
 269    * the uncompressed length. compressStream() must be passed exactly
 270    * uncompressedLength input bytes before the stream is ended.
 271    * uncompressStream() must be passed a compressed frame that uncompresses to
 272    * uncompressedLength.
 273    */
 274   void resetStream(folly::Optional<uint64_t> uncompressedLength = folly::none);
 275
 276   enum class FlushOp { NONE, FLUSH, END };
 277
 278   /**
 279    * Compresses some data from the input buffer and writes the compressed data
 280    * into the output buffer. It may read input without producing any output,
 281    * except when forced to flush.
 282    *
 283    * The input buffer is advanced to point to the range of data that hasn't yet
 284    * been read. Compression will resume at this point for the next call to
 285    * compressStream(). The output buffer is advanced one byte past the last byte
 286    * written.
 287    *
 288    * The default flushOp is NONE, which allows compressStream() complete
 289    * discretion in how much data to gather before writing any output.
 290    *
 291    * If flushOp is END, all pending and input data is flushed to the output
 292    * buffer, and the frame is ended. compressStream() must be called with the
 293    * same input and flushOp END until it returns true. At this point the caller
 294    * must call resetStream() to use the codec again.
 295    *
 296    * If flushOp is FLUSH, all pending and input data is flushed to the output
 297    * buffer, but the frame is not ended. compressStream() must be called with
 298    * the same input and flushOp END until it returns true. At this point the
 299    * caller can continue to compressStream() with any input data and flushOp.
 300    * The uncompressor, if passed all the produced output data, will be able to
 301    * uncompress all the input data passed to compressStream() so far. Excessive
 302    * use of flushOp FLUSH will deteriorate compression ratio. This is useful for
 303    * stateful streaming across a network. Most users don't need to use this
 304    * flushOp.
 305    *
 306    * A std::logic_error is thrown on incorrect usage of the API.
 307    * A std::runtime_error is thrown upon error conditions.
 308    */
 309   bool compressStream(
 310       folly::ByteRange& input,
 311       folly::MutableByteRange& output,
 312       FlushOp flushOp = StreamCodec::FlushOp::NONE);
 313
 314   /**
 315    * Uncompresses some data from the input buffer and writes the uncompressed
 316    * data into the output buffer. It may read input without producing any
 317    * output.
 318    *
 319    * The input buffer is advanced to point to the range of data that hasn't yet
 320    * been read. Uncompression will resume at this point for the next call to
 321    * uncompressStream(). The output buffer is advanced one byte past the last
 322    * byte written.
 323    *
 324    * The default flushOp is NONE, which allows uncompressStream() complete
 325    * discretion in how much output data to flush. The uncompressor may not make
 326    * maximum forward progress, but will make some forward progress when
 327    * possible.
 328    *
 329    * If flushOp is END, the caller guarantees that no more input will be
 330    * presented to uncompressStream(). uncompressStream() must be called with the
 331    * same input and flushOp END until it returns true. This is not mandatory,
 332    * but if the input is all available in one buffer, and there is enough output
 333    * space to write the entire frame, codecs can uncompress faster.
 334    *
 335    * If flushOp is FLUSH, uncompressStream() is guaranteed to make the maximum
 336    * amount of forward progress possible. When using this flushOp and
 337    * uncompressStream() returns with `!output.empty()` the caller knows that all
 338    * pending output has been flushed. This is useful for stateful streaming
 339    * across a network, and it should be used in conjunction with
 340    * compressStream() with flushOp FLUSH. Most users don't need to use this
 341    * flushOp.
 342    *
 343    * Returns true at the end of a frame. At this point resetStream() must be
 344    * called to reuse the codec.
 345    */
 346   bool uncompressStream(
 347       folly::ByteRange& input,
 348       folly::MutableByteRange& output,
 349       FlushOp flushOp = StreamCodec::FlushOp::NONE);
 350
 351  protected:
 352   explicit StreamCodec(CodecType type) : Codec(type) {}
 353
 354   // Returns the uncompressed length last passed to resetStream() or none if it
 355   // hasn't been called yet.
 356   folly::Optional<uint64_t> uncompressedLength() const {
 357     return uncompressedLength_;
 358   }
 359
 360  private:
 361   // default: Implemented using the streaming API.
 362   std::unique_ptr<IOBuf> doCompress(const folly::IOBuf* data) override;
 363   std::unique_ptr<IOBuf> doUncompress(
 364       const folly::IOBuf* data,
 365       folly::Optional<uint64_t> uncompressedLength) override;
 366
 367   // default: Returns false
 368   virtual bool doNeedsDataLength() const;
 369   virtual void doResetStream() = 0;
 370   virtual bool doCompressStream(
 371       folly::ByteRange& input,
 372       folly::MutableByteRange& output,
 373       FlushOp flushOp) = 0;
 374   virtual bool doUncompressStream(
 375       folly::ByteRange& input,
 376       folly::MutableByteRange& output,
 377       FlushOp flushOp) = 0;
 378
 379   enum class State {
 380     RESET,
 381     COMPRESS,
 382     COMPRESS_FLUSH,
 383     COMPRESS_END,
 384     UNCOMPRESS,
 385     END,
 386   };
 387   void assertStateIs(State expected) const;
 388
 389   CodecType type_;
 390   State state_{State::RESET};
 391   ByteRange previousInput_{};
 392   folly::Optional<uint64_t> uncompressedLength_{};
 393 };
 394
 395 constexpr int COMPRESSION_LEVEL_FASTEST = -1;
 396 constexpr int COMPRESSION_LEVEL_DEFAULT = -2;
 397 constexpr int COMPRESSION_LEVEL_BEST = -3;
 398
 399 /**
 400  * Return a codec for the given type. Throws on error.  The level
 401  * is a non-negative codec-dependent integer indicating the level of
 402  * compression desired, or one of the following constants:
 403  *
 404  * COMPRESSION_LEVEL_FASTEST is fastest (uses least CPU / memory,
 405  *   worst compression)
 406  * COMPRESSION_LEVEL_DEFAULT is the default (likely a tradeoff between
 407  *   FASTEST and BEST)
 408  * COMPRESSION_LEVEL_BEST is the best compression (uses most CPU / memory,
 409  *   best compression)
 410  *
 411  * When decompressing, the compression level is ignored. All codecs will
 412  * decompress all data compressed with the a codec of the same type, regardless
 413  * of compression level.
 414  */
 415 std::unique_ptr<Codec> getCodec(
 416     CodecType type,
 417     int level = COMPRESSION_LEVEL_DEFAULT);
 418
 419 /**
 420  * Return a codec for the given type. Throws on error.  The level
 421  * is a non-negative codec-dependent integer indicating the level of
 422  * compression desired, or one of the following constants:
 423  *
 424  * COMPRESSION_LEVEL_FASTEST is fastest (uses least CPU / memory,
 425  *   worst compression)
 426  * COMPRESSION_LEVEL_DEFAULT is the default (likely a tradeoff between
 427  *   FASTEST and BEST)
 428  * COMPRESSION_LEVEL_BEST is the best compression (uses most CPU / memory,
 429  *   best compression)
 430  *
 431  * When decompressing, the compression level is ignored. All codecs will
 432  * decompress all data compressed with the a codec of the same type, regardless
 433  * of compression level.
 434  */
 435 std::unique_ptr<StreamCodec> getStreamCodec(
 436     CodecType type,
 437     int level = COMPRESSION_LEVEL_DEFAULT);
 438
 439 /**
 440  * Returns a codec that can uncompress any of the given codec types as well as
 441  * {LZ4_FRAME, ZSTD, ZLIB, GZIP, LZMA2, BZIP2}. Appends each default codec to
 442  * customCodecs in order, so long as a codec with the same type() isn't already
 443  * present. When uncompress() is called, each codec's canUncompress() is called
 444  * in the order that they are given. Appended default codecs are checked last.
 445  * uncompress() is called on the first codec whose canUncompress() returns true.
 446  * An exception is thrown if no codec canUncompress() the data.
 447  * An exception is thrown if the chosen codec's uncompress() throws on the data.
 448  * An exception is thrown if compress() is called on the returned codec.
 449  *
 450  * Requirements are checked in debug mode and are as follows:
 451  * Let headers be the concatenation of every codec's validPrefixes().
 452  *  1. Each codec must override validPrefixes() and canUncompress().
 453  *  2. No codec's validPrefixes() may be empty.
 454  *  3. No header in headers may be empty.
 455  *  4. headers must not contain any duplicate elements.
 456  *  5. No strict non-empty prefix of any header in headers may be in headers.
 457  */
 458 std::unique_ptr<Codec> getAutoUncompressionCodec(
 459     std::vector<std::unique_ptr<Codec>> customCodecs = {});
 460
 461 /**
 462  * Check if a specified codec is supported.
 463  */
 464 bool hasCodec(CodecType type);
 465
 466 /**
 467  * Check if a specified codec is supported and supports streaming.
 468  */
 469 bool hasStreamCodec(CodecType type);
 470 } // namespace io
 471 } // namespace folly