Storages: Shutdown the LocalIndexScheduler before shutting down PageStorage/DeltaMergeStore

dbms/src/Interpreters/Context.cpp

@@ -271,6 +271,13 @@ struct ContextShared
             return;
         shutdown_called = true;
 
+        // The local index scheduler must be shutdown to stop all
+        // running tasks before shutting down `global_storage_pool`.
+        if (global_local_indexer_scheduler)
+        {
+            global_local_indexer_scheduler->shutdown();
+        }
+
         if (global_storage_pool)
         {
             // shutdown the gc task of global storage pool before

dbms/src/Interpreters/executeQuery.cpp

@@ -311,7 +311,7 @@ std::tuple<ASTPtr, BlockIO> executeQueryImpl(
 
                 if (elem.read_rows != 0)
                 {
-                    LOG_INFO(
+                    LOG_DEBUG(
                         execute_query_logger,
                         "Read {} rows, {} in {:.3f} sec., {} rows/sec., {}/sec.",
                         elem.read_rows,
@@ -421,7 +421,7 @@ void logQueryPipeline(const LoggerPtr & logger, const BlockInputStreamPtr & in)
         in->dumpTree(log_buffer);
         return log_buffer.toString();
     };
-    LOG_INFO(logger, pipeline_log_str());
+    LOG_DEBUG(logger, pipeline_log_str());
 }
 
 BlockIO executeQuery(const String & query, Context & context, bool internal, QueryProcessingStage::Enum stage)

dbms/src/Storages/DeltaMerge/LocalIndexerScheduler.cpp

@@ -64,9 +64,9 @@ LocalIndexerScheduler::LocalIndexerScheduler(const Options & options)
         start();
 }
 
-LocalIndexerScheduler::~LocalIndexerScheduler()
+void LocalIndexerScheduler::shutdown()
 {
-    LOG_INFO(logger, "LocalIndexerScheduler is destroying. Waiting scheduler and tasks to finish...");
+    LOG_INFO(logger, "LocalIndexerScheduler is shutting down. Waiting scheduler and tasks to finish...");
 
     // First quit the scheduler. Don't schedule more tasks.
     is_shutting_down = true;
@@ -81,7 +81,15 @@ LocalIndexerScheduler::~LocalIndexerScheduler()
 
     // Then wait all running tasks to finish.
     pool.reset();
+    LOG_INFO(logger, "LocalIndexerScheduler is shutdown.");
+}
 
+LocalIndexerScheduler::~LocalIndexerScheduler()
+{
+    if (!is_shutting_down)
+    {
+        shutdown();
+    }
     LOG_INFO(logger, "LocalIndexerScheduler is destroyed");
 }
 
@@ -295,7 +303,10 @@ bool LocalIndexerScheduler::tryAddTaskToPool(std::unique_lock<std::mutex> & lock
         }
     };
 
-    RUNTIME_CHECK(pool);
+    if (is_shutting_down || !pool)
+        // shutting down, retry again
+        return false;
+
     if (!pool->trySchedule(real_job))
         // Concurrent task limit reached
         return false;

dbms/src/Storages/DeltaMerge/LocalIndexerScheduler.h

@@ -93,6 +93,12 @@ class LocalIndexerScheduler
 
     ~LocalIndexerScheduler();
 
+    /**
+     * @brief Stop the scheduler and wait for running tasks to finish.
+     * Note that this method won't clear the task pushed.
+     */
+    void shutdown();
+
     /**
      * @brief Start the scheduler. In some tests we need to start scheduler
      * after some tasks are pushed.
@@ -101,7 +107,7 @@ class LocalIndexerScheduler
 
     /**
      * @brief Blocks until there is no tasks remaining in the queue and there is no running tasks.
-     * Should be only used in tests.
+     * **Should be only used in tests**.
      */
     void waitForFinish();
 
@@ -114,6 +120,7 @@ class LocalIndexerScheduler
 
     /**
     * @brief Drop all tasks matching specified keyspace id and table id.
+    * Note that this method won't drop the running tasks.
     */
     size_t dropTasks(KeyspaceID keyspace_id, TableID table_id);
 
@@ -147,9 +154,6 @@ class LocalIndexerScheduler
     void moveBackReadyTasks(std::unique_lock<std::mutex> & lock);
 
 private:
-    bool is_started = false;
-    std::thread scheduler_thread;
-
     /// Try to add a task to the pool. Returns false if the pool is full
     /// (for example, reaches concurrent task limit or memory limit).
     /// When pool is full, we will not try to schedule any more tasks at this moment.
@@ -160,6 +164,9 @@ class LocalIndexerScheduler
     /// heavy pressure.
     bool tryAddTaskToPool(std::unique_lock<std::mutex> & lock, const InternalTaskPtr & task);
 
+    std::thread scheduler_thread;
+    bool is_started = false;
+
     KeyspaceID last_schedule_keyspace_id = 0;
     std::map<KeyspaceID, TableID> last_schedule_table_id_by_ks;
 

dbms/src/Storages/DeltaMerge/Segment.cpp

@@ -417,24 +417,34 @@ SegmentPtr Segment::restoreSegment( //
     DMContext & context,
     PageIdU64 segment_id)
 {
-    Page page = context.storage_pool->metaReader()->read(segment_id); // not limit restore
-
-    ReadBufferFromMemory buf(page.data.begin(), page.data.size());
     Segment::SegmentMetaInfo segment_info;
-    readSegmentMetaInfo(buf, segment_info);
+    try
+    {
+        Page page = context.storage_pool->metaReader()->read(segment_id); // not limit restore
 
-    auto delta = DeltaValueSpace::restore(context, segment_info.range, segment_info.delta_id);
-    auto stable = StableValueSpace::restore(context, segment_info.stable_id);
-    auto segment = std::make_shared<Segment>(
-        parent_log,
-        segment_info.epoch,
-        segment_info.range,
-        segment_id,
-        segment_info.next_segment_id,
-        delta,
-        stable);
+        ReadBufferFromMemory buf(page.data.begin(), page.data.size());
+        readSegmentMetaInfo(buf, segment_info);
 
-    return segment;
+        auto delta = DeltaValueSpace::restore(context, segment_info.range, segment_info.delta_id);
+        auto stable = StableValueSpace::restore(context, segment_info.stable_id);
+        auto segment = std::make_shared<Segment>(
+            parent_log,
+            segment_info.epoch,
+            segment_info.range,
+            segment_id,
+            segment_info.next_segment_id,
+            delta,
+            stable);
+
+        return segment;
+    }
+    catch (DB::Exception & e)
+    {
+        e.addMessage(fmt::format("while restoreSegment, segment_id={}", segment_id));
+        e.rethrow();
+    }
+    RUNTIME_CHECK_MSG(false, "unreachable");
+    return {};
 }
 
 Segment::SegmentMetaInfos Segment::readAllSegmentsMetaInfoInRange( //

dbms/src/Storages/DeltaMerge/Segment.h

@@ -170,6 +170,19 @@ class Segment
         PageIdU64 next_segment_id{};
         PageIdU64 delta_id{};
         PageIdU64 stable_id{};
+
+        String toString() const
+        {
+            return fmt::format(
+                "{{version={} epoch={} range={} segment_id={} next_segment_id={} delta_id={} stable_id={}}}",
+                version,
+                epoch,
+                range.toString(),
+                segment_id,
+                next_segment_id,
+                delta_id,
+                stable_id);
+        }
     };
 
     using SegmentMetaInfos = std::vector<SegmentMetaInfo>;

dbms/src/Storages/DeltaMerge/tests/gtest_segment.cpp

@@ -48,11 +48,7 @@ namespace DB::ErrorCodes
 extern const int DT_DELTA_INDEX_ERROR;
 }
 
