Disagg: Fix the issue that lifetime of read-snapshots in WN may be lo…

…nger than expected (#9297) (#9317)

close #9298

1. In WN, release snapshots of SegmentReadTasks that do not have ColumnFileTiny and ColumnFileInMemory.
2. In CN, no matter all ColumnFileTinys and ColumnFileInMemorys are cached or not, call FetchDisaggPages to release snapshot in WN.

Signed-off-by: ti-chi-bot <[email protected]>

Co-authored-by: jinhelin <[email protected]>

dbms/src/Flash/Disaggregated/WNEstablishDisaggTaskHandler.cpp

@@ -103,6 +103,13 @@ void WNEstablishDisaggTaskHandler::execute(disaggregated::EstablishDisaggTaskRes
     snap->iterateTableSnapshots([&](const DM::Remote::DisaggPhysicalTableReadSnapshotPtr & snap) {
         response->add_tables(Serializer::serializeTo(snap, task_id, mem_tracker_wrapper).SerializeAsString());
     });
+
+    // Release SegmentReadTasks that do not have ColumnFileTiny and ColumnFileInMemory
+    // because these tasks will never call FetchDisaggPages.
+    auto to_release_tasks = snap->releaseNoNeedFetchTasks();
+    if (!to_release_tasks.empty())
+        LOG_INFO(log, "Release no need fetch tasks: count={} segments={}", to_release_tasks.size(), to_release_tasks);
+    snaps->unregisterSnapshotIfEmpty(task_id);
 }
 
 } // namespace DB

dbms/src/Flash/Disaggregated/WNFetchPagesStreamWriter.cpp

@@ -15,7 +15,6 @@
 #include <Common/Exception.h>
 #include <Common/MemoryTracker.h>
 #include <Common/Stopwatch.h>
-#include <Flash/Coprocessor/CHBlockChunkCodec.h>
 #include <Flash/Disaggregated/WNFetchPagesStreamWriter.h>
 #include <Flash/Mpp/TrackedMppDataPacket.h>
 #include <Interpreters/SharedContexts/Disagg.h>
@@ -27,8 +26,6 @@
 #include <Storages/DeltaMerge/Segment.h>
 #include <Storages/Page/PageUtil.h>
 #include <common/logger_useful.h>
-#include <kvproto/mpp.pb.h>
-#include <tipb/expression.pb.h>
 
 #include <memory>
 

dbms/src/Storages/DeltaMerge/Remote/DisaggSnapshot.cpp

@@ -72,6 +72,15 @@ bool DisaggReadSnapshot::empty() const
     return true;
 }
 
+SegmentReadTasks DisaggReadSnapshot::releaseNoNeedFetchTasks()
+{
+    SegmentReadTasks to_release_tasks;
+    std::unique_lock lock(mtx);
+    for (auto & [table_id, table_snap] : table_snapshots)
+        table_snap->releaseNoNeedFetchTasks(to_release_tasks);
+    return to_release_tasks;
+}
+
 DisaggPhysicalTableReadSnapshot::DisaggPhysicalTableReadSnapshot(
     KeyspaceTableID ks_table_id_,
     SegmentReadTasks && tasks_)
@@ -95,4 +104,20 @@ SegmentReadTaskPtr DisaggPhysicalTableReadSnapshot::popTask(const UInt64 segment
     return nullptr;
 }
 
+void DisaggPhysicalTableReadSnapshot::releaseNoNeedFetchTasks(SegmentReadTasks & to_release_tasks)
+{
+    std::unique_lock lock(mtx);
+    for (auto itr = tasks.begin(); itr != tasks.end();)
+    {
+        if (itr->second->hasColumnFileToFetch())
+        {
+            ++itr;
+        }
+        else
+        {
+            to_release_tasks.push_back(itr->second);
+            itr = tasks.erase(itr);
+        }
+    }
+}
 } // namespace DB::DM::Remote

dbms/src/Storages/DeltaMerge/Remote/DisaggSnapshot.h

