disagg: Fix unexpected object storage usage caused by pre-lock residue

dbms/src/Common/TiFlashMetrics.h

@@ -846,6 +846,15 @@ static_assert(RAFT_REGION_BIG_WRITE_THRES * 4 < RAFT_REGION_BIG_WRITE_MAX, "Inva
       F(type_to_finished, {"type", "to_finished"}),                                                                                 \
       F(type_to_error, {"type", "to_error"}),                                                                                       \
       F(type_to_cancelled, {"type", "to_cancelled"}))                                                                               \
+    M(tiflash_storage_s3_lock_mgr_status, "S3 Lock Manager", Gauge, F(type_prelock_keys, {{"type", "prelock_keys"}}))               \
+    M(tiflash_storage_s3_lock_mgr_counter,                                                                                          \
+      "S3 Lock Manager Counter",                                                                                                    \
+      Counter,                                                                                                                      \
+      F(type_create_lock_local, {{"type", "create_lock_local"}}),                                                                   \
+      F(type_create_lock_ingest, {{"type", "create_lock_ingest"}}),                                                                 \
+      F(type_clean_lock, {{"type", "clean_lock"}}),                                                                                 \
+      F(type_clean_lock_erase_hit, {{"type", "clean_lock_erase_hit"}}),                                                             \
+      F(type_clean_lock_erase_miss, {{"type", "clean_lock_erase_miss"}}))                                                           \
     M(tiflash_storage_s3_gc_status,                                                                                                 \
       "S3 GC status",                                                                                                               \
       Gauge,                                                                                                                        \

dbms/src/Storages/DeltaMerge/DeltaMergeStore_InternalSegment.cpp

@@ -1283,6 +1283,10 @@ SegmentPtr DeltaMergeStore::segmentMergeDelta(
 
         if (!isSegmentValid(lock, segment))
         {
+            LOG_INFO(
+                log,
+                "MergeDelta - Give up segmentMergeDelta because segment not valid, reason=concurrent_update segment={}",
+                segment->simpleInfo());
             LOG_DEBUG(
                 log,
                 "MergeDelta - Give up segmentMergeDelta because segment not valid, segment={}",

dbms/src/Storages/DeltaMerge/WriteBatchesImpl.h

@@ -171,6 +171,12 @@ struct WriteBatches : private boost::noncopyable
 
     void rollbackWrittenLogAndData()
     {
+        LOG_INFO(
+            Logger::get("WriteBatches"),
+            "Rollback written log/data, run_mode={} written_log_pages={} written_data_pages={}",
+            magic_enum::enum_name(run_mode),
+            written_log.size(),
+            written_data.size());
         WriteBatchWrapper log_wb(run_mode, keyspace_id, StorageType::Log, ns_id);
         for (auto p : written_log)
             log_wb.delPage(p);

dbms/src/Storages/Page/V3/BlobStore.cpp

@@ -405,6 +405,41 @@ typename BlobStore<Trait>::PageEntriesEdit BlobStore<Trait>::write(
     const size_t all_page_data_size = wb.getTotalDataSize();
 
     PageEntriesEdit edit;
+    auto log_remote_edit_summary = [&](const char * stage) {
+        size_t put_external_count = 0;
+        size_t put_remote_count = 0;
+        size_t checkpoint_info_count = 0;
+        std::vector<String> sample_data_file_ids;
+        for (const auto & rec : edit.getRecords())
+        {
+            if (rec.type == EditRecordType::PUT_EXTERNAL)
+                ++put_external_count;
+            else if (rec.type == EditRecordType::PUT)
+                ++put_remote_count;
+            if (rec.entry.checkpoint_info.has_value())
+            {
+                ++checkpoint_info_count;
+                if (sample_data_file_ids.size() < 8)
+                {
+                    sample_data_file_ids.emplace_back(*rec.entry.checkpoint_info.data_location.data_file_id);
+                }
+            }
+        }
+        if (checkpoint_info_count == 0)
+            return;
+        LOG_INFO(
+            log,
+            "BlobStore::write remote edit summary, stage={} wb_size={} edit_size={} put_external_count={} "
+            "put_remote_count={} checkpoint_info_count={} sample_data_file_ids={} wb={}",
+            stage,
+            wb.size(),
+            edit.size(),
+            put_external_count,
+            put_remote_count,
+            checkpoint_info_count,
+            sample_data_file_ids,
+            wb.toString());
+    };
 
     if (all_page_data_size == 0)
     {
@@ -491,6 +526,7 @@ typename BlobStore<Trait>::PageEntriesEdit BlobStore<Trait>::write(
                 throw Exception(fmt::format("Unknown write type: {}", magic_enum::enum_name(write.type)));
             }
         }
+        log_remote_edit_summary("zero_data_path");
         return edit;
     }
 
@@ -664,6 +700,7 @@ typename BlobStore<Trait>::PageEntriesEdit BlobStore<Trait>::write(
         throw e;
     }
 
+    log_remote_edit_summary("normal_path");
     return edit;
 }
 

dbms/src/Storages/Page/V3/PageDirectory.cpp

@@ -32,12 +32,14 @@
 #include <common/logger_useful.h>
 
 #include <magic_enum.hpp>
+#include <functional>
 #include <memory>
 #include <mutex>
 #include <optional>
 #include <shared_mutex>
 #include <type_traits>
 #include <utility>
+#include <vector>
 
 
 #ifdef FIU_ENABLE
@@ -1612,6 +1614,11 @@ std::unordered_set<String> PageDirectory<Trait>::apply(PageEntriesEdit && edit,
     CurrentMetrics::Increment pending_writer_size{CurrentMetrics::PSPendingWriterNum};
     Writer w;
     w.edit = &edit;
+    for (const auto & r : edit.getRecords())
+    {
+        if (r.entry.checkpoint_info.has_value())
+            w.applied_data_files.emplace(*r.entry.checkpoint_info.data_location.data_file_id);
+    }
 
     Stopwatch watch;
     std::unique_lock apply_lock(apply_mutex);
@@ -1640,8 +1647,8 @@ std::unordered_set<String> PageDirectory<Trait>::apply(PageEntriesEdit && edit,
             }
         }
         // the `applied_data_files` will be returned by the write
-        // group owner, others just return an empty set.
-        return {};
+        // group owner, followers return their own data_file_ids.
+        return std::move(w.applied_data_files);
     }
 
     /// This thread now is the write group owner, build the group. It will merge the
@@ -1704,12 +1711,53 @@ std::unordered_set<String> PageDirectory<Trait>::apply(PageEntriesEdit && edit,
 
     SYNC_FOR("before_PageDirectory::apply_to_memory");
     std::unordered_set<String> applied_data_files;
+    size_t edit_digest = 0;
+    size_t put_count = 0;
+    size_t put_external_count = 0;
+    size_t checkpoint_info_count = 0;
+    size_t put_external_with_checkpoint_info_count = 0;
+    std::vector<String> suspicious_records;
+    std::vector<String> candidate_records;
     {
         std::unique_lock table_lock(table_rw_mutex);
 
         // create entry version list for page_id.
         for (const auto & r : edit.getRecords())
         {
+            if (r.type == EditRecordType::PUT)
+                ++put_count;
+            else if (r.type == EditRecordType::PUT_EXTERNAL)
+                ++put_external_count;
+            if (r.entry.checkpoint_info.has_value())
+            {
+                ++checkpoint_info_count;
+                if (r.type == EditRecordType::PUT_EXTERNAL)
+                    ++put_external_with_checkpoint_info_count;
+            }
+            if (r.type == EditRecordType::PUT_EXTERNAL && !r.entry.checkpoint_info.has_value()
+                && suspicious_records.size() < 8)
+            {
+                suspicious_records.emplace_back(fmt::format("page_id={}", r.page_id));
+            }
+            if ((r.type == EditRecordType::PUT || r.type == EditRecordType::PUT_EXTERNAL) && candidate_records.size() < 12)
+            {
+                if (r.entry.checkpoint_info.has_value())
+                {
+                    candidate_records.emplace_back(fmt::format(
+                        "type={} page_id={} has_checkpoint=1 data_file_id={}",
+                        magic_enum::enum_name(r.type),
+                        r.page_id,
+                        *r.entry.checkpoint_info.data_location.data_file_id));
+                }
+                else
+                {
+                    candidate_records.emplace_back(fmt::format(
+                        "type={} page_id={} has_checkpoint=0",
+                        magic_enum::enum_name(r.type),
+                        r.page_id));
+                }
+            }
+
             // Protected in write_lock
             if (r.type == EditRecordType::DEL)
             {
@@ -1799,8 +1847,50 @@ std::unordered_set<String> PageDirectory<Trait>::apply(PageEntriesEdit && edit,
         sequence.fetch_add(edit_size);
     }
 
+    for (const auto & rec : candidate_records)
+    {
+        size_t rec_hash = std::hash<String>{}(rec);
+        edit_digest ^= rec_hash + 0x9e3779b97f4a7c15ULL + (edit_digest << 6) + (edit_digest >> 2);
+    }
+
+    if (checkpoint_info_count > 0)
+    {
+        // NOTE: test-only verbose diagnostics. Remove or lower level before production rollout.
+        LOG_INFO(
+            log,
+            "PageDirectory apply checkpoint summary, edit_size={} put_count={} put_external_count={} "
+            "checkpoint_info_count={} put_external_with_checkpoint_info_count={} edit_digest={} applied_data_files={} "
+            "candidate_records={} suspicious_records={}",
+            edit_size,
+            put_count,
+            put_external_count,
+            checkpoint_info_count,
+            put_external_with_checkpoint_info_count,
+            edit_digest,
+            applied_data_files,
+            candidate_records,
+            suspicious_records);
+    }
+
+    if (applied_data_files.empty() && checkpoint_info_count > 0)
+    {
+        LOG_INFO(
+            log,
+            "PageDirectory apply has no applied_data_files with checkpoint_info records, edit_size={} put_count={} "
+            "put_external_count={} checkpoint_info_count={} put_external_with_checkpoint_info_count={} edit_digest={} "
+            "suspicious_records={} candidate_records={}",
+            edit_size,
+            put_count,
+            put_external_count,
+            checkpoint_info_count,
+            put_external_with_checkpoint_info_count,
+            edit_digest,
+            suspicious_records,
+            candidate_records);
+    }
+
     success = true;
-    return applied_data_files;
+    return std::move(w.applied_data_files);
 }
 
 template <typename Trait>

dbms/src/Storages/Page/V3/PageDirectory.h

@@ -628,6 +628,7 @@ class PageDirectory
     struct Writer
     {
         PageEntriesEdit * edit;
+        std::unordered_set<String> applied_data_files;
         bool done = false; // The work has been performed by other thread
         bool success = false; // The work complete successfully
         std::unique_ptr<DB::Exception> exception;

dbms/src/Storages/Page/V3/Universal/S3LockLocalManager.cpp

@@ -14,6 +14,7 @@
 
 #include <Common/Exception.h>
 #include <Common/Stopwatch.h>
+#include <Common/TiFlashMetrics.h>
 #include <Flash/Disaggregated/S3LockClient.h>
 #include <Storages/Page/V3/CheckpointFile/CPManifestFileReader.h>
 #include <Storages/Page/V3/CheckpointFile/Proto/manifest_file.pb.h>
@@ -26,6 +27,8 @@
 
 #include <magic_enum.hpp>
 
+#include <vector>
+
 
 namespace DB::ErrorCodes
 {
@@ -40,7 +43,7 @@ S3LockLocalManager::S3LockLocalManager()
 {}
 
 // `store_id` is inited later because they may not
-// accessable when S3LockLocalManager is created.
+// accessible when S3LockLocalManager is created.
 std::optional<CheckpointProto::ManifestFilePrefix> S3LockLocalManager::initStoreInfo(
     StoreID actual_store_id,
     DB::S3::S3LockClientPtr s3lock_client_,
@@ -179,6 +182,18 @@ void S3LockLocalManager::createS3LockForWriteBatch(UniversalWriteBatch & write_b
         lock_key = std::make_shared<String>(lock_result);
     }
 
+    if (!s3_datafiles_to_lock.empty())
+    {
+        std::vector<String> lock_mappings;
+        lock_mappings.reserve(s3_datafiles_to_lock.size());
+        for (const auto & [input_key, lock_key] : s3_datafiles_to_lock)
+        {
+            if (lock_key)
+                lock_mappings.emplace_back(fmt::format("{} -> {}", input_key, *lock_key));
+        }
+        LOG_INFO(log, "S3 lock mapping for write batch, mapping_count={} mappings={}", lock_mappings.size(), lock_mappings);
+    }
+
     for (auto & w : write_batch.getMutWrites())
     {
         // Here we will replace the name to be the S3LockFile key name for later
@@ -224,7 +239,8 @@ String S3LockLocalManager::createS3Lock(
         // TODO: handle s3 network error and retry?
         auto s3_client = S3::ClientFactory::instance().sharedTiFlashClient();
         S3::uploadEmptyFile(*s3_client, lockkey);
-        LOG_DEBUG(log, "S3 lock created for local datafile, lockkey={}", lockkey);
+        GET_METRIC(tiflash_storage_s3_lock_mgr_counter, type_create_lock_local).Increment();
+        LOG_INFO(log, "S3 lock created for local datafile, datafile_key={} lockkey={}", datafile_key, lockkey);
     }
     else
     {
@@ -237,7 +253,8 @@ String S3LockLocalManager::createS3Lock(
         {
             throw Exception(ErrorCodes::S3_LOCK_CONFLICT, err_msg);
         }
-        LOG_DEBUG(log, "S3 lock created for ingest datafile, lockkey={}", lockkey);
+        GET_METRIC(tiflash_storage_s3_lock_mgr_counter, type_create_lock_ingest).Increment();
+        LOG_INFO(log, "S3 lock created for ingest datafile, datafile_key={} lockkey={}", datafile_key, lockkey);
     }
 
     // The related S3 data files in write batch is not applied into PageDirectory,
@@ -247,19 +264,58 @@ String S3LockLocalManager::createS3Lock(
     {
         std::unique_lock wlatch_keys(mtx_lock_keys);
         pre_lock_keys.emplace(lockkey);
+        GET_METRIC(tiflash_storage_s3_lock_mgr_status, type_prelock_keys).Set(pre_lock_keys.size());
     }
     return lockkey;
 }
 
 void S3LockLocalManager::cleanAppliedS3ExternalFiles(std::unordered_set<String> && applied_s3files)
 {
-    // After the entries applied into PageDirectory, manifest can get the S3 lock key
-    // from `VersionedPageEntries`, cleanup the pre lock files.
-    std::unique_lock wlatch_keys(mtx_lock_keys);
-    for (const auto & file : applied_s3files)
+    size_t erase_hit = 0;
+    size_t erase_miss = 0;
+    size_t remaining_pre_lock_keys = 0;
+    std::vector<String> erased_keys;
+    std::vector<String> missing_keys;
+    erased_keys.reserve(applied_s3files.size());
+    missing_keys.reserve(applied_s3files.size());
+    {
+        // After the entries applied into PageDirectory, manifest can get the S3 lock key
+        // from `VersionedPageEntries`, cleanup the pre lock files.
+        std::unique_lock wlatch_keys(mtx_lock_keys);
+        for (const auto & file : applied_s3files)
+        {
+            if (pre_lock_keys.erase(file) > 0)
+            {
+                ++erase_hit;
+                erased_keys.emplace_back(file);
+            }
+            else
+            {
+                ++erase_miss;
+                missing_keys.emplace_back(file);
+            }
+        }
+        remaining_pre_lock_keys = pre_lock_keys.size();
+        GET_METRIC(tiflash_storage_s3_lock_mgr_status, type_prelock_keys).Set(remaining_pre_lock_keys);
+    } // release the lock on mtx_lock_keys before logging
+    LOG_INFO(
+        log,
+        "Clean applied S3 external files, applied_count={} erase_hit={} erase_miss={} remaining_pre_lock_keys={}",
+        applied_s3files.size(),
+        erase_hit,
+        erase_miss,
+        remaining_pre_lock_keys);
+    for (const auto & key : erased_keys)
+    {
+        LOG_INFO(log, "Clean applied S3 external files, erase_hit lockkey={}", key);
+    }
+    for (const auto & key : missing_keys)
     {
-        pre_lock_keys.erase(file);
+        LOG_INFO(log, "Clean applied S3 external files, erase_miss lockkey={}", key);
     }
+    GET_METRIC(tiflash_storage_s3_lock_mgr_counter, type_clean_lock).Increment();
+    GET_METRIC(tiflash_storage_s3_lock_mgr_counter, type_clean_lock_erase_hit).Increment(erase_hit);
+    GET_METRIC(tiflash_storage_s3_lock_mgr_counter, type_clean_lock_erase_miss).Increment(erase_miss);
 }
 
 } // namespace DB::PS::V3