From 3c7ddfca346ca63d44f6a68ddc8ef7ff81a89c4e Mon Sep 17 00:00:00 2001
From: yzq <58433399+yangzq50@users.noreply.github.com>
Date: Thu, 2 Jan 2025 14:30:02 +0800
Subject: [PATCH] update TrySetOptimizing

---
 .../meta/entry/segment_index_entry.cpp        | 19 +++++++++++--------
 .../meta/entry/segment_index_entry.cppm       |  2 +-
 src/storage/meta/entry/table_entry.cpp        |  5 +++--
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/src/storage/meta/entry/segment_index_entry.cpp b/src/storage/meta/entry/segment_index_entry.cpp
index 1e5f2b48e9..cfa01e5281 100644
--- a/src/storage/meta/entry/segment_index_entry.cpp
+++ b/src/storage/meta/entry/segment_index_entry.cpp
@@ -865,7 +865,8 @@ void SegmentIndexEntry::ReplaceChunkIndexEntries(TxnTableStore *txn_table_store,
 ChunkIndexEntry *SegmentIndexEntry::RebuildChunkIndexEntries(TxnTableStore *txn_table_store, SegmentEntry *segment_entry) {
     const auto &index_name = *table_index_entry_->GetIndexName();
     Txn *txn = txn_table_store->GetTxn();
-    if (!TrySetOptimizing(txn)) {
+    const auto [result_b, add_segment_optimizing] = TrySetOptimizing(txn);
+    if (!result_b) {
         LOG_INFO(fmt::format("Index {} segment {} is optimizing, skip optimize.", index_name, segment_id_));
         return nullptr;
     }
@@ -899,6 +900,7 @@ ChunkIndexEntry *SegmentIndexEntry::RebuildChunkIndexEntries(TxnTableStore *txn_
             return nullptr;
         }
     }
+    add_segment_optimizing();
     RowID base_rowid(segment_id_, 0);
     SharedPtr<ChunkIndexEntry> merged_chunk_index_entry = nullptr;
     switch (index_base->index_type_) {
@@ -1225,17 +1227,18 @@ UniquePtr<SegmentIndexEntry> SegmentIndexEntry::Deserialize(const nlohmann::json
     return segment_index_entry;
 }
 
-bool SegmentIndexEntry::TrySetOptimizing(Txn *txn) {
+Pair<bool, std::function<void()>> SegmentIndexEntry::TrySetOptimizing(Txn *txn) {
     bool expected = false;
     bool success = optimizing_.compare_exchange_strong(expected, true);
     if (!success) {
-        return false;
+        return {false, nullptr};
     }
-    TableEntry *table_entry = table_index_entry_->table_index_meta()->GetTableEntry();
-    TxnTableStore *txn_table_store = txn->txn_store()->GetTxnTableStore(table_entry);
-    TxnIndexStore *txn_index_store = txn_table_store->GetIndexStore(table_index_entry_);
-    txn_index_store->AddSegmentOptimizing(this);
-    return true;
+    return {true, [this, txn] {
+        TableEntry *table_entry = table_index_entry_->table_index_meta()->GetTableEntry();
+        TxnTableStore *txn_table_store = txn->txn_store()->GetTxnTableStore(table_entry);
+        TxnIndexStore *txn_index_store = txn_table_store->GetIndexStore(table_index_entry_);
+        txn_index_store->AddSegmentOptimizing(this);
+    }};
 }
 
 void SegmentIndexEntry::ResetOptimizing() { optimizing_.store(false); }
diff --git a/src/storage/meta/entry/segment_index_entry.cppm b/src/storage/meta/entry/segment_index_entry.cppm
index 88242b9270..10bc288b81 100644
--- a/src/storage/meta/entry/segment_index_entry.cppm
+++ b/src/storage/meta/entry/segment_index_entry.cppm
@@ -277,7 +277,7 @@ private:
     Atomic<TxnTimeStamp> deprecate_ts_ = UNCOMMIT_TS;
 
 public:
-    bool TrySetOptimizing(Txn *txn);
+    Pair<bool, std::function<void()>> TrySetOptimizing(Txn *txn);
 
     void ResetOptimizing();
 
diff --git a/src/storage/meta/entry/table_entry.cpp b/src/storage/meta/entry/table_entry.cpp
index ddcdac7117..c8384f26c7 100644
--- a/src/storage/meta/entry/table_entry.cpp
+++ b/src/storage/meta/entry/table_entry.cpp
@@ -918,7 +918,8 @@ void TableEntry::OptimizeIndex(Txn *txn) {
                 const IndexFullText *index_fulltext = static_cast<const IndexFullText *>(index_base);
                 Map<SegmentID, SharedPtr<SegmentIndexEntry>> index_by_segment = table_index_entry->GetIndexBySegmentSnapshot(this, txn);
                 for (auto &[segment_id, segment_index_entry] : index_by_segment) {
-                    if (!segment_index_entry->TrySetOptimizing(txn)) {
+                    const auto [result_b, add_segment_optimizing] = segment_index_entry->TrySetOptimizing(txn);
+                    if (!result_b) {
                         LOG_INFO(fmt::format("Index {} segment {} is optimizing, skip optimize.", index_name, segment_id));
                         continue;
                     }
@@ -937,7 +938,7 @@ void TableEntry::OptimizeIndex(Txn *txn) {
                         opt_success = true;
                         continue;
                     }
-
+                    add_segment_optimizing();
                     String msg = fmt::format("merging {}", index_name);
                     Vector<String> base_names;
                     Vector<RowID> base_rowids;