Coverage Report

Created: 2026-11-12 08:06

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
  // No copying allowed
148
  FailureDecompressor(const FailureDecompressor&) = delete;
149
  FailureDecompressor& operator=(const FailureDecompressor&) = delete;
150
  FailureDecompressor(FailureDecompressor&&) = delete;
151
  FailureDecompressor& operator=(FailureDecompressor&&) = delete;
152
153
0
  const char* Name() const override { return "FailureDecompressor"; }
154
155
0
  Status ExtractUncompressedSize(Args& /*args*/) override { return status_; }
156
157
  Status DecompressBlock(const Args& /*args*/,
158
0
                         char* /*uncompressed_output*/) override {
159
0
    return status_;
160
0
  }
161
162
 protected:
163
  Status status_;
164
};
165
166
// Owns a decompression dictionary, and associated Decompressor, for storing
167
// in the block cache.
168
//
169
// Justification: for a "processed" dictionary to be saved in block cache, we
170
// also need a reference to the decompressor that processed it, to ensure it
171
// is recognized properly. At that point, we might as well have the dictionary
172
// part of the decompressor identity and track an associated decompressor along
173
// with a decompression dictionary in the block cache, and the decompressor
174
// hides potential details of processing the dictionary.
175
struct DecompressorDict {
176
  // Block containing the data for the compression dictionary in case the
177
  // constructor that takes a string parameter is used.
178
  std::string dict_str_;
179
180
  // Block containing the data for the compression dictionary in case the
181
  // constructor that takes a Slice parameter is used and the passed in
182
  // CacheAllocationPtr is not nullptr.
183
  CacheAllocationPtr dict_allocation_;
184
185
  // A Decompressor referencing and using the dictionary owned by this.
186
  std::unique_ptr<Decompressor> decompressor_;
187
188
  // Approximate owned memory usage
189
  size_t memory_usage_;
190
191
  DecompressorDict(std::string&& dict, Decompressor& from_decompressor)
192
0
      : dict_str_(std::move(dict)) {
193
0
    Populate(from_decompressor, dict_str_);
194
0
  }
195
196
  DecompressorDict(Slice slice, CacheAllocationPtr&& allocation,
197
                   Decompressor& from_decompressor)
198
0
      : dict_allocation_(std::move(allocation)) {
199
0
    Populate(from_decompressor, slice);
200
0
  }
201
202
  DecompressorDict(DecompressorDict&& rhs) noexcept
203
      : dict_str_(std::move(rhs.dict_str_)),
204
        dict_allocation_(std::move(rhs.dict_allocation_)),
205
        decompressor_(std::move(rhs.decompressor_)),
206
0
        memory_usage_(std::move(rhs.memory_usage_)) {}
207
208
0
  DecompressorDict& operator=(DecompressorDict&& rhs) noexcept {
209
0
    if (this == &rhs) {
210
0
      return *this;
211
0
    }
212
0
    dict_str_ = std::move(rhs.dict_str_);
213
0
    dict_allocation_ = std::move(rhs.dict_allocation_);
214
0
    decompressor_ = std::move(rhs.decompressor_);
215
0
    return *this;
216
0
  }
217
  // Disable copy
218
  DecompressorDict(const DecompressorDict&) = delete;
219
  DecompressorDict& operator=(const DecompressorDict&) = delete;
220
0
  ~DecompressorDict() = default;
221
222
  // The object is self-contained if the string constructor is used, or the
223
  // Slice constructor is invoked with a non-null allocation. Otherwise, it
224
  // is the caller's responsibility to ensure that the underlying storage
225
  // outlives this object.
226
0
  bool own_bytes() const { return !dict_str_.empty() || dict_allocation_; }
227
228
0
  const Slice& GetRawDict() const { return decompressor_->GetSerializedDict(); }
229
230
  // For TypedCacheInterface
231
0
  const Slice& ContentSlice() const { return GetRawDict(); }
232
  static constexpr CacheEntryRole kCacheEntryRole = CacheEntryRole::kOtherBlock;
233
  static constexpr BlockType kBlockType = BlockType::kCompressionDictionary;
234
235
0
  size_t ApproximateMemoryUsage() const { return memory_usage_; }
236
237
 private:
238
  void Populate(Decompressor& from_decompressor, Slice dict);
239
};
240
241
// Holds dictionary and related data, like ZSTD's digested compression
242
// dictionary.
243
struct CompressionDict {
244
#ifdef ZSTD
245
  ZSTD_CDict* zstd_cdict_ = nullptr;
246
#endif  // ZSTD
247
  std::string dict_;
248
249
 public:
250
0
  CompressionDict() = default;
251
0
  CompressionDict(std::string&& dict, CompressionType type, int level) {
252
0
    dict_ = std::move(dict);
253
#ifdef ZSTD
254
    zstd_cdict_ = nullptr;
255
    if (!dict_.empty() && type == kZSTD) {
256
      level = SanitizeZSTDCompressionLevel(level);
257
      // Should be safe (but slower) if below call fails as we'll use the
258
      // raw dictionary to compress.
259
      zstd_cdict_ = ZSTD_createCDict(dict_.data(), dict_.size(), level);
260
      assert(zstd_cdict_ != nullptr);
261
    }
262
#else
263
0
    (void)type;
264
0
    (void)level;
265
0
#endif  // ZSTD
266
0
  }
267
268
0
  CompressionDict(CompressionDict&& other) {
269
#ifdef ZSTD
270
    zstd_cdict_ = other.zstd_cdict_;
271
    other.zstd_cdict_ = nullptr;
272
#endif  // ZSTD
273
0
    dict_ = std::move(other.dict_);
274
0
  }
275
0
  CompressionDict& operator=(CompressionDict&& other) {
276
0
    if (this == &other) {
277
0
      return *this;
278
0
    }
279
0
#ifdef ZSTD
280
0
    zstd_cdict_ = other.zstd_cdict_;
281
0
    other.zstd_cdict_ = nullptr;
282
0
#endif  // ZSTD
283
0
    dict_ = std::move(other.dict_);
284
0
    return *this;
285
0
  }
286
287
0
  ~CompressionDict() {
288
#ifdef ZSTD
289
    size_t res = 0;
290
    if (zstd_cdict_ != nullptr) {
291
      res = ZSTD_freeCDict(zstd_cdict_);
292
    }
293
    assert(res == 0);  // Last I checked they can't fail
294
    (void)res;         // prevent unused var warning
295
#endif                 // ZSTD
296
0
  }
297
298
#ifdef ZSTD
299
  const ZSTD_CDict* GetDigestedZstdCDict() const { return zstd_cdict_; }
300
#endif  // ZSTD
301
302
0
  Slice GetRawDict() const { return dict_; }
303
0
  bool empty() const { return dict_.empty(); }
304
305
0
  static const CompressionDict& GetEmptyDict() {
306
0
    static CompressionDict empty_dict{};
307
0
    return empty_dict;
308
0
  }
309
310
  // Disable copy
311
  CompressionDict(const CompressionDict&) = delete;
312
  CompressionDict& operator=(const CompressionDict&) = delete;
313
};
314
315
class CompressionContext : public Compressor::WorkingArea {
316
 private:
317
#ifdef ZSTD
318
  ZSTD_CCtx* zstd_ctx_ = nullptr;
319
320
  ZSTD_CCtx* CreateZSTDContext() {
321
#ifdef ROCKSDB_ZSTD_CUSTOM_MEM
322
    return ZSTD_createCCtx_advanced(port::GetJeZstdAllocationOverrides());
323
#else   // ROCKSDB_ZSTD_CUSTOM_MEM
324
    return ZSTD_createCCtx();
325
#endif  // ROCKSDB_ZSTD_CUSTOM_MEM
326
  }
327
328
 public:
329
  // callable inside ZSTD_Compress
330
  ZSTD_CCtx* ZSTDPreallocCtx() const {
331
    assert(zstd_ctx_ != nullptr);
332
    return zstd_ctx_;
333
  }
334
335
 private:
336
#endif  // ZSTD
337
338
0
  void CreateNativeContext(CompressionType type, int level, bool checksum) {
339
0
#ifdef ZSTD
340
0
    if (type == kZSTD) {
341
0
      zstd_ctx_ = CreateZSTDContext();
342
0
      level = SanitizeZSTDCompressionLevel(level);
343
0
      size_t err =
344
0
          ZSTD_CCtx_setParameter(zstd_ctx_, ZSTD_c_compressionLevel, level);
345
0
      if (ZSTD_isError(err)) {
346
0
        assert(false);
347
0
        ZSTD_freeCCtx(zstd_ctx_);
348
0
        zstd_ctx_ = CreateZSTDContext();
349
0
      }
350
0
      if (checksum) {
351
0
        err = ZSTD_CCtx_setParameter(zstd_ctx_, ZSTD_c_checksumFlag, 1);
352
0
        if (ZSTD_isError(err)) {
353
0
          assert(false);
354
0
          ZSTD_freeCCtx(zstd_ctx_);
355
0
          zstd_ctx_ = CreateZSTDContext();
356
0
        }
357
0
      }
358
0
    }
359
0
#else
360
0
    (void)type;
361
0
    (void)level;
362
0
    (void)checksum;
363
0
#endif  // ZSTD
364
0
  }
365
0
  void DestroyNativeContext() {
366
0
#ifdef ZSTD
367
0
    if (zstd_ctx_ != nullptr) {
368
0
      ZSTD_freeCCtx(zstd_ctx_);
369
0
    }
370
0
#endif  // ZSTD
371
0
  }
372
373
 public:
374
  explicit CompressionContext(CompressionType type,
375
0
                              const CompressionOptions& options) {
376
0
    CreateNativeContext(type, options.level, options.checksum);
377
0
  }
378
0
  ~CompressionContext() { DestroyNativeContext(); }
379
  CompressionContext(const CompressionContext&) = delete;
380
  CompressionContext& operator=(const CompressionContext&) = delete;
381
  CompressionContext(CompressionContext&&) = delete;
382
  CompressionContext& operator=(CompressionContext&&) = delete;
383
};
384
385
// This is like a working area, reusable for different dicts, etc.
386
// TODO: refactor / consolidate
387
class UncompressionContext : public Decompressor::WorkingArea {
388
 private:
389
  CompressionContextCache* ctx_cache_ = nullptr;
390
  ZSTDUncompressCachedData uncomp_cached_data_;
391
392
 public:
393
0
  explicit UncompressionContext(CompressionType type) {
394
0
    if (type == kZSTD) {
395
0
      ctx_cache_ = CompressionContextCache::Instance();
396
0
      uncomp_cached_data_ = ctx_cache_->GetCachedZSTDUncompressData();
397
0
    }
398
0
  }
399
0
  ~UncompressionContext() {
400
0
    if (uncomp_cached_data_.GetCacheIndex() != -1) {
401
0
      assert(ctx_cache_ != nullptr);
402
0
      ctx_cache_->ReturnCachedZSTDUncompressData(
403
0
          uncomp_cached_data_.GetCacheIndex());
404
0
    }
405
0
  }
406
  UncompressionContext(const UncompressionContext&) = delete;
407
  UncompressionContext& operator=(const UncompressionContext&) = delete;
408
  UncompressionContext(UncompressionContext&&) = delete;
409
  UncompressionContext& operator=(UncompressionContext&&) = delete;
410
411
0
  ZSTDUncompressCachedData::ZSTDNativeContext GetZSTDContext() const {
412
0
    return uncomp_cached_data_.Get();
413
0
  }
414
};
415
416
1.04M
inline bool Snappy_Supported() {
417
#ifdef SNAPPY
418
  return true;
419
#else
420
1.04M
  return false;
421
1.04M
#endif
422
1.04M
}
423
424
44.8k
inline bool Zlib_Supported() {
425
44.8k
#ifdef ZLIB
426
44.8k
  return true;
427
#else
428
  return false;
429
#endif
430
44.8k
}
431
432
44.8k
inline bool BZip2_Supported() {
433
44.8k
#ifdef BZIP2
434
44.8k
  return true;
435
#else
436
  return false;
437
#endif
438
44.8k
}
439
440
1.08M
inline bool LZ4_Supported() {
441
#ifdef LZ4
442
  return true;
443
#else
444
1.08M
  return false;
445
1.08M
#endif
446
1.08M
}
447
448
996k
inline CompressionType GetDefaultCompressionType() {
449
996k
  return LZ4_Supported()
450
996k
             ? kLZ4Compression
451
996k
             : (Snappy_Supported() ? kSnappyCompression : kNoCompression);
452
996k
}
453
454
44.8k
inline bool XPRESS_Supported() {
455
#ifdef XPRESS
456
  return true;
457
#else
458
44.8k
  return false;
459
44.8k
#endif
460
44.8k
}
461
462
44.8k
inline bool ZSTD_Supported() {
463
#ifdef ZSTD
464
  // NB: ZSTD format is finalized since version 0.8.0. See ZSTD_VERSION_NUMBER
465
  // check above.
466
  return true;
467
#else
468
44.8k
  return false;
469
44.8k
#endif
470
44.8k
}
471
472
0
inline bool ZSTD_Streaming_Supported() {
473
#if defined(ZSTD)
474
  return true;
475
#else
476
0
  return false;
477
0
#endif
478
0
}
479
480
inline bool StreamingCompressionTypeSupported(
481
53.0k
    CompressionType compression_type) {
482
53.0k
  switch (compression_type) {
483
53.0k
    case kNoCompression:
484
53.0k
      return true;
485
0
    case kZSTD:
486
0
      return ZSTD_Streaming_Supported();
487
0
    default:
488
0
      return false;
489
53.0k
  }
490
53.0k
}
491
492
6.18M
inline bool CompressionTypeSupported(CompressionType compression_type) {
493
6.18M
  switch (compression_type) {
494
172k
    case kNoCompression:
495
172k
      return true;
496
44.8k
    case kSnappyCompression:
497
44.8k
      return Snappy_Supported();
498
44.8k
    case kZlibCompression:
499
44.8k
      return Zlib_Supported();
500
44.8k
    case kBZip2Compression:
501
44.8k
      return BZip2_Supported();
502
44.8k
    case kLZ4Compression:
503
44.8k
      return LZ4_Supported();
504
44.8k
    case kLZ4HCCompression:
505
44.8k
      return LZ4_Supported();
506
44.8k
    case kXpressCompression:
507
44.8k
      return XPRESS_Supported();
508
44.8k
    case kZSTD:
509
44.8k
      return ZSTD_Supported();
510
5.69M
    default:  // Including custom compression types
511
5.69M
      return false;
512
6.18M
  }
513
6.18M
}
514
515
0
inline bool DictCompressionTypeSupported(CompressionType compression_type) {
516
0
  switch (compression_type) {
517
0
    case kNoCompression:
518
0
      return false;
519
0
    case kSnappyCompression:
520
0
      return false;
521
0
    case kZlibCompression:
522
0
      return Zlib_Supported();
523
0
    case kBZip2Compression:
524
0
      return false;
525
0
    case kLZ4Compression:
526
0
    case kLZ4HCCompression:
527
#if LZ4_VERSION_NUMBER >= 10400  // r124+
528
      return LZ4_Supported();
529
#else
530
0
      return false;
531
0
#endif
532
0
    case kXpressCompression:
533
0
      return false;
534
0
    case kZSTD:
535
      // NB: dictionary supported since 0.5.0. See ZSTD_VERSION_NUMBER check
536
      // above.
537
0
      return ZSTD_Supported();
538
0
    default:  // Including custom compression types
539
0
      return false;
540
0
  }
541
0
}
542
543
// WART: does not match OptionsHelper::compression_type_string_map
544
std::string CompressionTypeToString(CompressionType compression_type);
545
546
// WART: does not match OptionsHelper::compression_type_string_map
547
CompressionType CompressionTypeFromString(std::string compression_type_str);
548
549
std::string CompressionOptionsToString(
550
    const CompressionOptions& compression_options);