@@ -81,6 +81,8 @@ class DisaggReadSnapshot
 
     bool empty() const;
 
+    SegmentReadTasks releaseNoNeedFetchTasks();
+
     DISALLOW_COPY(DisaggReadSnapshot);
 
 private:
@@ -104,6 +106,8 @@ class DisaggPhysicalTableReadSnapshot
         return tasks.empty();
     }
 
+    void releaseNoNeedFetchTasks(SegmentReadTasks & to_release_tasks);
+
     DISALLOW_COPY(DisaggPhysicalTableReadSnapshot);
 
 public:

dbms/src/Storages/DeltaMerge/Remote/RNWorkerFetchPages.cpp

@@ -94,6 +94,12 @@ disaggregated::FetchDisaggPagesRequest buildFetchPagesRequest(
 
 RNReadSegmentTaskPtr RNWorkerFetchPages::doWork(const RNReadSegmentTaskPtr & seg_task)
 {
+    if (seg_task->meta.delta_tinycf_page_ids.empty())
+    {
+        // No page need to be fetched or guarded.
+        return seg_task;
+    }
+
     MemoryTrackerSetter setter(true, fetch_pages_mem_tracker.get());
     Stopwatch watch_work{CLOCK_MONOTONIC_COARSE};
     SCOPE_EXIT({
@@ -170,10 +176,7 @@ void RNWorkerFetchPages::doFetchPages(
     const RNReadSegmentTaskPtr & seg_task,
     const disaggregated::FetchDisaggPagesRequest & request)
 {
-    // No page need to be fetched.
-    if (request.page_ids_size() == 0)
-        return;
-
+    // No matter all delta data is cached or not, call FetchDisaggPages to release snapshot in WN.
     Stopwatch sw_total;
     Stopwatch watch_rpc{CLOCK_MONOTONIC_COARSE};
     bool rpc_is_observed = false;
@@ -195,6 +198,10 @@ void RNWorkerFetchPages::doFetchPages(
         stream_resp->Finish();
     });
 
+    // All delta data is cached.
+    if (request.page_ids_size() == 0)
+        return;
+
     // Used to verify all pages are fetched.
     std::set<UInt64> remaining_pages_to_fetch;
     for (auto p : request.page_ids())

dbms/src/Storages/DeltaMerge/SegmentReadTaskPool.cpp

@@ -105,6 +105,17 @@ SegmentReadTasks SegmentReadTask::trySplitReadTasks(const SegmentReadTasks & tas
     return result_tasks;
 }
 
+bool SegmentReadTask::hasColumnFileToFetch() const
+{
+    auto need_to_fetch = [](const ColumnFilePtr & cf) {
+        // In v7.5, Only ColumnFileTiny need too fetch.
+        return cf->isTinyFile();
+    };
+    const auto & mem_cfs = read_snapshot->delta->getMemTableSetSnapshot()->getColumnFiles();
+    const auto & persisted_cfs = read_snapshot->delta->getPersistedFileSetSnapshot()->getColumnFiles();
+    return std::any_of(mem_cfs.cbegin(), mem_cfs.cend(), need_to_fetch)
+        || std::any_of(persisted_cfs.cbegin(), persisted_cfs.cend(), need_to_fetch);
+}
 
 SegmentReadTasksWrapper::SegmentReadTasksWrapper(bool enable_read_thread_, SegmentReadTasks && ordered_tasks_)
     : enable_read_thread(enable_read_thread_)

dbms/src/Storages/DeltaMerge/SegmentReadTaskPool.h

@@ -56,6 +56,9 @@ struct SegmentReadTask
     void mergeRanges() { ranges = DM::tryMergeRanges(std::move(ranges), 1); }
 
     static SegmentReadTasks trySplitReadTasks(const SegmentReadTasks & tasks, size_t expected_size);
+
+    // WN calls hasColumnFileToFetch to check whether a SegmentReadTask need to fetch column files from it
+    bool hasColumnFileToFetch() const;
 };
 
 class BlockStat