Coverage Report

Created: 2026-08-14 08:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/src/rocksdb/util/compression.h
Line
Count
Source
1
// Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2
//  This source code is licensed under both the GPLv2 (found in the
3
//  COPYING file in the root directory) and Apache 2.0 License
4
//  (found in the LICENSE.Apache file in the root directory).
5
//
6
// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
7
// Use of this source code is governed by a BSD-style license that can be
8
// found in the LICENSE file. See the AUTHORS file for names of contributors.
9
//
10
#pragma once
11
12
#include <algorithm>
13
14
#include "memory/memory_allocator_impl.h"
15
#include "rocksdb/advanced_compression.h"
16
#include "rocksdb/options.h"
17
#include "table/block_based/block_type.h"
18
#include "util/aligned_buffer.h"
19
#include "util/coding.h"
20
#include "util/compression_context_cache.h"
21
22
#ifdef ZSTD
23
#include <zstd.h>
24
#include <zstd_errors.h>
25
// ZSTD_Compress2(), ZSTD_compressStream2() and frame parameters all belong to
26
// advanced APIs and require v1.4.0+, which is from April 2019.
27
// https://github.com/facebook/zstd/blob/eb9f881eb810f2242f1ef36b3f3e7014eecb8fa6/lib/zstd.h#L297C40-L297C45
28
// To avoid a rat's nest of #ifdefs, we now require v1.4.0+ for ZSTD support.
29
#if ZSTD_VERSION_NUMBER < 10400
30
#error "ZSTD support requires version >= 1.4.0 (libzstd-devel)"
31
#endif  // ZSTD_VERSION_NUMBER
32
// The above release also includes digested dictionary support, but some
33
// required functions (ZSTD_createDDict_byReference) are still only exported
34
// with ZSTD_STATIC_LINKING_ONLY defined.
35
#if defined(ZSTD_STATIC_LINKING_ONLY)
36
#define ROCKSDB_ZSTD_DDICT
37
#endif  // defined(ZSTD_STATIC_LINKING_ONLY)
38
//  For ZDICT_* functions
39
#include <zdict.h>
40
// ZDICT_finalizeDictionary API is exported and stable since v1.4.5
41
#if ZSTD_VERSION_NUMBER >= 10405
42
#define ROCKSDB_ZDICT_FINALIZE
43
#endif  // ZSTD_VERSION_NUMBER >= 10405
44
#endif  // ZSTD
45
46
namespace ROCKSDB_NAMESPACE {
47
// Need this for the context allocation override
48
// On windows we need to do this explicitly
49
#if defined(ZSTD) && defined(ROCKSDB_JEMALLOC) && defined(OS_WIN) && \
50
    defined(ZSTD_STATIC_LINKING_ONLY)
51
#define ROCKSDB_ZSTD_CUSTOM_MEM
52
namespace port {
53
ZSTD_customMem GetJeZstdAllocationOverrides();
54
}  // namespace port
55
#endif  // defined(ZSTD) && defined(ROCKSDB_JEMALLOC) && defined(OS_WIN) &&
56
        // defined(ZSTD_STATIC_LINKING_ONLY)
57
58
#ifdef ZSTD
59
// Translates a configured compression_opts.level into the effective ZSTD
60
// compression level. Besides resolving the "use default" sentinel, this
61
// removes a discontinuity at level 0: ZSTD itself treats a requested level of
62
// 0 as "use the default level" (historically 3), which would make level 0
63
// more aggressive than levels 1 and 2. Mapping 0 to -1 keeps the level spectrum
64
// monotonic, which is friendlier to auto-tuning.
65
inline int SanitizeZSTDCompressionLevel(int level) {
66
  if (level == CompressionOptions::kDefaultCompressionLevel) {
67
    // NB: ZSTD_CLEVEL_DEFAULT is historically == 3
68
    return ZSTD_CLEVEL_DEFAULT;
69
  } else if (level == 0) {
70
    // Avoid library's discontinuity at level 0
71
    return -1;
72
  } else {
73
    return level;
74
  }
75
}
76
#endif  // ZSTD
77
78
// Cached data represents a portion that can be re-used
79
// If, in the future we have more than one native context to
80
// cache we can arrange this as a tuple
81
class ZSTDUncompressCachedData {
82
 public:
83
#if defined(ZSTD)
84
  using ZSTDNativeContext = ZSTD_DCtx*;
85
#else
86
  using ZSTDNativeContext = void*;
87
#endif  // ZSTD
88
64
  ZSTDUncompressCachedData() {}
89
  // Init from cache
90
  ZSTDUncompressCachedData(const ZSTDUncompressCachedData& o) = delete;
91
  ZSTDUncompressCachedData& operator=(const ZSTDUncompressCachedData&) = delete;
92
  ZSTDUncompressCachedData(ZSTDUncompressCachedData&& o) noexcept
93
0
      : ZSTDUncompressCachedData() {
94
0
    *this = std::move(o);
95
0
  }
96
0
  ZSTDUncompressCachedData& operator=(ZSTDUncompressCachedData&& o) noexcept {
97
0
    assert(zstd_ctx_ == nullptr);
98
0
    std::swap(zstd_ctx_, o.zstd_ctx_);
99
0
    std::swap(cache_idx_, o.cache_idx_);
100
0
    return *this;
101
0
  }
102
0
  ZSTDNativeContext Get() const { return zstd_ctx_; }
103
0
  int64_t GetCacheIndex() const { return cache_idx_; }
104
0
  void CreateIfNeeded() {
105
0
    if (zstd_ctx_ == nullptr) {
106
0
#if !defined(ZSTD)
107
0
      zstd_ctx_ = nullptr;
108
#elif defined(ROCKSDB_ZSTD_CUSTOM_MEM)
109
      zstd_ctx_ =
110
          ZSTD_createDCtx_advanced(port::GetJeZstdAllocationOverrides());
111
#else  // ZSTD && !ROCKSDB_ZSTD_CUSTOM_MEM
112
      zstd_ctx_ = ZSTD_createDCtx();
113
#endif
114
0
      cache_idx_ = -1;
115
0
    }
116
0
  }
117
0
  void InitFromCache(const ZSTDUncompressCachedData& o, int64_t idx) {
118
0
    zstd_ctx_ = o.zstd_ctx_;
119
0
    cache_idx_ = idx;
120
0
  }
121
64
  ~ZSTDUncompressCachedData() {
122
#if defined(ZSTD)
123
    if (zstd_ctx_ != nullptr && cache_idx_ == -1) {
124
      ZSTD_freeDCtx(zstd_ctx_);
125
    }
126
#endif  // ZSTD
127
64
  }
128
129
 private:
130
  ZSTDNativeContext zstd_ctx_ = nullptr;
131
  int64_t cache_idx_ = -1;  // -1 means this instance owns the context
132
};
133
}  // namespace ROCKSDB_NAMESPACE
134
135
#if defined(XPRESS)
136
#include "port/xpress.h"
137
#endif
138
139
namespace ROCKSDB_NAMESPACE {
140
141
class FailureDecompressor : public Decompressor {
142
 public:
143
0
  explicit FailureDecompressor(Status&& status) : status_(std::move(status)) {
144
0
    assert(!status_.ok());
145
0
  }
146
0
  ~FailureDecompressor() override { status_.PermitUncheckedError(); }
147
148
0
  const char* Name() const override { return "FailureDecompressor"; }
149
150
0
  Status ExtractUncompressedSize(Args& /*args*/) override { return status_; }
151
152
  Status DecompressBlock(const Args& /*args*/,
153
0
                         char* /*uncompressed_output*/) override {
154
0
    return status_;
155
0
  }
156
157
 protected:
158
  Status status_;
159
};
160
161
// Owns a decompression dictionary, and associated Decompressor, for storing
162
// in the block cache.
163
//
164
// Justification: for a "processed" dictionary to be saved in block cache, we
165
// also need a reference to the decompressor that processed it, to ensure it
166
// is recognized properly. At that point, we might as well have the dictionary
167
// part of the decompressor identity and track an associated decompressor along
168
// with a decompression dictionary in the block cache, and the decompressor
169
// hides potential details of processing the dictionary.
170
struct DecompressorDict {
171
  // Block containing the data for the compression dictionary in case the
172
  // constructor that takes a string parameter is used.
173
  std::string dict_str_;
174
175
  // Block containing the data for the compression dictionary in case the
176
  // constructor that takes a Slice parameter is used and the passed in
177
  // CacheAllocationPtr is not nullptr.
178
  CacheAllocationPtr dict_allocation_;
179
180
  // A Decompressor referencing and using the dictionary owned by this.
181
  std::unique_ptr<Decompressor> decompressor_;
182
183
  // Approximate owned memory usage
184
  size_t memory_usage_;
185
186
  DecompressorDict(std::string&& dict, Decompressor& from_decompressor)
187
0
      : dict_str_(std::move(dict)) {
188
0
    Populate(from_decompressor, dict_str_);
189
0
  }
190
191
  DecompressorDict(Slice slice, CacheAllocationPtr&& allocation,
192
                   Decompressor& from_decompressor)
193
0
      : dict_allocation_(std::move(allocation)) {
194
0
    Populate(from_decompressor, slice);
195
0
  }
196
197
  DecompressorDict(DecompressorDict&& rhs) noexcept
198
      : dict_str_(std::move(rhs.dict_str_)),
199
        dict_allocation_(std::move(rhs.dict_allocation_)),
200
        decompressor_(std::move(rhs.decompressor_)),
201
0
        memory_usage_(std::move(rhs.memory_usage_)) {}
202
203
0
  DecompressorDict& operator=(DecompressorDict&& rhs) noexcept {
204
0
    if (this == &rhs) {
205
0
      return *this;
206
0
    }
207
0
    dict_str_ = std::move(rhs.dict_str_);
208
0
    dict_allocation_ = std::move(rhs.dict_allocation_);
209
0
    decompressor_ = std::move(rhs.decompressor_);
210
0
    return *this;
211
0
  }
212
  // Disable copy
213
  DecompressorDict(const DecompressorDict&) = delete;
214
  DecompressorDict& operator=(const DecompressorDict&) = delete;
215
216
  // The object is self-contained if the string constructor is used, or the
217
  // Slice constructor is invoked with a non-null allocation. Otherwise, it
218
  // is the caller's responsibility to ensure that the underlying storage
219
  // outlives this object.
220
0
  bool own_bytes() const { return !dict_str_.empty() || dict_allocation_; }
221
222
0
  const Slice& GetRawDict() const { return decompressor_->GetSerializedDict(); }
223
224
  // For TypedCacheInterface
225
0
  const Slice& ContentSlice() const { return GetRawDict(); }
226
  static constexpr CacheEntryRole kCacheEntryRole = CacheEntryRole::kOtherBlock;
227
  static constexpr BlockType kBlockType = BlockType::kCompressionDictionary;
228
229
0
  size_t ApproximateMemoryUsage() const { return memory_usage_; }
230
231
 private:
232
  void Populate(Decompressor& from_decompressor, Slice dict);
233
};
234
235
// Holds dictionary and related data, like ZSTD's digested compression
236
// dictionary.
237
struct CompressionDict {
238
#ifdef ZSTD
239
  ZSTD_CDict* zstd_cdict_ = nullptr;
240
#endif  // ZSTD
241
  std::string dict_;
242
243
 public:
244
0
  CompressionDict() = default;
245
0
  CompressionDict(std::string&& dict, CompressionType type, int level) {
246
0
    dict_ = std::move(dict);
247
#ifdef ZSTD
248
    zstd_cdict_ = nullptr;
249
    if (!dict_.empty() && type == kZSTD) {
250
      level = SanitizeZSTDCompressionLevel(level);
251
      // Should be safe (but slower) if below call fails as we'll use the
252
      // raw dictionary to compress.
253
      zstd_cdict_ = ZSTD_createCDict(dict_.data(), dict_.size(), level);
254
      assert(zstd_cdict_ != nullptr);
255
    }
256
#else
257
0
    (void)type;
258
0
    (void)level;
259
0
#endif  // ZSTD
260
0
  }
261
262
0
  CompressionDict(CompressionDict&& other) {
263
#ifdef ZSTD
264
    zstd_cdict_ = other.zstd_cdict_;
265
    other.zstd_cdict_ = nullptr;
266
#endif  // ZSTD
267
0
    dict_ = std::move(other.dict_);
268
0
  }
269
0
  CompressionDict& operator=(CompressionDict&& other) {
270
0
    if (this == &other) {
271
0
      return *this;
272
0
    }
273
0
#ifdef ZSTD
274
0
    zstd_cdict_ = other.zstd_cdict_;
275
0
    other.zstd_cdict_ = nullptr;
276
0
#endif  // ZSTD
277
0
    dict_ = std::move(other.dict_);
278
0
    return *this;
279
0
  }
280
281
0
  ~CompressionDict() {
282
#ifdef ZSTD
283
    size_t res = 0;
284
    if (zstd_cdict_ != nullptr) {
285
      res = ZSTD_freeCDict(zstd_cdict_);
286
    }
287
    assert(res == 0);  // Last I checked they can't fail
288
    (void)res;         // prevent unused var warning
289
#endif                 // ZSTD
290
0
  }
291
292
#ifdef ZSTD
293
  const ZSTD_CDict* GetDigestedZstdCDict() const { return zstd_cdict_; }
294
#endif  // ZSTD
295
296
0
  Slice GetRawDict() const { return dict_; }
297
0
  bool empty() const { return dict_.empty(); }
298
299
0
  static const CompressionDict& GetEmptyDict() {
300
0
    static CompressionDict empty_dict{};
301
0
    return empty_dict;
302
0
  }
303
304
  // Disable copy
305
  CompressionDict(const CompressionDict&) = delete;
306
  CompressionDict& operator=(const CompressionDict&) = delete;
307
};
308
309
class CompressionContext : public Compressor::WorkingArea {
310
 private:
311
#ifdef ZSTD
312
  ZSTD_CCtx* zstd_ctx_ = nullptr;
313
314
  ZSTD_CCtx* CreateZSTDContext() {
315
#ifdef ROCKSDB_ZSTD_CUSTOM_MEM
316
    return ZSTD_createCCtx_advanced(port::GetJeZstdAllocationOverrides());
317
#else   // ROCKSDB_ZSTD_CUSTOM_MEM
318
    return ZSTD_createCCtx();
319
#endif  // ROCKSDB_ZSTD_CUSTOM_MEM
320
  }
321
322
 public:
323
  // callable inside ZSTD_Compress
324
  ZSTD_CCtx* ZSTDPreallocCtx() const {
325
    assert(zstd_ctx_ != nullptr);
326
    return zstd_ctx_;
327
  }
328
329
 private:
330
#endif  // ZSTD
331
332
0
  void CreateNativeContext(CompressionType type, int level, bool checksum) {
333
0
#ifdef ZSTD
334
0
    if (type == kZSTD) {
335
0
      zstd_ctx_ = CreateZSTDContext();
336
0
      level = SanitizeZSTDCompressionLevel(level);
337
0
      size_t err =
338
0
          ZSTD_CCtx_setParameter(zstd_ctx_, ZSTD_c_compressionLevel, level);
339
0
      if (ZSTD_isError(err)) {
340
0
        assert(false);
341
0
        ZSTD_freeCCtx(zstd_ctx_);
342
0
        zstd_ctx_ = CreateZSTDContext();
343
0
      }
344
0
      if (checksum) {
345
0
        err = ZSTD_CCtx_setParameter(zstd_ctx_, ZSTD_c_checksumFlag, 1);
346
0
        if (ZSTD_isError(err)) {
347
0
          assert(false);
348
0
          ZSTD_freeCCtx(zstd_ctx_);
349
0
          zstd_ctx_ = CreateZSTDContext();
350
0
        }
351
0
      }
352
0
    }
353
0
#else
354
0
    (void)type;
355
0
    (void)level;
356
0
    (void)checksum;
357
0
#endif  // ZSTD
358
0
  }
359
0
  void DestroyNativeContext() {
360
0
#ifdef ZSTD
361
0
    if (zstd_ctx_ != nullptr) {
362
0
      ZSTD_freeCCtx(zstd_ctx_);
363
0
    }
364
0
#endif  // ZSTD
365
0
  }
366
367
 public:
368
  explicit CompressionContext(CompressionType type,
369
0
                              const CompressionOptions& options) {
370
0
    CreateNativeContext(type, options.level, options.checksum);
371
0
  }
372
0
  ~CompressionContext() { DestroyNativeContext(); }
373
  CompressionContext(const CompressionContext&) = delete;
374
  CompressionContext& operator=(const CompressionContext&) = delete;
375
};
376
377
// This is like a working area, reusable for different dicts, etc.
378
// TODO: refactor / consolidate
379
class UncompressionContext : public Decompressor::WorkingArea {
380
 private:
381
  CompressionContextCache* ctx_cache_ = nullptr;
382
  ZSTDUncompressCachedData uncomp_cached_data_;
383
384
 public:
385
0
  explicit UncompressionContext(CompressionType type) {
386
0
    if (type == kZSTD) {
387
0
      ctx_cache_ = CompressionContextCache::Instance();
388
0
      uncomp_cached_data_ = ctx_cache_->GetCachedZSTDUncompressData();
389
0
    }
390
0
  }
391
0
  ~UncompressionContext() {
392
0
    if (uncomp_cached_data_.GetCacheIndex() != -1) {
393
0
      assert(ctx_cache_ != nullptr);
394
0
      ctx_cache_->ReturnCachedZSTDUncompressData(
395
0
          uncomp_cached_data_.GetCacheIndex());
396
0
    }
397
0
  }
398
  UncompressionContext(const UncompressionContext&) = delete;
399
  UncompressionContext& operator=(const UncompressionContext&) = delete;
400
401
0
  ZSTDUncompressCachedData::ZSTDNativeContext GetZSTDContext() const {
402
0
    return uncomp_cached_data_.Get();
403
0
  }
404
};
405
406
1.25M
inline bool Snappy_Supported() {
407
#ifdef SNAPPY
408
  return true;
409
#else
410
1.25M
  return false;
411
1.25M
#endif
412
1.25M
}
413
414
53.0k
inline bool Zlib_Supported() {
415
53.0k
#ifdef ZLIB
416
53.0k
  return true;
417
#else
418
  return false;
419
#endif
420
53.0k
}
421
422
53.0k
inline bool BZip2_Supported() {
423
53.0k
#ifdef BZIP2
424
53.0k
  return true;
425
#else
426
  return false;
427
#endif
428
53.0k
}
429
430
1.30M
inline bool LZ4_Supported() {
431
#ifdef LZ4
432
  return true;
433
#else
434
1.30M
  return false;
435
1.30M
#endif
436
1.30M
}
437
438
1.19M
inline CompressionType GetDefaultCompressionType() {
439
1.19M
  return LZ4_Supported()
440
1.19M
             ? kLZ4Compression
441
1.19M
             : (Snappy_Supported() ? kSnappyCompression : kNoCompression);
442
1.19M
}
443
444
53.0k
inline bool XPRESS_Supported() {
445
#ifdef XPRESS
446
  return true;
447
#else
448
53.0k
  return false;
449
53.0k
#endif
450
53.0k
}
451
452
53.0k
inline bool ZSTD_Supported() {
453
#ifdef ZSTD
454
  // NB: ZSTD format is finalized since version 0.8.0. See ZSTD_VERSION_NUMBER
455
  // check above.
456
  return true;
457
#else
458
53.0k
  return false;
459
53.0k
#endif
460
53.0k
}
461
462
0
inline bool ZSTD_Streaming_Supported() {
463
#if defined(ZSTD)
464
  return true;
465
#else
466
0
  return false;
467
0
#endif
468
0
}
469
470
inline bool StreamingCompressionTypeSupported(
471
61.4k
    CompressionType compression_type) {
472
61.4k
  switch (compression_type) {
473
61.4k
    case kNoCompression:
474
61.4k
      return true;
475
0
    case kZSTD:
476
0
      return ZSTD_Streaming_Supported();
477
0
    default:
478
0
      return false;
479
61.4k
  }
480
61.4k
}
481
482
7.32M
inline bool CompressionTypeSupported(CompressionType compression_type) {
483
7.32M
  switch (compression_type) {
484
208k
    case kNoCompression:
485
208k
      return true;
486
53.0k
    case kSnappyCompression:
487
53.0k
      return Snappy_Supported();
488
53.0k
    case kZlibCompression:
489
53.0k
      return Zlib_Supported();
490
53.0k
    case kBZip2Compression:
491
53.0k
      return BZip2_Supported();
492
53.0k
    case kLZ4Compression:
493
53.0k
      return LZ4_Supported();
494
53.0k
    case kLZ4HCCompression:
495
53.0k
      return LZ4_Supported();
496
53.0k
    case kXpressCompression:
497
53.0k
      return XPRESS_Supported();
498
53.0k
    case kZSTD:
499
53.0k
      return ZSTD_Supported();
500
6.74M
    default:  // Including custom compression types
501
6.74M
      return false;
502
7.32M
  }
503
7.32M
}
504
505
0
inline bool DictCompressionTypeSupported(CompressionType compression_type) {
506
0
  switch (compression_type) {
507
0
    case kNoCompression:
508
0
      return false;
509
0
    case kSnappyCompression:
510
0
      return false;
511
0
    case kZlibCompression:
512
0
      return Zlib_Supported();
513
0
    case kBZip2Compression:
514
0
      return false;
515
0
    case kLZ4Compression:
516
0
    case kLZ4HCCompression:
517
#if LZ4_VERSION_NUMBER >= 10400  // r124+
518
      return LZ4_Supported();
519
#else
520
0
      return false;
521
0
#endif
522
0
    case kXpressCompression:
523
0
      return false;
524
0
    case kZSTD:
525
      // NB: dictionary supported since 0.5.0. See ZSTD_VERSION_NUMBER check
526
      // above.
527
0
      return ZSTD_Supported();
528
0
    default:  // Including custom compression types
529
0
      return false;
530
0
  }
531
0
}
532
533
// WART: does not match OptionsHelper::compression_type_string_map
534
std::string CompressionTypeToString(CompressionType compression_type);
535
536
// WART: does not match OptionsHelper::compression_type_string_map
537
CompressionType CompressionTypeFromString(std::string compression_type_str);
538
539
std::string CompressionOptionsToString(
540
    const CompressionOptions& compression_options);
541
542
0
inline bool ZSTD_TrainDictionarySupported() {
543
#ifdef ZSTD
544
  // NB: Dictionary trainer is available since v0.6.1 for static linking, but
545
  // not available for dynamic linking until v1.1.3. See ZSTD_VERSION_NUMBER
546
  // check above.
547
  return true;
548
#else
549
0
  return false;
550
0
#endif
551
0
}
552
553
0
inline bool ZSTD_FinalizeDictionarySupported() {
554
#ifdef ROCKSDB_ZDICT_FINALIZE
555
  return true;
556
#else
557
0
  return false;
558
0
#endif
559
0
}
560
561
// Use to check whether compression types are related or unrelated
562
0
inline CompressionType CanonicalCompressionType(CompressionType type) {
563
0
  switch (type) {
564
0
    // Configuring LZ4 or LZ4HC can result in using the other, depending on
565
0
    // compression level.
566
0
    case kLZ4HCCompression:
567
0
      return kLZ4Compression;
568
0
    default:
569
0
      return type;
570
0
  }
571
0
}
572
573
// The new compression APIs intentionally make it difficult to generate
574
// compressed data larger than the original. (It is better to store the
575
// uncompressed version in that case.) For legacy cases that must store
576
// compressed data even when larger than the uncompressed, this is a convenient
577
// wrapper to support that, with a compressor from BuiltinCompressionManager and
578
// a GrowableBuffer.
579
Status LegacyForceBuiltinCompression(
580
    Compressor& builtin_compressor,
581
    Compressor::ManagedWorkingArea* working_area, Slice from,
582
    GrowableBuffer* to);
583
584
// Records the compression type for subsequent WAL records.
585
class CompressionTypeRecord {
586
 public:
587
  explicit CompressionTypeRecord(CompressionType compression_type)
588
0
      : compression_type_(compression_type) {}
589
590
0
  CompressionType GetCompressionType() const { return compression_type_; }
591
592
0
  inline void EncodeTo(std::string* dst) const {
593
0
    assert(dst != nullptr);
594
0
    PutFixed32(dst, compression_type_);
595
0
  }
596
597
0
  inline Status DecodeFrom(Slice* src) {
598
0
    constexpr char class_name[] = "CompressionTypeRecord";
599
600
0
    uint32_t val;
601
0
    if (!GetFixed32(src, &val)) {
602
0
      return Status::Corruption(class_name,
603
0
                                "Error decoding WAL compression type");
604
0
    }
605
0
    CompressionType compression_type = static_cast<CompressionType>(val);
606
0
    if (!StreamingCompressionTypeSupported(compression_type)) {
607
0
      return Status::Corruption(class_name,
608
0
                                "WAL compression type not supported");
609
0
    }
610
0
    compression_type_ = compression_type;
611
0
    return Status::OK();
612
0
  }
613
614
0
  inline std::string DebugString() const {
615
0
    return "compression_type: " + CompressionTypeToString(compression_type_);
616
0
  }
617
618
 private:
619
  CompressionType compression_type_;
620
};
621
622
// Base class to implement compression for a stream of buffers.
623
// Instantiate an implementation of the class using Create() with the
624
// compression type and use Compress() repeatedly.
625
// The output buffer needs to be at least max_output_len.
626
// Call Reset() in between frame boundaries or in case of an error.
627
// NOTE: This class is not thread safe.
628
class StreamingCompress {
629
 public:
630
  StreamingCompress(CompressionType compression_type,
631
                    const CompressionOptions& opts,
632
                    uint32_t compress_format_version, size_t max_output_len)
633
0
      : compression_type_(compression_type),
634
0
        opts_(opts),
635
0
        compress_format_version_(compress_format_version),
636
0
        max_output_len_(max_output_len) {}
637
0
  virtual ~StreamingCompress() = default;
638
  // compress should be called repeatedly with the same input till the method
639
  // returns 0
640
  // Parameters:
641
  // input - buffer to compress
642
  // input_size - size of input buffer
643
  // output - compressed buffer allocated by caller, should be at least
644
  // max_output_len
645
  // output_size - size of the output buffer
646
  // Returns -1 for errors, the remaining size of the input buffer that needs
647
  // to be compressed
648
  virtual int Compress(const char* input, size_t input_size, char* output,
649
                       size_t* output_pos) = 0;
650
  // static method to create object of a class inherited from
651
  // StreamingCompress based on the actual compression type.
652
  static std::unique_ptr<StreamingCompress> Create(
653
      CompressionType compression_type, const CompressionOptions& opts,
654
      uint32_t compress_format_version, size_t max_output_len);
655
  virtual void Reset() = 0;
656
657
 protected:
658
  const CompressionType compression_type_;
659
  const CompressionOptions opts_;
660
  const uint32_t compress_format_version_;
661
  const size_t max_output_len_;
662
};
663
664
// Base class to uncompress a stream of compressed buffers.
665
// Instantiate an implementation of the class using Create() with the
666
// compression type and use Uncompress() repeatedly.
667
// The output buffer needs to be at least max_output_len.
668
// Call Reset() in between frame boundaries or in case of an error.
669
// NOTE: This class is not thread safe.
670
class StreamingUncompress {
671
 public:
672
  StreamingUncompress(CompressionType compression_type,
673
                      uint32_t compress_format_version, size_t max_output_len)
674
0
      : compression_type_(compression_type),
675
0
        compress_format_version_(compress_format_version),
676
0
        max_output_len_(max_output_len) {}
677
0
  virtual ~StreamingUncompress() = default;
678
  // Uncompress can be called repeatedly to progressively process the same
679
  // input buffer, or can be called with a new input buffer. When the input
680
  // buffer is not fully consumed, the return value is > 0 or output_size
681
  // == max_output_len. When calling uncompress to continue processing the
682
  // same input buffer, the input argument should be nullptr.
683
  // Parameters:
684
  // input - buffer to uncompress
685
  // input_size - size of input buffer
686
  // output - uncompressed buffer allocated by caller, should be at least
687
  // max_output_len
688
  // output_size - size of the output buffer
689
  // Returns -1 for errors, remaining input to be processed otherwise.
690
  virtual int Uncompress(const char* input, size_t input_size, char* output,
691
                         size_t* output_pos) = 0;
692
  static std::unique_ptr<StreamingUncompress> Create(
693
      CompressionType compression_type, uint32_t compress_format_version,
694
      size_t max_output_len);
695
  virtual void Reset() = 0;
696
697
 protected:
698
  CompressionType compression_type_;
699
  uint32_t compress_format_version_;
700
  size_t max_output_len_;
701
};
702
703
class ZSTDStreamingCompress final : public StreamingCompress {
704
 public:
705
  explicit ZSTDStreamingCompress(const CompressionOptions& opts,
706
                                 uint32_t compress_format_version,
707
                                 size_t max_output_len)
708
0
      : StreamingCompress(kZSTD, opts, compress_format_version,
709
0
                          max_output_len) {
710
#ifdef ZSTD
711
    cctx_ = ZSTD_createCCtx();
712
    // Each compressed frame will have a checksum
713
    ZSTD_CCtx_setParameter(cctx_, ZSTD_c_checksumFlag, 1);
714
    assert(cctx_ != nullptr);
715
    input_buffer_ = {/*src=*/nullptr, /*size=*/0, /*pos=*/0};
716
#endif
717
0
  }
718
0
  ~ZSTDStreamingCompress() override {
719
0
#ifdef ZSTD
720
0
    ZSTD_freeCCtx(cctx_);
721
0
#endif
722
0
  }
723
  int Compress(const char* input, size_t input_size, char* output,
724
               size_t* output_pos) override;
725
  void Reset() override;
726
#ifdef ZSTD
727
  ZSTD_CCtx* cctx_;
728
  ZSTD_inBuffer input_buffer_;
729
#endif
730
};
731
732
class ZSTDStreamingUncompress final : public StreamingUncompress {
733
 public:
734
  explicit ZSTDStreamingUncompress(uint32_t compress_format_version,
735
                                   size_t max_output_len)
736
0
      : StreamingUncompress(kZSTD, compress_format_version, max_output_len) {
737
#ifdef ZSTD
738
    dctx_ = ZSTD_createDCtx();
739
    assert(dctx_ != nullptr);
740
    input_buffer_ = {/*src=*/nullptr, /*size=*/0, /*pos=*/0};
741
#endif
742
0
  }
743
0
  ~ZSTDStreamingUncompress() override {
744
0
#ifdef ZSTD
745
0
    ZSTD_freeDCtx(dctx_);
746
0
#endif
747
0
  }
748
  int Uncompress(const char* input, size_t input_size, char* output,
749
                 size_t* output_size) override;
750
  void Reset() override;
751
752
 private:
753
#ifdef ZSTD
754
  ZSTD_DCtx* dctx_;
755
  ZSTD_inBuffer input_buffer_;
756
#endif
757
};
758
759
}  // namespace ROCKSDB_NAMESPACE