@nxtedition/rocksdb 13.5.13 → 14.0.0

package/deps/rocksdb/rocksdb/cache/compressed_secondary_cache.cc

@@ -16,6 +16,31 @@
 #include "util/string_util.h"
 
 namespace ROCKSDB_NAMESPACE {
+namespace {
+// Format of values in CompressedSecondaryCache:
+// If enable_custom_split_merge:
+//  * A chain of CacheValueChunk representing the sequence of bytes for a tagged
+//    value. The overall length of the tagged value is determined by the chain
+//    of CacheValueChunks.
+// If !enable_custom_split_merge:
+//  * A LengthPrefixedSlice (starts with varint64 size) of a tagged value.
+//
+// A tagged value has a 2-byte header before the "saved" or compressed block
+// data:
+//  * 1 byte for "source" CacheTier indicating which tier is responsible for
+//    compression/decompression.
+//  * 1 byte for compression type which is generated/used by
+//    CompressedSecondaryCache iff source == CacheTier::kVolatileCompressedTier
+//    (original entry passed in was uncompressed). Otherwise, the compression
+//    type is preserved from the entry passed in.
+constexpr uint32_t kTagSize = 2;
+
+// Size of tag + varint size prefix when applicable
+uint32_t GetHeaderSize(size_t data_size, bool enable_split_merge) {
+  return (enable_split_merge ? 0 : VarintLength(kTagSize + data_size)) +
+         kTagSize;
+}
+}  // namespace
 
 CompressedSecondaryCache::CompressedSecondaryCache(
     const CompressedSecondaryCacheOptions& opts)
@@ -40,13 +65,9 @@ std::unique_ptr<SecondaryCacheResultHandle> CompressedSecondaryCache::Lookup(
     Cache::CreateContext* create_context, bool /*wait*/, bool advise_erase,
     Statistics* stats, bool& kept_in_sec_cache) {
   assert(helper);
-  // This is a minor optimization. Its ok to skip it in TSAN in order to
-  // avoid a false positive.
-#ifndef __SANITIZE_THREAD__
-  if (disable_cache_) {
+  if (disable_cache_.LoadRelaxed()) {
     return nullptr;
   }
-#endif
 
   std::unique_ptr<SecondaryCacheResultHandle> handle;
   kept_in_sec_cache = false;
@@ -62,75 +83,58 @@ std::unique_ptr<SecondaryCacheResultHandle> CompressedSecondaryCache::Lookup(
     return nullptr;
   }
 
-  CacheAllocationPtr* ptr{nullptr};
-  CacheAllocationPtr merged_value;
-  size_t handle_value_charge{0};
-  const char* data_ptr = nullptr;
-  CacheTier source = CacheTier::kVolatileCompressedTier;
-  CompressionType type = cache_options_.compression_type;
+  std::string merged_value;
+  Slice tagged_data;
   if (cache_options_.enable_custom_split_merge) {
     CacheValueChunk* value_chunk_ptr =
-        reinterpret_cast<CacheValueChunk*>(handle_value);
-    merged_value = MergeChunksIntoValue(value_chunk_ptr, handle_value_charge);
-    ptr = &merged_value;
-    data_ptr = ptr->get();
+        static_cast<CacheValueChunk*>(handle_value);
+    merged_value = MergeChunksIntoValue(value_chunk_ptr);
+    tagged_data = Slice(merged_value);
   } else {
-    uint32_t type_32 = static_cast<uint32_t>(type);
-    uint32_t source_32 = static_cast<uint32_t>(source);
-    ptr = reinterpret_cast<CacheAllocationPtr*>(handle_value);
-    handle_value_charge = cache_->GetCharge(lru_handle);
-    data_ptr = ptr->get();
-    const char* limit = ptr->get() + handle_value_charge;
-    data_ptr =
-        GetVarint32Ptr(data_ptr, limit, static_cast<uint32_t*>(&type_32));
-    type = static_cast<CompressionType>(type_32);
-    data_ptr =
-        GetVarint32Ptr(data_ptr, limit, static_cast<uint32_t*>(&source_32));
-    source = static_cast<CacheTier>(source_32);
-    uint64_t data_size = 0;
-    data_ptr =
-        GetVarint64Ptr(data_ptr, limit, static_cast<uint64_t*>(&data_size));
-    assert(handle_value_charge > data_size);
-    handle_value_charge = data_size;
+    tagged_data = GetLengthPrefixedSlice(static_cast<char*>(handle_value));
   }
-  MemoryAllocator* allocator = cache_options_.memory_allocator.get();
 
-  Status s;
-  Cache::ObjectPtr value{nullptr};
-  size_t charge{0};
+  auto source = lossless_cast<CacheTier>(tagged_data[0]);
+  auto type = lossless_cast<CompressionType>(tagged_data[1]);
+
+  std::unique_ptr<char[]> uncompressed;
+  Slice saved(tagged_data.data() + kTagSize, tagged_data.size() - kTagSize);
   if (source == CacheTier::kVolatileCompressedTier) {
-    if (cache_options_.compression_type == kNoCompression ||
-        cache_options_.do_not_compress_roles.Contains(helper->role)) {
-      s = helper->create_cb(Slice(data_ptr, handle_value_charge),
-                            kNoCompression, CacheTier::kVolatileTier,
-                            create_context, allocator, &value, &charge);
-    } else {
-      // TODO: can we work some magic with create_cb, which might be based on
-      // custom compression, to decompress without an extra copy in create_cb?
+    if (type != kNoCompression) {
+      // TODO: can we do something to avoid yet another allocation?
       Decompressor::Args args;
-      args.compressed_data = Slice(data_ptr, handle_value_charge);
-      args.compression_type = cache_options_.compression_type;
-      s = decompressor_->ExtractUncompressedSize(args);
-      assert(s.ok());
+      args.compressed_data = saved;
+      args.compression_type = type;
+      Status s = decompressor_->ExtractUncompressedSize(args);
+      assert(s.ok());  // in-memory data
       if (s.ok()) {
-        auto uncompressed = std::make_unique<char[]>(args.uncompressed_size);
+        uncompressed = std::make_unique<char[]>(args.uncompressed_size);
         s = decompressor_->DecompressBlock(args, uncompressed.get());
-        assert(s.ok());
-        if (s.ok()) {
-          s = helper->create_cb(
-              Slice(uncompressed.get(), args.uncompressed_size), kNoCompression,
-              CacheTier::kVolatileTier, create_context, allocator, &value,
-              &charge);
-        }
+        assert(s.ok());  // in-memory data
       }
+      if (!s.ok()) {
+        cache_->Release(lru_handle, /*erase_if_last_ref=*/true);
+        return nullptr;
+      }
+      saved = Slice(uncompressed.get(), args.uncompressed_size);
+      type = kNoCompression;
+      // Free temporary compressed data as early as we can. This could matter
+      // for unusually large blocks because we also have
+      // * Another compressed copy above (from lru_cache).
+      // * The uncompressed copy in `uncompressed`.
+      // * Another uncompressed copy in `result_value` below.
+      // Let's try to max out at 3 copies instead of 4.
+      merged_value = std::string();
     }
-  } else {
-    // The item was not compressed by us. Let the helper create_cb
-    // uncompress it
-    s = helper->create_cb(Slice(data_ptr, handle_value_charge), type, source,
-                          create_context, allocator, &value, &charge);
+    // Reduced as if it came from primary cache
+    source = CacheTier::kVolatileTier;
   }
 
+  Cache::ObjectPtr result_value = nullptr;
+  size_t result_charge = 0;
+  Status s = helper->create_cb(saved, type, source, create_context,
+                               cache_options_.memory_allocator.get(),
+                               &result_value, &result_charge);
   if (!s.ok()) {
     cache_->Release(lru_handle, /*erase_if_last_ref=*/true);
     return nullptr;
@@ -148,7 +152,8 @@ std::unique_ptr<SecondaryCacheResultHandle> CompressedSecondaryCache::Lookup(
     kept_in_sec_cache = true;
     cache_->Release(lru_handle, /*erase_if_last_ref=*/false);
   }
-  handle.reset(new CompressedSecondaryCacheResultHandle(value, charge));
+  handle.reset(
+      new CompressedSecondaryCacheResultHandle(result_value, result_charge));
   RecordTick(stats, COMPRESSED_SECONDARY_CACHE_HITS);
   return handle;
 }
@@ -171,85 +176,111 @@ bool CompressedSecondaryCache::MaybeInsertDummy(const Slice& key) {
 
 Status CompressedSecondaryCache::InsertInternal(
     const Slice& key, Cache::ObjectPtr value,
-    const Cache::CacheItemHelper* helper, CompressionType type,
+    const Cache::CacheItemHelper* helper, CompressionType from_type,
     CacheTier source) {
-  if (source != CacheTier::kVolatileCompressedTier &&
-      cache_options_.enable_custom_split_merge) {
-    // We don't support custom split/merge for the tiered case
-    return Status::OK();
-  }
-
-  auto internal_helper = GetHelper(cache_options_.enable_custom_split_merge);
-  char header[20];
-  char* payload = header;
-  payload = EncodeVarint32(payload, static_cast<uint32_t>(type));
-  payload = EncodeVarint32(payload, static_cast<uint32_t>(source));
-  size_t data_size = (*helper->size_cb)(value);
-  char* data_size_ptr = payload;
-  payload = EncodeVarint64(payload, data_size);
-
-  size_t header_size = payload - header;
-  size_t total_size = data_size + header_size;
-  CacheAllocationPtr ptr =
-      AllocateBlock(total_size, cache_options_.memory_allocator.get());
-  char* data_ptr = ptr.get() + header_size;
-
-  Status s = (*helper->saveto_cb)(value, 0, data_size, data_ptr);
+  bool enable_split_merge = cache_options_.enable_custom_split_merge;
+  const Cache::CacheItemHelper* internal_helper = GetHelper(enable_split_merge);
+
+  // TODO: variant of size_cb that also returns a pointer to the data if
+  // already available. Saves an allocation if we keep the compressed version.
+  const size_t data_size_original = (*helper->size_cb)(value);
+
+  // Allocate enough memory for header/tag + original data because (a) we might
+  // not be attempting compression at all, and (b) we might keep the original if
+  // compression is insufficient. But we don't need the length prefix with
+  // enable_split_merge. TODO: be smarter with CacheValueChunk to save an
+  // allocation in the enable_split_merge case.
+  size_t header_size = GetHeaderSize(data_size_original, enable_split_merge);
+  CacheAllocationPtr allocation = AllocateBlock(
+      header_size + data_size_original, cache_options_.memory_allocator.get());
+  char* data_ptr = allocation.get() + header_size;
+  Slice tagged_data(data_ptr - kTagSize, data_size_original + kTagSize);
+  assert(tagged_data.data() >= allocation.get());
+
+  Status s = (*helper->saveto_cb)(value, 0, data_size_original, data_ptr);
   if (!s.ok()) {
     return s;
   }
-  Slice val(data_ptr, data_size);
 
-  std::string compressed_val;
-  if (cache_options_.compression_type != kNoCompression &&
-      type == kNoCompression &&
+  std::unique_ptr<char[]> tagged_compressed_data;
+  CompressionType to_type = kNoCompression;
+  if (compressor_ && from_type == kNoCompression &&
       !cache_options_.do_not_compress_roles.Contains(helper->role)) {
-    PERF_COUNTER_ADD(compressed_sec_cache_uncompressed_bytes, data_size);
-
-    CompressionType to_type = kNoCompression;
-    s = compressor_->CompressBlock(val, &compressed_val, &to_type,
+    assert(source == CacheTier::kVolatileCompressedTier);
+
+    // TODO: consider malloc sizes for max acceptable compressed size
+    // Or maybe max_compressed_bytes_per_kb
+    size_t data_size_compressed = data_size_original - 1;
+    tagged_compressed_data =
+        std::make_unique<char[]>(data_size_compressed + kTagSize);
+    s = compressor_->CompressBlock(Slice(data_ptr, data_size_original),
+                                   tagged_compressed_data.get() + kTagSize,
+                                   &data_size_compressed, &to_type,
                                    nullptr /*working_area*/);
     if (!s.ok()) {
       return s;
     }
-    // TODO: allow values not compressed when there's no size savings?
-    assert(to_type == cache_options_.compression_type);
-    if (to_type != cache_options_.compression_type) {
-      return Status::Corruption("Failed to compress value.");
-    }
-
-    val = Slice(compressed_val);
-    data_size = compressed_val.size();
-    payload = EncodeVarint64(data_size_ptr, data_size);
-    header_size = payload - header;
-    total_size = header_size + data_size;
-    PERF_COUNTER_ADD(compressed_sec_cache_compressed_bytes, data_size);
-
-    if (!cache_options_.enable_custom_split_merge) {
-      ptr = AllocateBlock(total_size, cache_options_.memory_allocator.get());
-      data_ptr = ptr.get() + header_size;
-      memcpy(data_ptr, compressed_val.data(), data_size);
+    PERF_COUNTER_ADD(compressed_sec_cache_uncompressed_bytes,
+                     data_size_original);
+    if (to_type == kNoCompression) {
+      // Compression rejected or otherwise aborted/failed
+      to_type = kNoCompression;
+      tagged_compressed_data.reset();
+      // TODO: consider separate counters for rejected compressions
+      PERF_COUNTER_ADD(compressed_sec_cache_compressed_bytes,
+                       data_size_original);
+    } else {
+      PERF_COUNTER_ADD(compressed_sec_cache_compressed_bytes,
+                       data_size_compressed);
+      if (enable_split_merge) {
+        // Only need tagged_data for copying into CacheValueChunks.
+        tagged_data = Slice(tagged_compressed_data.get(),
+                            data_size_compressed + kTagSize);
+        allocation.reset();
+      } else {
+        // Replace allocation with compressed version, copied from string
+        header_size = GetHeaderSize(data_size_compressed, enable_split_merge);
+        allocation = AllocateBlock(header_size + data_size_compressed,
+                                   cache_options_.memory_allocator.get());
+        data_ptr = allocation.get() + header_size;
+        // Ignore unpopulated tag on tagged_compressed_data; will only be
+        // populated on the new allocation.
+        std::memcpy(data_ptr, tagged_compressed_data.get() + kTagSize,
+                    data_size_compressed);
+        tagged_data =
+            Slice(data_ptr - kTagSize, data_size_compressed + kTagSize);
+        assert(tagged_data.data() >= allocation.get());
+      }
     }
   }
 
   PERF_COUNTER_ADD(compressed_sec_cache_insert_real_count, 1);
-  if (cache_options_.enable_custom_split_merge) {
+
+  // Save the tag fields
+  const_cast<char*>(tagged_data.data())[0] = lossless_cast<char>(source);
+  const_cast<char*>(tagged_data.data())[1] = lossless_cast<char>(
+      source == CacheTier::kVolatileCompressedTier ? to_type : from_type);
+
+  if (enable_split_merge) {
     size_t split_charge{0};
-    CacheValueChunk* value_chunks_head = SplitValueIntoChunks(
-        val, cache_options_.compression_type, split_charge);
-    return cache_->Insert(key, value_chunks_head, internal_helper,
-                          split_charge);
+    CacheValueChunk* value_chunks_head =
+        SplitValueIntoChunks(tagged_data, split_charge);
+    s = cache_->Insert(key, value_chunks_head, internal_helper, split_charge);
+    assert(s.ok());  // LRUCache::Insert() with handle==nullptr always OK
   } else {
+    // Save the size prefix
+    char* ptr = allocation.get();
+    ptr = EncodeVarint64(ptr, tagged_data.size());
+    assert(ptr == tagged_data.data());
 #ifdef ROCKSDB_MALLOC_USABLE_SIZE
-    size_t charge = malloc_usable_size(ptr.get());
+    size_t charge = malloc_usable_size(allocation.get());
 #else
-    size_t charge = total_size;
+    size_t charge = tagged_data.size();
 #endif
-    std::memcpy(ptr.get(), header, header_size);
-    CacheAllocationPtr* buf = new CacheAllocationPtr(std::move(ptr));
-    charge += sizeof(CacheAllocationPtr);
-    return cache_->Insert(key, buf, internal_helper, charge);
+    s = cache_->Insert(key, allocation.release(), internal_helper, charge);
+    assert(s.ok());  // LRUCache::Insert() with handle==nullptr always OK
   }
+  return Status::OK();
 }
 
 Status CompressedSecondaryCache::Insert(const Slice& key,
@@ -271,7 +302,17 @@ Status CompressedSecondaryCache::Insert(const Slice& key,
 Status CompressedSecondaryCache::InsertSaved(
     const Slice& key, const Slice& saved, CompressionType type = kNoCompression,
     CacheTier source = CacheTier::kVolatileTier) {
+  if (source == CacheTier::kVolatileCompressedTier) {
+    // Unexpected, would violate InsertInternal preconditions
+    assert(source != CacheTier::kVolatileCompressedTier);
+    return Status::OK();
+  }
   if (type == kNoCompression) {
+    // Not currently supported (why?)
+    return Status::OK();
+  }
+  if (cache_options_.enable_custom_split_merge) {
+    // We don't support custom split/merge for the tiered case (why?)
     return Status::OK();
   }
 
@@ -291,7 +332,7 @@ Status CompressedSecondaryCache::SetCapacity(size_t capacity) {
   MutexLock l(&capacity_mutex_);
   cache_options_.capacity = capacity;
   cache_->SetCapacity(capacity);
-  disable_cache_ = capacity == 0;
+  disable_cache_.StoreRelaxed(capacity == 0);
   return Status::OK();
 }
 
@@ -321,9 +362,14 @@ std::string CompressedSecondaryCache::GetPrintableOptions() const {
   return ret;
 }
 
+// FIXME: this could use a lot of attention, including:
+// * Use allocator
+// * We shouldn't be worse than non-split; be more pro-actively aware of
+// internal fragmentation
+// * Consider a unified object/chunk structure that may or may not split
+// * Optimize size overhead of chunks
 CompressedSecondaryCache::CacheValueChunk*
 CompressedSecondaryCache::SplitValueIntoChunks(const Slice& value,
-                                               CompressionType compression_type,
                                                size_t& charge) {
   assert(!value.empty());
   const char* src_ptr = value.data();
@@ -344,15 +390,14 @@ CompressedSecondaryCache::SplitValueIntoChunks(const Slice& value,
     // size, or there is no compression.
     if (upper == malloc_bin_sizes_.begin() ||
         upper == malloc_bin_sizes_.end() ||
-        *upper - predicted_chunk_size < malloc_bin_sizes_.front() ||
-        compression_type == kNoCompression) {
+        *upper - predicted_chunk_size < malloc_bin_sizes_.front()) {
       tmp_size = predicted_chunk_size;
     } else {
       tmp_size = *(--upper);
     }
 
     CacheValueChunk* new_chunk =
-        reinterpret_cast<CacheValueChunk*>(new char[tmp_size]);
+        static_cast<CacheValueChunk*>(static_cast<void*>(new char[tmp_size]));
     current_chunk->next = new_chunk;
     current_chunk = current_chunk->next;
     actual_chunk_size = tmp_size - sizeof(CacheValueChunk) + 1;
@@ -367,28 +412,24 @@ CompressedSecondaryCache::SplitValueIntoChunks(const Slice& value,
   return dummy_head.next;
 }
 
-CacheAllocationPtr CompressedSecondaryCache::MergeChunksIntoValue(
-    const void* chunks_head, size_t& charge) {
-  const CacheValueChunk* head =
-      reinterpret_cast<const CacheValueChunk*>(chunks_head);
+std::string CompressedSecondaryCache::MergeChunksIntoValue(
+    const CacheValueChunk* head) {
   const CacheValueChunk* current_chunk = head;
-  charge = 0;
+  size_t total_size = 0;
   while (current_chunk != nullptr) {
-    charge += current_chunk->size;
+    total_size += current_chunk->size;
     current_chunk = current_chunk->next;
   }
 
-  CacheAllocationPtr ptr =
-      AllocateBlock(charge, cache_options_.memory_allocator.get());
+  std::string result;
+  result.reserve(total_size);
   current_chunk = head;
-  size_t pos{0};
   while (current_chunk != nullptr) {
-    memcpy(ptr.get() + pos, current_chunk->data, current_chunk->size);
-    pos += current_chunk->size;
+    result.append(current_chunk->data, current_chunk->size);
     current_chunk = current_chunk->next;
   }
-
-  return ptr;
+  assert(result.size() == total_size);
+  return result;
 }
 
 const Cache::CacheItemHelper* CompressedSecondaryCache::GetHelper(
@@ -402,16 +443,16 @@ const Cache::CacheItemHelper* CompressedSecondaryCache::GetHelper(
             CacheValueChunk* tmp_chunk = chunks_head;
             chunks_head = chunks_head->next;
             tmp_chunk->Free();
-            obj = nullptr;
           }
         }};
     return &kHelper;
   } else {
     static const Cache::CacheItemHelper kHelper{
         CacheEntryRole::kMisc,
-        [](Cache::ObjectPtr obj, MemoryAllocator* /*alloc*/) {
-          delete static_cast<CacheAllocationPtr*>(obj);
-          obj = nullptr;
+        [](Cache::ObjectPtr obj, MemoryAllocator* alloc) {
+          if (obj != nullptr) {
+            CacheAllocationDeleter{alloc}(static_cast<char*>(obj));
+          }
         }};
     return &kHelper;
   }
@@ -422,12 +463,7 @@ size_t CompressedSecondaryCache::TEST_GetCharge(const Slice& key) {
   if (lru_handle == nullptr) {
     return 0;
   }
-
   size_t charge = cache_->GetCharge(lru_handle);
-  if (cache_->Value(lru_handle) != nullptr &&
-      !cache_options_.enable_custom_split_merge) {
-    charge -= 10;
-  }
   cache_->Release(lru_handle, /*erase_if_last_ref=*/false);
   return charge;
 }

package/deps/rocksdb/rocksdb/cache/compressed_secondary_cache.h

@@ -10,13 +10,12 @@
 #include <memory>
 
 #include "cache/cache_reservation_manager.h"
-#include "cache/lru_cache.h"
 #include "memory/memory_allocator_impl.h"
+#include "rocksdb/advanced_compression.h"
 #include "rocksdb/secondary_cache.h"
 #include "rocksdb/slice.h"
 #include "rocksdb/status.h"
-#include "util/compression.h"
-#include "util/mutexlock.h"
+#include "util/atomic.h"
 
 namespace ROCKSDB_NAMESPACE {
 
@@ -124,14 +123,9 @@ class CompressedSecondaryCache : public SecondaryCache {
   // Split value into chunks to better fit into jemalloc bins. The chunks
   // are stored in CacheValueChunk and extra charge is needed for each chunk,
   // so the cache charge is recalculated here.
-  CacheValueChunk* SplitValueIntoChunks(const Slice& value,
-                                        CompressionType compression_type,
-                                        size_t& charge);
+  CacheValueChunk* SplitValueIntoChunks(const Slice& value, size_t& charge);
 
-  // After merging chunks, the extra charge for each chunk is removed, so
-  // the charge is recalculated.
-  CacheAllocationPtr MergeChunksIntoValue(const void* chunks_head,
-                                          size_t& charge);
+  std::string MergeChunksIntoValue(const CacheValueChunk* head);
 
   bool MaybeInsertDummy(const Slice& key);
 
@@ -149,7 +143,7 @@ class CompressedSecondaryCache : public SecondaryCache {
   std::shared_ptr<Decompressor> decompressor_;
   mutable port::Mutex capacity_mutex_;
   std::shared_ptr<ConcurrentCacheReservationManager> cache_res_mgr_;
-  bool disable_cache_;
+  RelaxedAtomic<bool> disable_cache_;
 };
 
 }  // namespace ROCKSDB_NAMESPACE