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