Fix the memory leak with pinned partitioned filters

Summary:
The existing unit test did not set the level so the check for pinned partitioned filter/index being properly released from the block cache was not properly exercised as they only take effect in level 0. As a result a memory leak in pinned partitioned filters was hidden. The patch fix the test as well as the bug.
Closes https://github.com/facebook/rocksdb/pull/3692

Differential Revision: D7559763

Pulled By: maysamyabandeh

fbshipit-source-id: 55eff274945838af983c764a7d71e8daff092e4a

diff --git a/HISTORY.md b/HISTORY.md
index 18b6a5ab2170cfae72fcaec83b83891b67827587..ec422414e0b00ca4adce9d01db0960600eb395e2 100644 (file)
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -2,6 +2,7 @@
 ## 5.12.2 (3/23/2018)
 ### Bug Fixes
 * Fsync after writing global seq number to the ingestion file in ExternalSstFileIngestionJob.
+* Fix memory leak when pin_l0_filter_and_index_blocks_in_cache is used with partitioned filters
 
 ### Java API Changes
 * Add `BlockBasedTableConfig.setBlockCache` to allow sharing a block cache across DB instances.

diff --git a/table/partitioned_filter_block.cc b/table/partitioned_filter_block.cc
index 8d6df35e14485862177494dbea1d4897b2b7fc78..a7939321a5555e2c1b34b0dbae92a96e8550962d 100644 (file)
--- a/table/partitioned_filter_block.cc
+++ b/table/partitioned_filter_block.cc
@@ -244,6 +244,13 @@ size_t PartitionedFilterBlockReader::ApproximateMemoryUsage() const {
   return idx_on_fltr_blk_->size();
 }
 
+// Release the cached entry and decrement its ref count.
+void ReleaseFilterCachedEntry(void* arg, void* h) {
+  Cache* cache = reinterpret_cast<Cache*>(arg);
+  Cache::Handle* handle = reinterpret_cast<Cache::Handle*>(h);
+  cache->Release(handle);
+}
+
 // TODO(myabandeh): merge this with the same function in IndexReader
 void PartitionedFilterBlockReader::CacheDependencies(bool pin) {
   // Before read partitions, prefetch them to avoid lots of IOs
@@ -301,6 +308,8 @@ void PartitionedFilterBlockReader::CacheDependencies(bool pin) {
     if (LIKELY(filter.IsSet())) {
       if (pin) {
         filter_map_[handle.offset()] = std::move(filter);
+        RegisterCleanup(&ReleaseFilterCachedEntry, block_cache,
+                        filter.cache_handle);
       } else {
         block_cache->Release(filter.cache_handle);
       }

diff --git a/table/partitioned_filter_block.h b/table/partitioned_filter_block.h
index 1a00a86e6ceaef334b2f139da4e0770374c44ad4..72d86b645a5da0c8db4a3d6a8a2a18adbb9c4d61 100644 (file)
--- a/table/partitioned_filter_block.h
+++ b/table/partitioned_filter_block.h
@@ -61,7 +61,8 @@ class PartitionedFilterBlockBuilder : public FullFilterBlockBuilder {
   uint32_t filters_in_partition_;
 };
 
-class PartitionedFilterBlockReader : public FilterBlockReader {
+class PartitionedFilterBlockReader : public FilterBlockReader,
+                                     public Cleanable {
  public:
   explicit PartitionedFilterBlockReader(const SliceTransform* prefix_extractor,
                                         bool whole_key_filtering,

diff --git a/table/table_test.cc b/table/table_test.cc
index 4bdf6ba1950f3fa79cdf196577146f1a5e9bfdf2..f41bbb303d8e998a1ca43bfe63e48cc2143bd5da 100644 (file)
--- a/table/table_test.cc
+++ b/table/table_test.cc
@@ -298,9 +298,11 @@ class KeyConvertingIterator : public InternalIterator {
 class TableConstructor: public Constructor {
  public:
   explicit TableConstructor(const Comparator* cmp,
-                            bool convert_to_internal_key = false)
+                            bool convert_to_internal_key = false,
+                            int level = -1)
       : Constructor(cmp),
-        convert_to_internal_key_(convert_to_internal_key) {}
+        convert_to_internal_key_(convert_to_internal_key),
+        level_(level) {}
   ~TableConstructor() { Reset(); }
 
   virtual Status FinishImpl(const Options& options,
@@ -315,14 +317,12 @@ class TableConstructor: public Constructor {
     std::vector<std::unique_ptr<IntTblPropCollectorFactory>>
         int_tbl_prop_collector_factories;
     std::string column_family_name;
-    int unknown_level = -1;
     builder.reset(ioptions.table_factory->NewTableBuilder(
-        TableBuilderOptions(ioptions, internal_comparator,
-                            &int_tbl_prop_collector_factories,
-                            options.compression, CompressionOptions(),
-                            nullptr /* compression_dict */,
-                            false /* skip_filters */, column_family_name,
-                            unknown_level),
+        TableBuilderOptions(
+            ioptions, internal_comparator, &int_tbl_prop_collector_factories,
+            options.compression, CompressionOptions(),
+            nullptr /* compression_dict */, false /* skip_filters */,
+            column_family_name, level_),
         TablePropertiesCollectorFactory::Context::kUnknownColumnFamily,
         file_writer_.get()));
 
@@ -347,8 +347,10 @@ class TableConstructor: public Constructor {
     uniq_id_ = cur_uniq_id_++;
     file_reader_.reset(test::GetRandomAccessFileReader(new test::StringSource(
         GetSink()->contents(), uniq_id_, ioptions.allow_mmap_reads)));
+    const bool skip_filters = false;
     return ioptions.table_factory->NewTableReader(
-        TableReaderOptions(ioptions, soptions, internal_comparator),
+        TableReaderOptions(ioptions, soptions, internal_comparator,
+                           skip_filters, level_),
         std::move(file_reader_), GetSink()->contents().size(), &table_reader_);
   }
 
@@ -408,6 +410,7 @@ class TableConstructor: public Constructor {
   unique_ptr<RandomAccessFileReader> file_reader_;
   unique_ptr<TableReader> table_reader_;
   bool convert_to_internal_key_;
+  int level_;
 
   TableConstructor();
 
@@ -2244,6 +2247,7 @@ std::map<std::string, size_t> MockCache::marked_data_in_cache_;
 // table is closed. This test makes sure that the only items remains in the
 // cache after the table is closed are raw data blocks.
 TEST_F(BlockBasedTableTest, NoObjectInCacheAfterTableClose) {
+  for (int level: {-1, 0, 1, 10}) {
   for (auto index_type :
        {BlockBasedTableOptions::IndexType::kBinarySearch,
         BlockBasedTableOptions::IndexType::kTwoLevelIndexSearch}) {
@@ -2280,7 +2284,9 @@ TEST_F(BlockBasedTableTest, NoObjectInCacheAfterTableClose) {
                 rocksdb::NewBloomFilterPolicy(10, block_based_filter));
             opt.table_factory.reset(NewBlockBasedTableFactory(table_options));
 
-            TableConstructor c(BytewiseComparator());
+            bool convert_to_internal_key = false;
+            TableConstructor c(BytewiseComparator(), convert_to_internal_key,
+                               level);
             std::string user_key = "k01";
             std::string key =
                 InternalKey(user_key, 0, kTypeValue).Encode().ToString();
@@ -2321,6 +2327,7 @@ TEST_F(BlockBasedTableTest, NoObjectInCacheAfterTableClose) {
       }
     }
   }
+  } // level
 }
 
 TEST_F(BlockBasedTableTest, BlockCacheLeak) {