disagg: Fix TiFlash crash randomly when fails connect to s3 (#10676) (#10677)

close #10674

* Updates the async upload paths to capture shared resources by value and wait for all tasks to finish before returning, preventing early‑return lifetimes from leaving in‑flight tasks with dangling s3_client/file_provider references.
* Propagates the same safety pattern to `putCheckpointFiles`, `getDataFilesInfo`, `copyToLocal`, `setTaggingsForKeys`

Signed-off-by: JaySon-Huang <tshent@qq.com>

Co-authored-by: JaySon-Huang <tshent@qq.com>

Author: ti-chi-bot

dbms/src/Common/FailPoint.cpp

@@ -107,6 +107,7 @@ namespace DB
     M(force_fail_in_flush_region_data)                       \
     M(force_use_dmfile_format_v3)                            \
     M(force_set_mocked_s3_object_mtime)                      \
+    M(force_syncpoint_on_s3_upload)                          \
     M(force_stop_background_checkpoint_upload)               \
     M(force_schema_sync_diff_fail)                           \
     M(exception_after_large_write_exceed)                    \

dbms/src/Flash/Mpp/MPPTask.cpp

@@ -223,7 +223,7 @@ void MPPTask::finishWrite()
 
         // The finish of pushing all blocks not means that cte sink job has been done.
         // Execution summary statistic also need to be sent. So we can release cte
-        // only when execution sumary statistic has been sent.
+        // only when execution summary statistic has been sent.
         this->context->getCTEManager()->releaseCTEBySink(resp, this->dag_context->getQueryIDAndCTEIDForSink());
         this->notify_cte_sink_finish = true;
     }

dbms/src/IO/IOThreadPools.h

