Fix heap use after free bug

Fix close db slowly when force evict disabled

Author: mm304321141

db/db_impl.cc

@@ -287,7 +287,6 @@ DBImpl::DBImpl(const DBOptions& options, const std::string& dbname,
       initial_db_options_(SanitizeOptions(dbname, options)),
       immutable_db_options_(initial_db_options_),
       mutable_db_options_(initial_db_options_),
-      table_evict_type_(mutable_db_options_.table_evict_type),
       stats_(immutable_db_options_.statistics.get()),
       db_lock_(nullptr),
       mutex_(stats_, env_, DB_MUTEX_WAIT_MICROS,
@@ -457,6 +456,7 @@ DBImpl::DBImpl(const DBOptions& options, const std::string& dbname,
   versions_.reset(new VersionSet(dbname_, &immutable_db_options_, &env_options_,
                                  seq_per_batch, table_cache_.get(),
                                  write_buffer_manager_, &write_controller_));
+  versions_->set_table_evict_type(mutable_db_options_.table_evict_type);
   column_family_memtables_.reset(
       new ColumnFamilyMemTablesImpl(versions_->GetColumnFamilySet()));
 
@@ -1591,7 +1591,7 @@ Status DBImpl::SetDBOptions(
         new_options.bytes_per_sync = 1024 * 1024;
       }
       mutable_db_options_ = new_options;
-      table_evict_type_ = mutable_db_options_.table_evict_type;
+      versions_->set_table_evict_type(mutable_db_options_.table_evict_type);
       env_options_for_compaction_ = EnvOptions(
           BuildDBOptions(immutable_db_options_, mutable_db_options_));
       env_options_for_compaction_ = env_->OptimizeForCompactionTableWrite(

db/db_impl.h

@@ -827,7 +827,6 @@ class DBImpl : public DB {
   const DBOptions initial_db_options_;
   const ImmutableDBOptions immutable_db_options_;
   MutableDBOptions mutable_db_options_;
-  std::atomic<int> table_evict_type_;
   Statistics* stats_;
   std::unordered_map<std::string, RecoveredTransaction*>
       recovered_transactions_;

db/db_impl_files.cc

@@ -445,7 +445,7 @@ void DBImpl::PurgeObsoleteFiles(JobContext& state, bool schedule_only) {
     }
     ++i;
   }
-  int table_evict_type = table_evict_type_;
+  int table_evict_type = versions_->table_evict_type();
 
   for (const auto& candidate_file : candidate_files) {
     const std::string& to_delete = candidate_file.file_name;

db/dbformat.h

@@ -10,6 +10,7 @@
 #pragma once
 #include <stdio.h>
 
+#include <memory>
 #include <numeric>
 #include <string>
 #include <utility>
@@ -613,6 +614,8 @@ extern Status ReadRecordFromWriteBatch(Slice* input, char* tag,
                                        uint32_t* column_family, Slice* key,
                                        Slice* value, Slice* blob, Slice* xid);
 
+class LifeCycle : public std::enable_shared_from_this<LifeCycle> {};
+
 // When user call DeleteRange() to delete a range of keys,
 // we will store a serialized RangeTombstone in MemTable and SST.
 // the struct here is a easy-understood form

db/table_cache.cc

@@ -140,7 +140,7 @@ TableCache::TableCache(const ColumnFamilyOptions& initial_cf_options,
                        const ImmutableDBOptions& db_options,
                        const EnvOptions* env_options, Cache* const cache)
     : initial_cf_options_(initial_cf_options),
-      ioptions_(db_options, initial_cf_options_),
+      ioptions_(db_options, initial_cf_options_, shared_from_this()),
       env_options_(*env_options),
       cache_(cache),
       immortal_tables_(false) {}

db/table_cache.h

@@ -35,7 +35,7 @@ struct FileDescriptor;
 class GetContext;
 class HistogramImpl;
 
-class TableCache {
+class TableCache : public LifeCycle {
  public:
   TableCache(const ColumnFamilyOptions& initial_cf_options,
              const ImmutableDBOptions& db_options,

db/version_set.cc

@@ -2996,12 +2996,15 @@ VersionSet::~VersionSet() {
   // VersionSet
   Cache* table_cache = column_family_set_->get_table_cache();
   column_family_set_.reset();
+  int table_evict_type = table_evict_type();
   for (auto& file : obsolete_files_) {
     if (file.metadata->table_reader_handle) {
       table_cache->Release(file.metadata->table_reader_handle);
     }
-    if (file.table_cache) {
-      file.table_cache->ForceEvict(file.metadata->fd.GetNumber(), nullptr);
+    if (file.table_cache && table_evict_type != kSkipForceEvict) {
+      file.table_cache->ForceEvict(
+          file.metadata->fd.GetNumber(),
+          table_evict_type == kAlwaysForceEvict ? file.metadata : nullptr);
     } else {
       TableCache::Evict(table_cache, file.metadata->fd.GetNumber());
     }

db/version_set.h

@@ -1121,6 +1121,9 @@ class VersionSet {
         new_options.writable_file_max_buffer_size;
   }
 
+  int table_evict_type() const { return table_evict_type_; }
+  void set_table_evict_type(int type) { table_evict_type_ = type; }
+
   const ImmutableDBOptions* db_options() const { return db_options_; }
 
   static uint64_t GetNumLiveVersions(Version* dummy_versions);
@@ -1223,6 +1226,7 @@ class VersionSet {
   std::vector<std::string> obsolete_manifests_;
 
   const bool seq_per_batch_;
+  std::atomic<int> table_evict_type_{0};
 
   // env options for all reads and writes except compactions
   EnvOptions env_options_;

options/cf_options.cc

@@ -95,6 +95,13 @@ ImmutableCFOptions::ImmutableCFOptions(const ImmutableDBOptions& db_options,
   }
 }
 
+ImmutableCFOptions::ImmutableCFOptions(const ImmutableDBOptions& db_options,
+                                       const ColumnFamilyOptions& cf_options,
+                                       std::shared_ptr<LifeCycle> _life_cycle)
+    : ImmutableCFOptions(db_options, cf_options) {
+  life_cycle = _life_cycle;
+}
+
 // Multiple two operands. If they overflow, return op1.
 uint64_t MultiplyCheckOverflow(uint64_t op1, double op2) {
   if (op1 == 0 || op2 <= 0) {

options/cf_options.h

@@ -28,6 +28,10 @@ struct ImmutableCFOptions {
   ImmutableCFOptions(const ImmutableDBOptions& db_options,
                      const ColumnFamilyOptions& cf_options);
 
+  ImmutableCFOptions(const ImmutableDBOptions& db_options,
+                     const ColumnFamilyOptions& cf_options,
+                     std::shared_ptr<LifeCycle> life_cycle);
+
   CompactionStyle compaction_style;
 
   CompactionPri compaction_pri;
@@ -135,6 +139,8 @@ struct ImmutableCFOptions {
 
   std::shared_ptr<std::vector<std::unique_ptr<IntTblPropCollectorFactory>>>
       int_tbl_prop_collector_factories_for_blob;
+
+  std::weak_ptr<LifeCycle> life_cycle;
 };
 
 struct BlobConfig {

table/block_based_table_reader.cc

@@ -3205,9 +3205,13 @@ void BlockBasedTable::Close() {
                                 rep_->dummy_index_reader_offset, cache_key);
     erased_count += rep_->table_options.block_cache.get()->Erase(key);
 
-    RecordTick(rep_->ioptions.statistics, BLOCK_CACHE_ERASE, 2);
-    RecordTick(rep_->ioptions.statistics, BLOCK_CACHE_ERASE_FAILURES,
-               2 - erased_count);
+    auto life_cycle = rep_->cf_life_cycle.lock();
+
+    if (life_cycle) {
+      RecordTick(rep_->ioptions.statistics, BLOCK_CACHE_ERASE, 2);
+      RecordTick(rep_->ioptions.statistics, BLOCK_CACHE_ERASE_FAILURES,
+                 2 - erased_count);
+    }
   }
   rep_->closed = true;
 }

table/block_based_table_reader.h

@@ -445,6 +445,7 @@ struct BlockBasedTable::Rep {
       const InternalKeyComparator& _internal_comparator, bool skip_filters,
       int _level, const bool _immortal_table)
       : ioptions(_ioptions),
+        cf_life_cycle(_ioptions.life_cycle),
         env_options(_env_options),
         table_options(_table_opt),
         filter_policy(skip_filters ? nullptr : _table_opt.filter_policy.get()),
@@ -460,6 +461,7 @@ struct BlockBasedTable::Rep {
         immortal_table(_immortal_table) {}
 
   const ImmutableCFOptions& ioptions;
+  std::weak_ptr<LifeCycle> cf_life_cycle;
   const EnvOptions& env_options;
   const BlockBasedTableOptions table_options;
   const FilterPolicy* const filter_policy;