551
552
0
inline bool ZSTD_TrainDictionarySupported() {
553
#ifdef ZSTD
554
  // NB: Dictionary trainer is available since v0.6.1 for static linking, but
555
  // not available for dynamic linking until v1.1.3. See ZSTD_VERSION_NUMBER
556
  // check above.
557
  return true;
558
#else
559
0
  return false;
560
0
#endif
561
0
}
562
563
0
inline bool ZSTD_FinalizeDictionarySupported() {
564
#ifdef ROCKSDB_ZDICT_FINALIZE
565
  return true;
566
#else
567
0
  return false;
568
0
#endif
569
0
}
570
571
// Use to check whether compression types are related or unrelated
572
0
inline CompressionType CanonicalCompressionType(CompressionType type) {
573
0
  switch (type) {
574
0
    // Configuring LZ4 or LZ4HC can result in using the other, depending on
575
0
    // compression level.
576
0
    case kLZ4HCCompression:
577
0
      return kLZ4Compression;
578
0
    default:
579
0
      return type;
580
0
  }
581
0
}
582
583
// The new compression APIs intentionally make it difficult to generate
584
// compressed data larger than the original. (It is better to store the
585
// uncompressed version in that case.) For legacy cases that must store
586
// compressed data even when larger than the uncompressed, this is a convenient
587
// wrapper to support that, with a compressor from BuiltinCompressionManager and
588
// a GrowableBuffer.
589
Status LegacyForceBuiltinCompression(
590
    Compressor& builtin_compressor,
591
    Compressor::ManagedWorkingArea* working_area, Slice from,
592
    GrowableBuffer* to);