@@ -51,7 +51,7 @@ struct BuildReadTaskTrait
 class FutureContainer
 {
 public:
-    FutureContainer(const LoggerPtr & log_, size_t reserve_size = 0)
+    explicit FutureContainer(const LoggerPtr & log_, size_t reserve_size = 0)
         : log(log_)
     {
         futures.reserve(reserve_size);

dbms/src/Storages/DeltaMerge/Remote/DataStore/DataStoreS3.cpp

@@ -71,54 +71,55 @@ void DataStoreS3::putDMFileLocalFiles(
     auto s3_client = S3::ClientFactory::instance().sharedTiFlashClient();
 
     // First, upload non-meta files.
-    std::vector<std::future<void>> upload_results;
-    upload_results.reserve(local_files.size() - 1);
+    IOPoolHelper::FutureContainer upload_results(log, local_files.size());
     for (const auto & fname : local_files)
     {
         if (DMFileMetaV2::isMetaFileName(fname))
             continue;
 
         auto local_fname = fmt::format("{}/{}", local_dir, fname);
         auto remote_fname = fmt::format("{}/{}", remote_dir, fname);
-        auto task = std::make_shared<std::packaged_task<void()>>(
-            [&, local_fname = std::move(local_fname), remote_fname = std::move(remote_fname)]() -> void {
-                S3::uploadFile(
-                    *s3_client,
-                    local_fname,
-                    remote_fname,
-                    EncryptionPath(local_dir, fname, oid.keyspace_id),
-                    file_provider);
+        auto encryption_path = EncryptionPath(local_dir, fname, oid.keyspace_id);
+        // Capture shared resources by value in tasks to avoid dangling references on early errors.
+        auto task = std::make_shared<std::packaged_task<void()>>( //
+            [s3_client,
+             provider = file_provider,
+             local_fname = std::move(local_fname),
+             remote_fname = std::move(remote_fname),
+             encryption_path = std::move(encryption_path)]() -> void {
+                S3::uploadFile(*s3_client, local_fname, remote_fname, encryption_path, provider);
             });
-        upload_results.push_back(task->get_future());
+        upload_results.add(task->get_future());
         DataStoreS3Pool::get().scheduleOrThrowOnError([task]() { (*task)(); });
     }
-    for (auto & f : upload_results)
-        f.get();
+    // Wait for all tasks to finish before returning to keep captured resources alive.
+    upload_results.getAllResults();
 
     // Then, upload meta files.
     // Only when the meta upload is successful, the dmfile upload can be considered successful.
-    upload_results.clear();
+    IOPoolHelper::FutureContainer meta_upload_results(log, local_files.size());
     for (const auto & fname : local_files)
     {
         if (!DMFileMetaV2::isMetaFileName(fname))
             continue;
 
         auto local_fname = fmt::format("{}/{}", local_dir, fname);
         auto remote_fname = fmt::format("{}/{}", remote_dir, fname);
-        auto task = std::make_shared<std::packaged_task<void()>>(
-            [&, local_fname = std::move(local_fname), remote_fname = std::move(remote_fname)]() {
-                S3::uploadFile(
-                    *s3_client,
-                    local_fname,
-                    remote_fname,
-                    EncryptionPath(local_dir, fname, oid.keyspace_id),
-                    file_provider);
+        auto encryption_path = EncryptionPath(local_dir, fname, oid.keyspace_id);
+        // Capture shared resources by value in tasks to avoid dangling references on early errors.
+        auto task = std::make_shared<std::packaged_task<void()>>( //
+            [s3_client,
+             provider = file_provider,
+             local_fname = std::move(local_fname),
+             remote_fname = std::move(remote_fname),
+             encryption_path = std::move(encryption_path)]() {
+                S3::uploadFile(*s3_client, local_fname, remote_fname, encryption_path, provider);
             });
-        upload_results.push_back(task->get_future());
+        meta_upload_results.add(task->get_future());
         DataStoreS3Pool::get().scheduleOrThrowOnError([task]() { (*task)(); });
     }
-    for (auto & f : upload_results)
-        f.get();
+    // Wait for all tasks to finish before returning to keep captured resources alive.
+    meta_upload_results.getAllResults();
 
     LOG_INFO(log, "Upload DMFile finished, key={}, cost={}ms", remote_dir, sw.elapsedMilliseconds());
 }
@@ -134,30 +135,31 @@ bool DataStoreS3::putCheckpointFiles(
     /// then upload the CheckpointManifest to make the files within
     /// `upload_seq` public to S3GCManager.
 
-    std::vector<std::future<void>> upload_results;
-    // upload in parallel
+    // Upload in parallel.
     // Note: Local checkpoint files are always not encrypted.
+    IOPoolHelper::FutureContainer upload_results(log, local_files.data_files.size());
     for (size_t file_idx = 0; file_idx < local_files.data_files.size(); ++file_idx)
     {
-        auto task = std::make_shared<std::packaged_task<void()>>([&, idx = file_idx] {
-            const auto & local_datafile = local_files.data_files[idx];
-            auto s3key = S3::S3Filename::newCheckpointData(store_id, upload_seq, idx);
-            auto lock_key = s3key.toView().getLockKey(store_id, upload_seq);
-            S3::uploadFile(
-                *s3_client,
-                local_datafile,
-                s3key.toFullKey(),
-                EncryptionPath(local_datafile, "", NullspaceID),
-                file_provider);
-            S3::uploadEmptyFile(*s3_client, lock_key);
-        });
-        upload_results.push_back(task->get_future());
+        auto local_datafile = local_files.data_files[file_idx];
+        auto s3key = S3::S3Filename::newCheckpointData(store_id, upload_seq, file_idx);
+        auto remote_key = s3key.toFullKey();
+        auto lock_key = s3key.toView().getLockKey(store_id, upload_seq);
+        auto encryption_path = EncryptionPath(local_datafile, "", NullspaceID);
+        // Capture by value to avoid dangling references.
+        auto task = std::make_shared<std::packaged_task<void()>>( //
+            [s3_client,
+             provider = file_provider,
+             local_datafile = std::move(local_datafile),
+             remote_key = std::move(remote_key),
+             lock_key = std::move(lock_key),
+             encryption_path = std::move(encryption_path)] {
+                S3::uploadFile(*s3_client, local_datafile, remote_key, encryption_path, provider);
+                S3::uploadEmptyFile(*s3_client, lock_key);
+            });
+        upload_results.add(task->get_future());
         DataStoreS3Pool::get().scheduleOrThrowOnError([task] { (*task)(); });
     }
-    for (auto & f : upload_results)
-    {
-        f.get();
-    }
+    upload_results.getAllResults();
 
     // upload manifest after all CheckpointData uploaded
     auto s3key = S3::S3Filename::newCheckpointManifest(store_id, upload_seq);
@@ -180,7 +182,7 @@ std::unordered_map<String, IDataStore::DataFileInfo> DataStoreS3::getDataFilesIn
     for (const auto & lock_key : lock_keys)
     {
         auto task = std::make_shared<std::packaged_task<std::tuple<String, DataFileInfo>()>>(
-            [&s3_client, lock_key = lock_key, log = this->log]() noexcept {
+            [s3_client, lock_key = lock_key, log = this->log]() noexcept {
                 auto key_view = S3::S3FilenameView::fromKey(lock_key);
                 auto datafile_key = key_view.asDataFile().toFullKey();
                 try
@@ -232,43 +234,37 @@ void DataStoreS3::copyToLocal(
 {
     auto s3_client = S3::ClientFactory::instance().sharedTiFlashClient();
     const auto remote_dir = S3::S3Filename::fromDMFileOID(remote_oid).toFullKey();
-    std::vector<std::future<void>> results;
-    results.reserve(target_short_fnames.size());
+    IOPoolHelper::FutureContainer results(Logger::get("DataStoreS3"), target_short_fnames.size());
     for (const auto & fname : target_short_fnames)
     {
         auto remote_fname = fmt::format("{}/{}", remote_dir, fname);
         auto local_fname = fmt::format("{}/{}", local_dir, fname);
-        auto task = std::make_shared<std::packaged_task<void()>>(
-            [&, local_fname = std::move(local_fname), remote_fname = std::move(remote_fname)]() {
+        auto task = std::make_shared<std::packaged_task<void()>>( //
+            [s3_client, local_fname = std::move(local_fname), remote_fname = std::move(remote_fname)]() {
                 auto tmp_fname = fmt::format("{}.tmp", local_fname);
                 S3::downloadFile(*s3_client, tmp_fname, remote_fname);
                 Poco::File(tmp_fname).renameTo(local_fname);
             });
-        results.push_back(task->get_future());
+        results.add(task->get_future());
         DataStoreS3Pool::get().scheduleOrThrowOnError([task]() { (*task)(); });
     }
-    for (auto & f : results)
-    {
-        f.get();
-    }
+    results.getAllResults();
 }
 
 void DataStoreS3::setTaggingsForKeys(const std::vector<String> & keys, std::string_view tagging)
 {
     auto s3_client = S3::ClientFactory::instance().sharedTiFlashClient();
-    std::vector<std::future<void>> results;
-    results.reserve(keys.size());
+    IOPoolHelper::FutureContainer results(log, keys.size());
     for (const auto & k : keys)
     {
-        auto task = std::make_shared<std::packaged_task<void()>>(
-            [&s3_client, &tagging, key = k] { rewriteObjectWithTagging(*s3_client, key, String(tagging)); });
-        results.emplace_back(task->get_future());
+        auto task = std::make_shared<std::packaged_task<void()>>( //
+            [s3_client, tagging_str = String(tagging), key = k] {
+                rewriteObjectWithTagging(*s3_client, key, tagging_str);
+            });
+        results.add(task->get_future());
         DataStoreS3Pool::get().scheduleOrThrowOnError([task] { (*task)(); });
     }
-    for (auto & f : results)
-    {
-        f.get();
-    }
+    results.getAllResults();
 }
 
 IPreparedDMFileTokenPtr DataStoreS3::prepareDMFile(const S3::DMFileOID & oid, UInt64 page_id)

dbms/src/Storages/S3/S3Common.cpp

@@ -19,6 +19,7 @@
 #include <Common/RemoteHostFilter.h>
 #include <Common/Stopwatch.h>
 #include <Common/StringUtils/StringUtils.h>
+#include <Common/SyncPoint/SyncPoint.h>
 #include <Common/TiFlashMetrics.h>
 #include <IO/BaseFile/PosixRandomAccessFile.h>
 #include <IO/Buffer/ReadBufferFromRandomAccessFile.h>
@@ -156,6 +157,7 @@ class AWSLogger final : public Aws::Utils::Logging::LogSystemInterface
 namespace DB::FailPoints
 {
 extern const char force_set_mocked_s3_object_mtime[];
+extern const char force_syncpoint_on_s3_upload[];
 } // namespace DB::FailPoints
 namespace DB::S3
 {
@@ -659,6 +661,14 @@ static bool doUploadFile(
     {
         ProfileEvents::increment(is_dmfile ? ProfileEvents::S3PutDMFileRetry : ProfileEvents::S3PutObjectRetry);
     }
+#ifdef FIU_ENABLE
+    if (auto v = FailPointHelper::getFailPointVal(FailPoints::force_syncpoint_on_s3_upload); v)
+    {
+        const auto & prefix = std::any_cast<String>(v.value());
+        if (!prefix.empty() && startsWith(remote_fname, prefix))
+            SYNC_FOR("before_S3Common::uploadFile");
+    }
+#endif
     auto result = client.PutObject(req);
     if (!result.IsSuccess())
     {

dbms/src/Storages/S3/tests/gtest_s3file.cpp

@@ -14,9 +14,11 @@
 
 #include <Common/Exception.h>
 #include <Common/FailPoint.h>
+#include <Common/SyncPoint/Ctl.h>
 #include <IO/BaseFile/PosixWritableFile.h>
 #include <IO/Buffer/ReadBufferFromRandomAccessFile.h>
 #include <IO/Encryption/MockKeyManager.h>
+#include <IO/IOThreadPools.h>
 #include <Interpreters/Context.h>
 #include <Poco/DigestStream.h>
 #include <Poco/MD5Engine.h>
@@ -49,7 +51,9 @@
 #include <chrono>
 #include <ext/scope_guard.h>
 #include <fstream>
+#include <future>
 #include <memory>
+#include <thread>
 
 
 using namespace std::chrono_literals;
@@ -60,6 +64,7 @@ using namespace DB::S3;
 namespace DB::FailPoints
 {
 extern const char force_set_mocked_s3_object_mtime[];
+extern const char force_syncpoint_on_s3_upload[];
 extern const char force_s3_random_access_file_init_fail[];
 extern const char force_s3_random_access_file_read_fail[];
 } // namespace DB::FailPoints
@@ -774,6 +779,74 @@ try
 }
 CATCH
 
+TEST_P(S3FileTest, PutDMFileLocalFilesWaitsForAllTasks)
+try
+{
+    // This test requires MockS3Client control hooks.
+    if (dynamic_cast<DB::S3::tests::MockS3Client *>(s3_client.get()) == nullptr)
+        return;
+
+    auto & pool = DataStoreS3Pool::get();
+    // Ensure at least two threads so the blocked and failed uploads can run concurrently.
+    const auto old_max_threads = pool.getMaxThreads();
+    const auto old_queue_size = pool.getQueueSize();
+    SCOPE_EXIT({
+        pool.setMaxThreads(old_max_threads);
+        pool.setQueueSize(old_queue_size);
+    });
+    pool.setMaxThreads(std::max<size_t>(2, old_max_threads));
+    pool.setQueueSize(std::max<size_t>(2, old_queue_size));
+
+    const String local_dir = fmt::format("{}/dmf_local", getTemporaryPath());
+    createIfNotExist(local_dir);
+    std::vector<String> local_files{"failed.dat", "blocked.dat"};
+    writeLocalFile(fmt::format("{}/{}", local_dir, local_files[0]), 16);
+    writeLocalFile(fmt::format("{}/{}", local_dir, local_files[1]), 16);
+
+    const S3::DMFileOID oid{
+        .store_id = 1,
+        .keyspace_id = keyspace_id,
+        .table_id = 100,
+        .file_id = 1,
+    };
+    const auto remote_dir = S3::S3Filename::fromDMFileOID(oid).toFullKey();
+
+    // Set up a syncpoint to block the "blocked.dat" from being uploaded.
+    const auto blocked_key = fmt::format("{}/{}", remote_dir, local_files[1]);
+    FailPointHelper::enableFailPoint(FailPoints::force_syncpoint_on_s3_upload, blocked_key);
+    SCOPE_EXIT({ FailPointHelper::disableFailPoint(FailPoints::force_syncpoint_on_s3_upload); });
+
+    auto sp_upload = SyncPointCtl::enableInScope("before_S3Common::uploadFile");
+
+    // The "failed.dat" upload should fail.
+    DB::S3::tests::MockS3Client::setPutObjectStatus(DB::S3::tests::MockS3Client::S3Status::FAILED);
+    SCOPE_EXIT({ DB::S3::tests::MockS3Client::setPutObjectStatus(DB::S3::tests::MockS3Client::S3Status::NORMAL); });
+
+    std::promise<void> done;
+    auto done_future = done.get_future();
+    // Run in a separate thread so we can observe whether it returns before "blocked.dat" upload is unblocked.
+    std::thread worker([&] {
+        try
+        {
+            data_store->putDMFileLocalFiles(local_dir, local_files, oid);
+        }
+        catch (...)
+        {}
+        done.set_value();
+    });
+
+    sp_upload.waitAndPause();
+
+    // The `putDMFileLocalFiles` should not return while "blocked.dat" is still blocked.
+    EXPECT_EQ(done_future.wait_for(200ms), std::future_status::timeout);
+
+    // Avoid pausing again on retries after unblocking.
+    FailPointHelper::disableFailPoint(FailPoints::force_syncpoint_on_s3_upload);
+    sp_upload.next();
+    worker.join();
+}
+CATCH
+
 INSTANTIATE_TEST_CASE_P(S3File, S3FileTest, testing::Values(false, true));
 
 } // namespace DB::tests