-namespace DB
-{
-namespace DM
-{
-namespace GC
+namespace DB::DM::GC
 {
 bool shouldCompactStableWithTooMuchDataOutOfSegmentRange(
     const DMContext & context, //
@@ -63,7 +59,7 @@ bool shouldCompactStableWithTooMuchDataOutOfSegmentRange(
     double invalid_data_ratio_threshold,
     const LoggerPtr & log);
 }
-namespace tests
+namespace DB::DM::tests
 {
 
 class SegmentOperationTest : public SegmentTestBasic
@@ -1366,7 +1362,4 @@ try
 }
 CATCH
 
-
-} // namespace tests
-} // namespace DM
-} // namespace DB
+} // namespace DB::DM::tests

dbms/src/Storages/Page/tools/PageCtl/PageStorageCtlV3.cpp

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Common/Checksum.h>
 #include <IO/Encryption/MockKeyManager.h>
 #include <IO/FileProvider/FileProvider.h>
 #include <Interpreters/Context.h>
@@ -871,7 +872,7 @@ class PageStorageControlV3
         ChecksumClass digest;
         digest.update(buffer, size);
         auto checksum = digest.checksum();
-        fmt::print("checksum: 0x{:X}\n", checksum);
+        fmt::println("checksum: 0x{:X}", checksum);
 
         auto hex_str = Redact::keyToHexString(buffer, size);
         delete[] buffer;