fix(storage): make block map access thread-safe and look up block metas by block id

src/executor/operator/physical_scan/physical_match_sparse_scan_impl.cpp

@@ -434,6 +434,9 @@ void PhysicalMatchSparseScan::ExecuteInnerT(DistFunc *dist_func,
         Bitmask bitmask;
 
         BlockMeta *block_meta = block_index->GetBlockMeta(segment_id, block_id);
+        if (block_meta == nullptr) [[unlikely]] {
+            UnrecoverableError(fmt::format("MatchSparseScan: block not found. segment_id: {}, block_id: {}", segment_id, block_id));
+        }
         LOG_DEBUG(fmt::format("MatchSparseScan: segment_id: {}, block_id: {}", segment_id, block_id));
 
         std::tie(row_cnt, status) = block_meta->GetRowCnt1();

src/executor/operator/physical_scan/physical_match_tensor_scan_impl.cpp

@@ -1270,7 +1270,10 @@ void GetRerankerScore(std::vector<MatchTensorRerankDoc> &rerank_docs,
         const BlockID block_id = segment_offset / DEFAULT_BLOCK_CAPACITY;
         const BlockOffset block_offset = segment_offset % DEFAULT_BLOCK_CAPACITY;
         ColumnVector column_vec;
-        BlockMeta *block_meta = block_index->new_segment_block_index_.at(segment_id).block_map().at(block_id).get();
+        BlockMeta *block_meta = block_index->GetBlockMeta(segment_id, block_id);
+        if (block_meta == nullptr) [[unlikely]] {
+            UnrecoverableError(fmt::format("MatchTensor rerank: block not found. segment_id: {}, block_id: {}", segment_id, block_id));
+        }
         ColumnMeta column_meta(column_id, *block_meta);
         auto [block_row_cnt, status] = block_meta->GetRowCnt1();
         if (!status.ok()) {

src/executor/operator/physical_scan/physical_table_scan_impl.cpp

@@ -135,6 +135,9 @@ void PhysicalTableScan::ExecuteInternal(QueryContext *query_context, TableScanOp
         BlockMeta *current_block_meta = block_index->GetBlockMeta(segment_id, block_id);
 
         std::optional<NewTxnGetVisibleRangeState> &range_state = table_scan_function_data_ptr->get_visible_range_state_;
+        if (current_block_meta == nullptr) [[unlikely]] {
+            UnrecoverableError(fmt::format("TableScan: block not found. segment_id: {}, block_id: {}", segment_id, block_id));
+        }
         if (!range_state) {
             // TODO
             range_state = NewTxnGetVisibleRangeState();
@@ -153,6 +156,9 @@ void PhysicalTableScan::ExecuteInternal(QueryContext *query_context, TableScanOp
                                           block_ids_idx,
                                           block_ids_count));
                     ++block_ids_idx;
+                    // Drop the range state of the skipped block, otherwise the next block would be
+                    // read with this block's visible ranges (the state is only rebuilt when empty).
+                    range_state.reset();
                     continue;
                 } else {
                     LOG_TRACE(fmt::format("TableScan: block_ids_idx: {}, block_ids.size(): {}, not skipped after apply FastRoughFilter",

src/storage/catalog/meta/segment_meta.cppm

@@ -84,7 +84,6 @@ public:
     Status CommitBlock(BlockID block_id, TxnTimeStamp commit_ts);
 
     std::tuple<std::vector<BlockID> *, Status> GetBlockIDs1();
-    std::tuple<std::vector<BlockID> *, Status> GetBlockIDs1(TxnTimeStamp commit_ts);
 
     std::tuple<size_t, Status> GetRowCnt1();
 

src/storage/catalog/meta/segment_meta_impl.cpp

@@ -189,22 +189,27 @@ Status SegmentMeta::AddBlockWithID(TxnTimeStamp commit_ts, BlockID block_id) {
 }
 
 std::pair<BlockID, Status> SegmentMeta::AddBlockID1(TxnTimeStamp commit_ts) {
-    Status status;
-
     BlockID block_id = 0;
     {
-        std::vector<BlockID> *block_ids_ptr = nullptr;
-        std::tie(block_ids_ptr, status) = GetBlockIDs1(commit_ts);
-        if (!status.ok()) {
-            return {INVALID_BLOCK_ID, status};
+        // Hold the lock across the whole read-modify-write: otherwise two appends of the same
+        // segment can both read the same last id and hand out the same one, and push_back can
+        // grow the vector while another thread is iterating it.
+        std::lock_guard<std::mutex> lock(mtx_);
+        if (!block_ids1_) {
+            block_ids1_ = infinity::GetTableSegmentBlocks(&kv_instance_,
+                                                          table_meta_.db_id_str(),
+                                                          table_meta_.table_id_str(),
+                                                          segment_id_,
+                                                          begin_ts_,
+                                                          commit_ts);
         }
-        block_id = block_ids_ptr->empty() ? 0 : block_ids_ptr->back() + 1;
+        block_id = block_ids1_->empty() ? 0 : block_ids1_->back() + 1;
         block_ids1_->push_back(block_id);
     }
 
     std::string block_id_key = KeyEncode::CatalogTableSegmentBlockKey(table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_, block_id);
     std::string commit_ts_str = fmt::format("{}", commit_ts);
-    status = kv_instance_.Put(block_id_key, commit_ts_str);
+    Status status = kv_instance_.Put(block_id_key, commit_ts_str);
     if (!status.ok()) {
         return {0, status};
     }
@@ -230,14 +235,6 @@ std::tuple<std::vector<BlockID> *, Status> SegmentMeta::GetBlockIDs1() {
     return {&*block_ids1_, Status::OK()};
 }
 
-std::tuple<std::vector<BlockID> *, Status> SegmentMeta::GetBlockIDs1(TxnTimeStamp commit_ts) {
-    if (!block_ids1_) {
-        block_ids1_ =
-            infinity::GetTableSegmentBlocks(&kv_instance_, table_meta_.db_id_str(), table_meta_.table_id_str(), segment_id_, begin_ts_, commit_ts);
-    }
-    return {&*block_ids1_, Status::OK()};
-}
-
 std::tuple<size_t, Status> SegmentMeta::GetRowCnt1() {
     std::lock_guard<std::mutex> lock(mtx_);
     if (row_cnt_) {

src/storage/common/block_index.cppm

@@ -35,7 +35,8 @@ public:
     std::unique_ptr<SegmentMeta> segment_meta_;
 
 private:
-    mutable std::vector<std::unique_ptr<BlockMeta>> block_map_;
+    mutable std::mutex block_map_mtx_;
+    mutable std::vector<std::unique_ptr<BlockMeta>> block_map_; // only used in function block_map()
 };
 
 export struct BlockIndex {

src/storage/common/block_index_impl.cpp

@@ -41,6 +41,7 @@ SegmentOffset NewSegmentSnapshot::segment_offset() const {
 }
 
 const std::vector<std::unique_ptr<BlockMeta>> &NewSegmentSnapshot::block_map() const {
+    std::lock_guard<std::mutex> lock(block_map_mtx_);
     if (block_map_.size()) {
         return block_map_;
     }
@@ -68,7 +69,7 @@ void BlockIndex::NewInit(const std::shared_ptr<TableMeta> &table_meta) {
         RecoverableError(status);
     }
     for (SegmentID segment_id : *segment_ids_ptr) {
-        NewSegmentSnapshot &segment_snapshot = new_segment_block_index_.emplace(segment_id, NewSegmentSnapshot()).first->second;
+        NewSegmentSnapshot &segment_snapshot = new_segment_block_index_[segment_id];
         segment_snapshot.segment_meta_ = std::make_unique<SegmentMeta>(segment_id, *table_meta_);
     }
 }
@@ -105,11 +106,16 @@ BlockMeta *BlockIndex::GetBlockMeta(u32 segment_id, u16 block_id) const {
     if (seg_it == new_segment_block_index_.end()) {
         return nullptr;
     }
-    const auto &blocks_info = seg_it->second;
-    if (blocks_info.block_map().size() <= block_id) {
+    // block_map() is sorted by id (GetTableSegmentBlocks sorts the ids) but holds only the blocks
+    // visible to this transaction, so block_id is a key, not a position in the vector.
+    const std::vector<std::unique_ptr<BlockMeta>> &block_map = seg_it->second.block_map();
+    auto it = std::lower_bound(block_map.begin(), block_map.end(), block_id, [](const std::unique_ptr<BlockMeta> &block_meta, BlockID id) {
+        return block_meta->block_id() < id;
+    });
+    if (it == block_map.end() || (*it)->block_id() != block_id) {
         return nullptr;
     }
-    return blocks_info.block_map()[block_id].get();
+    return it->get();
 }
 void IndexIndex::Insert(std::string index_name, std::shared_ptr<NewIndexSnapshot> new_index_snapshot) {
     new_index_snapshots_vec_.push_back(new_index_snapshot.get());