593
594
// Records the compression type for subsequent WAL records.
595
class CompressionTypeRecord {
596
 public:
597
  explicit CompressionTypeRecord(CompressionType compression_type)
598
0
      : compression_type_(compression_type) {}
599
600
0
  CompressionType GetCompressionType() const { return compression_type_; }
601
602
0
  inline void EncodeTo(std::string* dst) const {
603
0
    assert(dst != nullptr);
604
0
    PutFixed32(dst, compression_type_);
605
0
  }
606
607
0
  inline Status DecodeFrom(Slice* src) {
608
0
    constexpr char class_name[] = "CompressionTypeRecord";
609
610
0
    uint32_t val;
611
0
    if (!GetFixed32(src, &val)) {
612
0
      return Status::Corruption(class_name,
613
0
                                "Error decoding WAL compression type");
614
0
    }
615
0
    CompressionType compression_type = static_cast<CompressionType>(val);
616
0
    if (!StreamingCompressionTypeSupported(compression_type)) {
617
0
      return Status::Corruption(class_name,
618
0
                                "WAL compression type not supported");
619
0
    }
620
0
    compression_type_ = compression_type;
621
0
    return Status::OK();
622
0
  }
623
624
0
  inline std::string DebugString() const {
625
0
    return "compression_type: " + CompressionTypeToString(compression_type_);
626
0
  }
627
628
 private:
629
  CompressionType compression_type_;
630
};
631
632
// Base class to implement compression for a stream of buffers.
633
// Instantiate an implementation of the class using Create() with the
634
// compression type and use Compress() repeatedly.
635
// The output buffer needs to be at least max_output_len.
636
// Call Reset() in between frame boundaries or in case of an error.
637
// NOTE: This class is not thread safe.
638
class StreamingCompress {
639
 public:
640
  StreamingCompress(CompressionType compression_type,
641
                    const CompressionOptions& opts,
642
                    uint32_t compress_format_version, size_t max_output_len)
643
0
      : compression_type_(compression_type),
644
0
        opts_(opts),
645
0
        compress_format_version_(compress_format_version),
646
0
        max_output_len_(max_output_len) {}
647
0
  virtual ~StreamingCompress() = default;
648
  // compress should be called repeatedly with the same input till the method
649
  // returns 0
650
  // Parameters:
651
  // input - buffer to compress
652
  // input_size - size of input buffer
653
  // output - compressed buffer allocated by caller, should be at least
654
  // max_output_len
655
  // output_size - size of the output buffer
656
  // Returns -1 for errors, the remaining size of the input buffer that needs
657
  // to be compressed
658
  virtual int Compress(const char* input, size_t input_size, char* output,
659
                       size_t* output_pos) = 0;
660
  // static method to create object of a class inherited from
661
  // StreamingCompress based on the actual compression type.
662
  static std::unique_ptr<StreamingCompress> Create(
663
      CompressionType compression_type, const CompressionOptions& opts,
664
      uint32_t compress_format_version, size_t max_output_len);
665
  virtual void Reset() = 0;
666
667
 protected:
668
  const CompressionType compression_type_;
669
  const CompressionOptions opts_;
670
  const uint32_t compress_format_version_;
671
  const size_t max_output_len_;
672
};
673
674
// Base class to uncompress a stream of compressed buffers.
675
// Instantiate an implementation of the class using Create() with the
676
// compression type and use Uncompress() repeatedly.
677
// The output buffer needs to be at least max_output_len.
678
// Call Reset() in between frame boundaries or in case of an error.
679
// NOTE: This class is not thread safe.
680
class StreamingUncompress {
681
 public:
682
  StreamingUncompress(CompressionType compression_type,
683
                      uint32_t compress_format_version, size_t max_output_len)
684
0
      : compression_type_(compression_type),
685
0
        compress_format_version_(compress_format_version),
686
0
        max_output_len_(max_output_len) {}
687
0
  virtual ~StreamingUncompress() = default;
688
  // Uncompress can be called repeatedly to progressively process the same
689
  // input buffer, or can be called with a new input buffer. When the input
690
  // buffer is not fully consumed, the return value is > 0 or output_size
691
  // == max_output_len. When calling uncompress to continue processing the
692
  // same input buffer, the input argument should be nullptr.
693
  // Parameters:
694
  // input - buffer to uncompress
695
  // input_size - size of input buffer
696
  // output - uncompressed buffer allocated by caller, should be at least
697
  // max_output_len
698
  // output_size - size of the output buffer
699
  // Returns -1 for errors, remaining input to be processed otherwise.
700
  virtual int Uncompress(const char* input, size_t input_size, char* output,
701
                         size_t* output_pos) = 0;
702
  static std::unique_ptr<StreamingUncompress> Create(
703
      CompressionType compression_type, uint32_t compress_format_version,
704
      size_t max_output_len);
705
  virtual void Reset() = 0;
706
707
 protected:
708
  CompressionType compression_type_;
709
  uint32_t compress_format_version_;
710
  size_t max_output_len_;
711
};
712
713
class ZSTDStreamingCompress final : public StreamingCompress {
714
 public:
715
  explicit ZSTDStreamingCompress(const CompressionOptions& opts,
716
                                 uint32_t compress_format_version,
717
                                 size_t max_output_len)
718
0
      : StreamingCompress(kZSTD, opts, compress_format_version,
719
0
                          max_output_len) {
720
#ifdef ZSTD
721
    cctx_ = ZSTD_createCCtx();
722
    // Each compressed frame will have a checksum
723
    ZSTD_CCtx_setParameter(cctx_, ZSTD_c_checksumFlag, 1);
724
    assert(cctx_ != nullptr);
725
    input_buffer_ = {/*src=*/nullptr, /*size=*/0, /*pos=*/0};
726
#endif
727
0
  }
728
0
  ~ZSTDStreamingCompress() override {
729
0
#ifdef ZSTD
730
0
    ZSTD_freeCCtx(cctx_);
731
0
#endif
732
0
  }
733
  // No copying allowed
734
  ZSTDStreamingCompress(const ZSTDStreamingCompress&) = delete;
735
  ZSTDStreamingCompress& operator=(const ZSTDStreamingCompress&) = delete;
736
  ZSTDStreamingCompress(ZSTDStreamingCompress&&) = delete;
737
  ZSTDStreamingCompress& operator=(ZSTDStreamingCompress&&) = delete;
738
  int Compress(const char* input, size_t input_size, char* output,
739
               size_t* output_pos) override;
740
  void Reset() override;
741
#ifdef ZSTD
742
  ZSTD_CCtx* cctx_;
743
  ZSTD_inBuffer input_buffer_;
744
#endif
745
};
746
747
class ZSTDStreamingUncompress final : public StreamingUncompress {
748
 public:
749
  explicit ZSTDStreamingUncompress(uint32_t compress_format_version,
750
                                   size_t max_output_len)
751
0
      : StreamingUncompress(kZSTD, compress_format_version, max_output_len) {
752
#ifdef ZSTD
753
    dctx_ = ZSTD_createDCtx();
754
    assert(dctx_ != nullptr);
755
    input_buffer_ = {/*src=*/nullptr, /*size=*/0, /*pos=*/0};
756
#endif
757
0
  }
758
0
  ~ZSTDStreamingUncompress() override {
759
0
#ifdef ZSTD
760
0
    ZSTD_freeDCtx(dctx_);
761
0
#endif
762
0
  }
763
  // No copying allowed
764
  ZSTDStreamingUncompress(const ZSTDStreamingUncompress&) = delete;
765
  ZSTDStreamingUncompress& operator=(const ZSTDStreamingUncompress&) = delete;
766
  ZSTDStreamingUncompress(ZSTDStreamingUncompress&&) = delete;
767
  ZSTDStreamingUncompress& operator=(ZSTDStreamingUncompress&&) = delete;
768
  int Uncompress(const char* input, size_t input_size, char* output,
769
                 size_t* output_size) override;
770
  void Reset() override;
771
772
 private:
773
#ifdef ZSTD
774
  ZSTD_DCtx* dctx_;
775
  ZSTD_inBuffer input_buffer_;
776
#endif
777
};
778
779
}  // namespace ROCKSDB_NAMESPACE