[fix](load) change tablet schema pointer to shared_ptr in memtable (#…

…37927) ## Proposed changes Fix potential use-after-free in memtable->_tablet_schema.
apache · Jul 17, 2024 · 7a7839d · 7a7839d
1 parent aa968d6
commit 7a7839d
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 8 deletions.
diff --git a/be/src/olap/base_tablet.cpp b/be/src/olap/base_tablet.cpp
@@ -830,7 +830,7 @@ Status BaseTablet::sort_block(vectorized::Block& in_block, vectorized::Block& ou
  vectorized::MutableBlock::build_mutable_block(&output_block);
 
  std::shared_ptr<RowInBlockComparator> vec_row_comparator =
- std::make_shared<RowInBlockComparator>(_tablet_meta->tablet_schema().get());
+ std::make_shared<RowInBlockComparator>(_tablet_meta->tablet_schema());
  vec_row_comparator->set_block(&mutable_input_block);
 
  std::vector<std::unique_ptr<RowInBlock>> row_in_blocks;

diff --git a/be/src/olap/memtable.cpp b/be/src/olap/memtable.cpp
@@ -48,7 +48,7 @@ bvar::Adder<int64_t> g_memtable_input_block_allocated_size("memtable_input_block
 
 using namespace ErrorCode;
 
-MemTable::MemTable(int64_t tablet_id, const TabletSchema* tablet_schema,
+MemTable::MemTable(int64_t tablet_id, std::shared_ptr<TabletSchema> tablet_schema,
  const std::vector<SlotDescriptor*>* slot_descs, TupleDescriptor* tuple_desc,
  bool enable_unique_key_mow, PartialUpdateInfo* partial_update_info,
  const std::shared_ptr<MemTracker>& insert_mem_tracker,

diff --git a/be/src/olap/memtable.h b/be/src/olap/memtable.h
@@ -129,15 +129,16 @@ class Tie {
 
 class RowInBlockComparator {
 public:
- RowInBlockComparator(const TabletSchema* tablet_schema) : _tablet_schema(tablet_schema) {}
+ RowInBlockComparator(std::shared_ptr<TabletSchema> tablet_schema)
+ : _tablet_schema(tablet_schema) {}
  // call set_block before operator().
  // only first time insert block to create _input_mutable_block,
  // so can not Comparator of construct to set pblock
  void set_block(vectorized::MutableBlock* pblock) { _pblock = pblock; }
  int operator()(const RowInBlock* left, const RowInBlock* right) const;
 
 private:
- const TabletSchema* _tablet_schema = nullptr;
+ std::shared_ptr<TabletSchema> _tablet_schema;
  vectorized::MutableBlock* _pblock = nullptr; // corresponds to Memtable::_input_mutable_block
 };
 
@@ -168,7 +169,7 @@ class MemTableStat {
 
 class MemTable {
 public:
- MemTable(int64_t tablet_id, const TabletSchema* tablet_schema,
+ MemTable(int64_t tablet_id, std::shared_ptr<TabletSchema> tablet_schema,
  const std::vector<SlotDescriptor*>* slot_descs, TupleDescriptor* tuple_desc,
  bool enable_unique_key_mow, PartialUpdateInfo* partial_update_info,
  const std::shared_ptr<MemTracker>& insert_mem_tracker,
@@ -209,7 +210,7 @@ class MemTable {
  bool _enable_unique_key_mow = false;
  bool _is_partial_update = false;
  const KeysType _keys_type;
- const TabletSchema* _tablet_schema = nullptr;
+ std::shared_ptr<TabletSchema> _tablet_schema;
 
  std::shared_ptr<RowInBlockComparator> _vec_row_comparator;
 

diff --git a/be/src/olap/memtable_writer.cpp b/be/src/olap/memtable_writer.cpp
@@ -213,8 +213,8 @@ void MemTableWriter::_reset_mem_table() {
  }
  {
  std::lock_guard<SpinLock> l(_mem_table_ptr_lock);
- _mem_table.reset(new MemTable(_req.tablet_id, _tablet_schema.get(), _req.slots,
- _req.tuple_desc, _unique_key_mow, _partial_update_info.get(),
+ _mem_table.reset(new MemTable(_req.tablet_id, _tablet_schema, _req.slots, _req.tuple_desc,
+ _unique_key_mow, _partial_update_info.get(),
  mem_table_insert_tracker, mem_table_flush_tracker));
  }