Better PlainTableReader

db/table_cache.cc

@@ -219,6 +219,7 @@ Status TableCache::FindTable(const EnvOptions& env_options,
       s = cache_->Insert(key, table_reader.get(), 1, &DeleteEntry<TableReader>,
                          handle);
       if (s.ok()) {
+        table_reader->SetTableCacheHandle(cache_, *handle);
         // Release ownership of table reader.
         table_reader.release();
       }

include/rocksdb/lazy_buffer.h

@@ -317,6 +317,11 @@ inline LazyBufferContext* LazyBufferState::get_context(LazyBuffer* buffer) {
   return &buffer->context_;
 }
 
+#ifdef __GNUC__
+# pragma GCC diagnostic push
+# pragma GCC diagnostic ignored "-Wmissing-field-initializers"
+#endif
+
 inline LazyBuffer::LazyBuffer() noexcept
     : state_(LazyBufferState::light_state()),
       context_{},
@@ -370,6 +375,10 @@ inline LazyBuffer::LazyBuffer(const LazyBufferState* _state,
   assert(_state != nullptr);
 }
 
+#ifdef __GNUC__
+# pragma GCC diagnostic pop
+#endif
+
 inline void LazyBuffer::destroy() {
   if (state_ != nullptr) {
     state_->destroy(this);

table/plain_table_reader.cc

@@ -91,14 +91,24 @@ class PlainTableIterator : public InternalIterator {
 };
 
 extern const uint64_t kPlainTableMagicNumber;
-PlainTableReader::PlainTableReader(const ImmutableCFOptions& ioptions,
-                                   std::unique_ptr<RandomAccessFileReader>&& file,
-                                   const EnvOptions& storage_options,
-                                   const InternalKeyComparator& icomparator,
-                                   EncodingType encoding_type,
-                                   uint64_t file_number, uint64_t file_size,
-                                   const TableProperties* table_properties,
-                                   const SliceTransform* prefix_extractor)
+
+void PlainTableReader::SetTableCacheHandle(Cache* table_cache,
+                                           Cache::Handle *handle) {
+  MutexLock lock(&table_cache_mutex_);
+  assert(table_cache != nullptr);
+  assert(handle != nullptr);
+  table_cache_ = table_cache;
+  table_cache_handle_ = handle;
+}
+
+PlainTableReader::PlainTableReader(
+    const ImmutableCFOptions& ioptions,
+    std::unique_ptr<RandomAccessFileReader>&& file,
+    const EnvOptions& storage_options,
+    const InternalKeyComparator& icomparator, EncodingType encoding_type,
+    uint64_t file_number, uint64_t file_size,
+    const TableProperties* table_properties,
+    const SliceTransform* prefix_extractor)
     : internal_comparator_(icomparator),
       encoding_type_(encoding_type),
       full_scan_mode_(false),
@@ -111,7 +121,9 @@ PlainTableReader::PlainTableReader(const ImmutableCFOptions& ioptions,
       ioptions_(ioptions),
       file_number_(file_number),
       file_size_(file_size),
-      table_properties_(nullptr) {}
+      table_properties_(nullptr),
+      table_cache_(nullptr),
+      table_cache_handle_(nullptr) {}
 
 PlainTableReader::~PlainTableReader() {
 }
@@ -601,10 +613,26 @@ Status PlainTableReader::Get(const ReadOptions& /*ro*/, const Slice& target,
     // can we enable the fast path?
     if (internal_comparator_.Compare(found_key, parsed_target) >= 0) {
       bool dont_care __attribute__((__unused__));
-      if (!get_context->SaveValue(found_key,
-                                  LazyBuffer(found_value, false, file_number_),
-                                  &dont_care)) {
-        break;
+      if (found_value.size() <= sizeof(LazyBufferContext) ||
+          table_cache_handle_ == nullptr ||
+          !table_cache_->Ref(table_cache_handle_)) {
+        if (!get_context->SaveValue(found_key,
+                                    LazyBuffer(found_value, false,
+                                               file_number_),
+                                    &dont_care)) {
+          break;
+        }
+      } else {
+        if (!get_context->SaveValue(
+            found_key, LazyBuffer(
+                found_value, Cleanable([](void* arg1, void* arg2) {
+                  auto c = static_cast<Cache*>(arg1);
+                  auto h = static_cast<Cache::Handle*>(arg2);
+                  c->Release(h);
+                }, table_cache_, table_cache_handle_), file_number_),
+            &dont_care)) {
+          break;
+        }
       }
     }
   }

table/plain_table_reader.h

@@ -108,6 +108,8 @@ class PlainTableReader: public TableReader {
     return file_number_;
   }
 
+  void SetTableCacheHandle(Cache* table_cache, Cache::Handle *handle) override;
+
   PlainTableReader(const ImmutableCFOptions& ioptions,
                    std::unique_ptr<RandomAccessFileReader>&& file,
                    const EnvOptions& env_options,
@@ -165,6 +167,9 @@ class PlainTableReader: public TableReader {
   uint64_t file_number_;
   uint64_t file_size_;
   std::shared_ptr<const TableProperties> table_properties_;
+  Cache* table_cache_;
+  Cache::Handle *table_cache_handle_;
+  port::Mutex table_cache_mutex_;
 
   bool IsFixedLength() const {
     return user_key_len_ != kPlainTableVariableLength;

table/table_reader.h

@@ -10,6 +10,7 @@
 #pragma once
 #include <memory>
 #include "db/range_tombstone_fragmenter.h"
+#include "rocksdb/cache.h"
 #include "rocksdb/slice_transform.h"
 #include "table/internal_iterator.h"
 
@@ -119,6 +120,9 @@ class TableReader {
     return Status::NotSupported("VerifyChecksum() not supported");
   }
 
+  virtual void SetTableCacheHandle(Cache* /*table_cache*/,
+                                   Cache::Handle */*handle*/) {}
+
   void UpdateMaxCoveringTombstoneSeq(
       const ReadOptions& readOptions, const Slice& user_key,
       SequenceNumber* max_covering_tombstone_seq);

util/lazy_buffer.cc

@@ -402,6 +402,21 @@ const LazyBufferState* LazyBufferState::cleanable_state() {
   return &static_state;
 }
 
+void LazyBuffer::fix_light_state(const LazyBuffer& other) {
+  assert(state_ == LazyBufferState::light_state());
+  assert(other.size_ <= sizeof(LazyBufferContext));
+  data_ = union_cast<LightLazyBufferState::Context>(&context_)->data;
+  size_ = other.size_;
+  if (!other.empty()) {
+    ::memmove(data_, other.data_, size_);
+  }
+}
+
+#ifdef __GNUC__
+# pragma GCC diagnostic push
+# pragma GCC diagnostic ignored "-Wmissing-field-initializers"
+#endif
+
 LazyBuffer::LazyBuffer(size_t _size) noexcept
     : context_{},
       file_number_(uint64_t(-1)) {
@@ -424,16 +439,6 @@ LazyBuffer::LazyBuffer(size_t _size) noexcept
   }
 }
 
-void LazyBuffer::fix_light_state(const LazyBuffer& other) {
-  assert(state_ == LazyBufferState::light_state());
-  assert(other.size_ <= sizeof(LazyBufferContext));
-  data_ = union_cast<LightLazyBufferState::Context>(&context_)->data;
-  size_ = other.size_;
-  if (!other.empty()) {
-    ::memmove(data_, other.data_, size_);
-  }
-}
-
 LazyBuffer::LazyBuffer(LazyBufferCustomizeBuffer _buffer) noexcept
     : state_(LazyBufferState::buffer_state()),
       context_{reinterpret_cast<uint64_t>(_buffer.handle),
@@ -453,6 +458,10 @@ LazyBuffer::LazyBuffer(std::string* _string) noexcept
   ::new(&union_cast<StringLazyBufferState::Context>(&context_)->status) Status;
 }
 
+#ifdef __GNUC__
+# pragma GCC diagnostic pop
+#endif
+
 void LazyBuffer::reset(LazyBufferCustomizeBuffer _buffer) {
   assert(_buffer.handle != nullptr);
   assert(_buffer.uninitialized_resize != nullptr);