Merge pull request ClickHouse#76847 from ClickHouse/stoprep

Fix race in DatabaseReplicated between stopReplication and startupDatabaseAsync

Author: al13n321

programs/server/Server.cpp

@@ -2398,7 +2398,7 @@ try
         LOG_INFO(log, "Stopping AsyncLoader.");
 
         // Waits for all currently running jobs to finish and do not run any other pending jobs.
-        global_context->getAsyncLoader().stop();
+        global_context->getAsyncLoader().shutdown();
     );
 
     try

src/Common/AsyncLoader.cpp

@@ -241,14 +241,7 @@ AsyncLoader::~AsyncLoader()
     // When all jobs are done we could still have finalizing workers.
     // These workers could call updateCurrentPriorityAndSpawn() that scans all pools.
     // We need to stop all of them before destructing any of them.
-    stop();
-}
-
-void AsyncLoader::start()
-{
-    std::unique_lock lock{mutex};
-    is_running = true;
-    updateCurrentPriorityAndSpawn(lock);
+    shutdown();
 }
 
 void AsyncLoader::wait()
@@ -277,7 +270,27 @@ void AsyncLoader::wait()
     }
 }
 
-void AsyncLoader::stop()
+void AsyncLoader::shutdown()
+{
+    LoadJobSet jobs;
+
+    {
+        std::unique_lock lock{mutex};
+        shutdown_requested = true;
+        is_running = false;
+
+        for (const auto & [job, _] : scheduled_jobs)
+            jobs.insert(job);
+    }
+
+    // Cancel scheduled jobs, wait for currently running jobs to finish.
+    remove(jobs);
+
+    for (auto & p : pools)
+        p.thread_pool->wait();
+}
+
+void AsyncLoader::pause()
 {
     {
         std::unique_lock lock{mutex};
@@ -289,6 +302,13 @@ void AsyncLoader::stop()
         p.thread_pool->wait();
 }
 
+void AsyncLoader::unpause()
+{
+    std::unique_lock lock{mutex};
+    is_running = true;
+    updateCurrentPriorityAndSpawn(lock);
+}
+
 void AsyncLoader::schedule(LoadTask & task)
 {
     chassert(this == &task.loader);
@@ -331,6 +351,12 @@ void AsyncLoader::schedule(const LoadJobSet & jobs_to_schedule)
     for (const auto & job : jobs_to_schedule)
         gatherNotScheduled(job, jobs, lock);
 
+    if (jobs.empty())
+        return;
+
+    if (shutdown_requested)
+        throw Exception(ErrorCodes::ASYNC_LOAD_CANCELED, "AsyncLoader was shut down");
+
     // Ensure scheduled_jobs graph will have no cycles. The only way to get a cycle is to add a cycle, assuming old jobs cannot reference new ones.
     checkCycle(jobs, lock);
 

src/Common/AsyncLoader.h

@@ -395,17 +395,20 @@ class AsyncLoader : private boost::noncopyable
     // WARNING: all tasks instances should be destructed before associated AsyncLoader.
     ~AsyncLoader();
 
-    // Start workers to execute scheduled load jobs. Note that AsyncLoader is constructed as already started.
-    void start();
-
     // Wait for all load jobs to finish, including all new jobs. So at first take care to stop adding new jobs.
     void wait();
 
+    // Wait for currently executing jobs to finish, cancel pending jobs with an exception,
+    // prevent scheduling of new jobs.
+    void shutdown();
+
     // Wait for currently executing jobs to finish, but do not run any other pending jobs.
     // Not finished jobs are left in pending state:
-    //  - they can be executed by calling start() again;
+    //  - they can be executed by calling unpause() again;
     //  - or canceled using ~Task() or remove() later.
-    void stop();
+    // Currently only used in tests.
+    void pause();
+    void unpause();
 
     // Schedule all jobs of given `task` and their dependencies (even if they are not in task).
     // All dependencies of a scheduled job inherit its pool if it has higher priority. This way higher priority job
@@ -484,6 +487,7 @@ class AsyncLoader : private boost::noncopyable
 
     mutable std::mutex mutex; // Guards all the fields below.
     bool is_running = true;
+    bool shutdown_requested = false;
     std::optional<Priority> current_priority; // highest priority among active pools
     UInt64 last_ready_seqno = 0; // Increasing counter for ready queue keys.
     UInt64 last_job_id = 0; // Increasing counter for job IDs

src/Common/FailPoint.cpp

@@ -87,6 +87,7 @@ static struct InitFiu
     REGULAR(zero_copy_unlock_zk_fail_after_op) \
     REGULAR(plain_rewritable_object_storage_azure_not_found_on_init) \
     PAUSEABLE(storage_merge_tree_background_clear_old_parts_pause) \
+    PAUSEABLE(database_replicated_startup_pause) \
 
 
 namespace FailPoints

src/Common/tests/gtest_async_loader.cpp

@@ -53,7 +53,7 @@ struct AsyncLoaderTest
     explicit AsyncLoaderTest(std::vector<Initializer> initializers)
         : loader(getPoolInitializers(initializers), /* log_failures = */ false, /* log_progress = */ false, /* log_events = */ false)
     {
-        loader.stop(); // All tests call `start()` manually to better control ordering
+        loader.pause(); // All tests call `unpause()` manually to better control ordering
     }
 
     explicit AsyncLoaderTest(size_t max_threads = 1)
@@ -181,7 +181,7 @@ TEST(AsyncLoader, Smoke)
 
         std::thread waiter_thread([&t, job5] { t.loader.wait(job5); });
 
-        t.loader.start();
+        t.loader.unpause();
 
         t.loader.wait(job3);
         t.loader.wait();
@@ -196,7 +196,7 @@ TEST(AsyncLoader, Smoke)
     ASSERT_EQ(jobs_done, 5);
     ASSERT_EQ(low_priority_jobs_done, 1);
 
-    t.loader.stop();
+    t.loader.pause();
 }
 
 TEST(AsyncLoader, CycleDetection)
@@ -348,7 +348,7 @@ TEST(AsyncLoader, CancelPendingDependency)
 TEST(AsyncLoader, CancelExecutingJob)
 {
     AsyncLoaderTest t;
-    t.loader.start();
+    t.loader.unpause();
 
     std::barrier sync(2);
 
@@ -380,7 +380,7 @@ TEST(AsyncLoader, CancelExecutingJob)
 TEST(AsyncLoader, CancelExecutingTask)
 {
     AsyncLoaderTest t(16);
-    t.loader.start();
+    t.loader.unpause();
     std::barrier sync(2);
 
     auto blocker_job_func = [&] (AsyncLoader &, const LoadJobPtr &)
@@ -436,7 +436,7 @@ TEST(AsyncLoader, CancelExecutingTask)
 TEST(AsyncLoader, JobFailure)
 {
     AsyncLoaderTest t;
-    t.loader.start();
+    t.loader.unpause();
 
     std::string error_message = "test job failure";
 
@@ -466,7 +466,7 @@ TEST(AsyncLoader, JobFailure)
 TEST(AsyncLoader, ScheduleJobWithFailedDependencies)
 {
     AsyncLoaderTest t;
-    t.loader.start();
+    t.loader.unpause();
 
     std::string_view error_message = "test job failure";
 
@@ -522,7 +522,7 @@ TEST(AsyncLoader, ScheduleJobWithCanceledDependencies)
     auto canceled_task = t.schedule({ canceled_job });
     canceled_task->remove();
 
-    t.loader.start();
+    t.loader.unpause();
 
     auto job_func = [&] (AsyncLoader &, const LoadJobPtr &) {};
     auto job1 = makeLoadJob({ canceled_job }, "job1", job_func);
@@ -559,7 +559,7 @@ TEST(AsyncLoader, IgnoreDependencyFailure)
 {
     AsyncLoaderTest t;
     std::atomic<bool> success{false};
-    t.loader.start();
+    t.loader.unpause();
 
     std::string_view error_message = "test job failure";
 
@@ -588,7 +588,7 @@ TEST(AsyncLoader, CustomDependencyFailure)
     int error_count = 0;
     std::atomic<size_t> good_count{0};
     std::barrier canceled_sync(4);
-    t.loader.start();
+    t.loader.unpause();
 
     std::string_view error_message = "test job failure";
 
@@ -675,7 +675,7 @@ TEST(AsyncLoader, WaitersLimit)
     };
 
     std::barrier sync(2);
-    t.loader.start();
+    t.loader.unpause();
 
     auto job_func = [&] (AsyncLoader &, const LoadJobPtr &) {
         sync.arrive_and_wait(); // (A)
@@ -723,7 +723,7 @@ TEST(AsyncLoader, WaitersLimit)
 TEST(AsyncLoader, TestConcurrency)
 {
     AsyncLoaderTest t(10);
-    t.loader.start();
+    t.loader.unpause();
 
     for (int concurrency = 1; concurrency <= 10; concurrency++)
     {
@@ -750,7 +750,7 @@ TEST(AsyncLoader, TestConcurrency)
 TEST(AsyncLoader, TestOverload)
 {
     AsyncLoaderTest t(3);
-    t.loader.start();
+    t.loader.unpause();
 
     size_t max_threads = t.loader.getMaxThreads(/* pool = */ 0);
     std::atomic<int> executing{0};
@@ -765,12 +765,12 @@ TEST(AsyncLoader, TestOverload)
             executing--;
         };
 
-        t.loader.stop();
+        t.loader.pause();
         std::vector<LoadTaskPtr> tasks;
         tasks.reserve(concurrency);
         for (int i = 0; i < concurrency; i++)
             tasks.push_back(t.schedule(t.chainJobSet(5, job_func)));
-        t.loader.start();
+        t.loader.unpause();
         t.loader.wait();
         ASSERT_EQ(executing, 0);
     }
@@ -818,7 +818,7 @@ TEST(AsyncLoader, StaticPriorities)
     jobs.push_back(makeLoadJob({ jobs[6] }, 9, "H", job_func)); // 7
     auto task = t.schedule({ jobs.begin(), jobs.end() });
 
-    t.loader.start();
+    t.loader.unpause();
     t.loader.wait();
     ASSERT_TRUE(schedule == "A9E9D9F9G9H9C4B3" || schedule == "A9D9E9F9G9H9C4B3");
 }
@@ -831,7 +831,7 @@ TEST(AsyncLoader, SimplePrioritization)
         {.max_threads = 1, .priority{-2}},
     });
 
-    t.loader.start();
+    t.loader.unpause();
 
     std::atomic<int> executed{0}; // Number of previously executed jobs (to test execution order)
     LoadJobPtr job_to_prioritize;
@@ -951,9 +951,9 @@ TEST(AsyncLoader, DynamicPriorities)
 
         job_to_prioritize = jobs[6]; // G
 
-        t.loader.start();
+        t.loader.unpause();
         t.loader.wait();
-        t.loader.stop();
+        t.loader.pause();
 
         if (prioritize)
         {
@@ -1000,7 +1000,7 @@ TEST(AsyncLoader, JobPrioritizedWhileWaited)
 
     job_to_wait = jobs[1];
 
-    t.loader.start();
+    t.loader.unpause();
 
     while (job_to_wait->waitersCount() == 0)
         std::this_thread::yield();
@@ -1011,15 +1011,15 @@ TEST(AsyncLoader, JobPrioritizedWhileWaited)
     sync.arrive_and_wait();
 
     t.loader.wait();
-    t.loader.stop();
+    t.loader.pause();
     ASSERT_EQ(t.loader.suspendedWorkersCount(1), 0);
     ASSERT_EQ(t.loader.suspendedWorkersCount(0), 0);
 }
 
 TEST(AsyncLoader, RandomIndependentTasks)
 {
     AsyncLoaderTest t(16);
-    t.loader.start();
+    t.loader.unpause();
 
     auto job_func = [&] (AsyncLoader &, const LoadJobPtr & self)
     {
@@ -1041,7 +1041,7 @@ TEST(AsyncLoader, RandomIndependentTasks)
 TEST(AsyncLoader, RandomDependentTasks)
 {
     AsyncLoaderTest t(16);
-    t.loader.start();
+    t.loader.unpause();
 
     std::mutex mutex;
     std::condition_variable cv;
@@ -1108,7 +1108,7 @@ TEST(AsyncLoader, SetMaxThreads)
     for (int i = 0; i < 1000; i++)
         tasks.push_back(t.schedule({makeLoadJob({}, "job", job_func)}));
 
-    t.loader.start();
+    t.loader.unpause();
     while (sync_index < syncs.size())
     {
         // Wait for `max_threads` jobs to start executing
@@ -1132,7 +1132,7 @@ TEST(AsyncLoader, SetMaxThreads)
 TEST(AsyncLoader, SubJobs)
 {
     AsyncLoaderTest t(1);
-    t.loader.start();
+    t.loader.unpause();
 
     // An example of component with an asynchronous loading interface
     class MyComponent : boost::noncopyable {
@@ -1195,7 +1195,7 @@ TEST(AsyncLoader, SubJobs)
 TEST(AsyncLoader, RecursiveJob)
 {
     AsyncLoaderTest t(1);
-    t.loader.start();
+    t.loader.unpause();
 
     // An example of component with an asynchronous loading interface (a complicated one)
     class MyComponent : boost::noncopyable {

src/Databases/DatabaseReplicated.cpp

@@ -40,6 +40,7 @@
 #include <base/defines.h>
 #include <base/getFQDNOrHostName.h>
 #include <Common/Exception.h>
+#include <Common/FailPoint.h>
 #include <Common/Macros.h>
 #include <Common/OpenTelemetryTraceContext.h>
 #include <Common/PoolId.h>
@@ -95,6 +96,12 @@ namespace ErrorCodes
     extern const int CANNOT_RESTORE_TABLE;
     extern const int QUERY_IS_PROHIBITED;
     extern const int SUPPORT_IS_DISABLED;
+    extern const int ASYNC_LOAD_CANCELED;
+}
+
+namespace FailPoints
+{
+    extern const char database_replicated_startup_pause[];
 }
 
 static constexpr const char * REPLICATED_DATABASE_MARK = "DatabaseReplicated";
@@ -767,6 +774,8 @@ LoadTaskPtr DatabaseReplicated::startupDatabaseAsync(AsyncLoader & async_loader,
             if (is_probably_dropped)
                 return;
 
+            FailPointInjection::pauseFailPoint(FailPoints::database_replicated_startup_pause);
+
             {
                 std::lock_guard lock{ddl_worker_mutex};
                 ddl_worker = std::make_unique<DatabaseReplicatedDDLWorker>(this, getContext());
@@ -1695,6 +1704,17 @@ void DatabaseReplicated::renameDatabase(ContextPtr query_context, const String &
 
 void DatabaseReplicated::stopReplication()
 {
+    try
+    {
+        /// Make sure startupDatabaseAsync doesn't start ddl_worker after stopReplication().
+        waitDatabaseStarted();
+    }
+    catch (Exception & e)
+    {
+        if (e.code() != ErrorCodes::ASYNC_LOAD_CANCELED)
+            tryLogCurrentException("DatabaseReplicated", "Async loading failed", LogsLevel::warning);
+    }
+
     std::lock_guard lock{ddl_worker_mutex};
     if (ddl_worker)
         ddl_worker->shutdown();