Storage: Refactor DMContext to avoid being misused

dbms/src/Operators/DMSegmentThreadSourceOp.cpp

@@ -76,7 +76,7 @@ OperatorStatus DMSegmentThreadSourceOp::readImpl(Block & block)
 
 OperatorStatus DMSegmentThreadSourceOp::executeIOImpl()
 {
-    if unlikely (done)
+    if (unlikely(done))
         return OperatorStatus::HAS_OUTPUT;
 
     while (!cur_stream)
@@ -92,7 +92,7 @@ OperatorStatus DMSegmentThreadSourceOp::executeIOImpl()
 
         auto block_size = std::max(
             expected_block_size,
-            static_cast<size_t>(dm_context->db_context.getSettingsRef().dt_segment_stable_pack_rows));
+            static_cast<size_t>(dm_context->global_context.getSettingsRef().dt_segment_stable_pack_rows));
         cur_stream = task->segment->getInputStream(
             read_mode,
             *dm_context,

dbms/src/Operators/DMSegmentThreadSourceOp.h

@@ -15,7 +15,7 @@
 #pragma once
 
 #include <Operators/Operator.h>
-#include <Storages/DeltaMerge/DMContext.h>
+#include <Storages/DeltaMerge/DMContext_fwd.h>
 #include <Storages/DeltaMerge/Segment.h>
 #include <Storages/DeltaMerge/SegmentReadTaskPool.h>
 

dbms/src/Server/DTTool/DTToolBench.cpp

@@ -365,7 +365,7 @@ int benchEntry(const std::vector<std::string> & opts)
         auto storage_pool
             = std::make_shared<DB::DM::StoragePool>(*db_context, NullspaceID, /*ns_id*/ 1, *path_pool, "test.t1");
         auto dm_settings = DB::DM::DeltaMergeStore::Settings{};
-        auto dm_context = std::make_unique<DB::DM::DMContext>( //
+        auto dm_context = DB::DM::DMContext::createUnique(
             *db_context,
             path_pool,
             storage_pool,

dbms/src/Server/tests/gtest_dttool.cpp

@@ -84,7 +84,7 @@ struct DTToolTest : public DB::base::TiFlashStorageTestBasic
         auto storage_pool
             = std::make_shared<DB::DM::StoragePool>(*db_context, NullspaceID, /*ns_id*/ 1, *path_pool, "test.t1");
         auto dm_settings = DB::DM::DeltaMergeStore::Settings{};
-        auto dm_context = std::make_unique<DB::DM::DMContext>( //
+        auto dm_context = DB::DM::DMContext::createUnique(
             *db_context,
             path_pool,
             storage_pool,

dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileBig.cpp

@@ -36,9 +36,9 @@ ColumnFileBig::ColumnFileBig(const DMContext & context, const DMFilePtr & file_,
     calculateStat(context);
 }
 
-void ColumnFileBig::calculateStat(const DMContext & context)
+void ColumnFileBig::calculateStat(const DMContext & dm_context)
 {
-    auto index_cache = context.db_context.getGlobalContext().getMinMaxIndexCache();
+    auto index_cache = dm_context.global_context.getMinMaxIndexCache();
 
     auto pack_filter = DMFilePackFilter::loadFrom(
         file,
@@ -47,10 +47,10 @@ void ColumnFileBig::calculateStat(const DMContext & context)
         {segment_range},
         EMPTY_RS_OPERATOR,
         {},
-        context.db_context.getFileProvider(),
-        context.getReadLimiter(),
-        context.scan_context,
-        /*tracing_id*/ context.tracing_id);
+        dm_context.global_context.getFileProvider(),
+        dm_context.getReadLimiter(),
+        dm_context.scan_context,
+        /*tracing_id*/ dm_context.tracing_id);
 
     std::tie(valid_rows, valid_bytes) = pack_filter.validRowsAndBytes();
 }
@@ -79,7 +79,7 @@ void ColumnFileBig::serializeMetadata(WriteBuffer & buf, bool /*save_schema*/) c
 }
 
 ColumnFilePersistedPtr ColumnFileBig::deserializeMetadata(
-    const DMContext & context, //
+    const DMContext & dm_context, //
     const RowKeyRange & segment_range,
     ReadBuffer & buf)
 {
@@ -92,13 +92,16 @@ ColumnFilePersistedPtr ColumnFileBig::deserializeMetadata(
 
     String file_parent_path;
     DMFilePtr dmfile;
-    auto remote_data_store = context.db_context.getSharedContextDisagg()->remote_data_store;
-    auto path_delegate = context.path_pool->getStableDiskDelegator();
+    auto remote_data_store = dm_context.global_context.getSharedContextDisagg()->remote_data_store;
+    auto path_delegate = dm_context.path_pool->getStableDiskDelegator();
     if (remote_data_store)
     {
-        auto wn_ps = context.db_context.getWriteNodePageStorage();
+        auto wn_ps = dm_context.global_context.getWriteNodePageStorage();
         auto full_page_id = UniversalPageIdFormat::toFullPageId(
-            UniversalPageIdFormat::toFullPrefix(context.keyspace_id, StorageType::Data, context.physical_table_id),
+            UniversalPageIdFormat::toFullPrefix(
+                dm_context.keyspace_id,
+                StorageType::Data,
+                dm_context.physical_table_id),
             file_page_id);
         auto full_external_id = wn_ps->getNormalPageId(full_page_id);
         auto local_external_id = UniversalPageIdFormat::getU64ID(full_external_id);
@@ -112,10 +115,10 @@ ColumnFilePersistedPtr ColumnFileBig::deserializeMetadata(
     }
     else
     {
-        auto file_id = context.storage_pool->dataReader()->getNormalPageId(file_page_id);
-        file_parent_path = context.path_pool->getStableDiskDelegator().getDTFilePath(file_id);
+        auto file_id = dm_context.storage_pool->dataReader()->getNormalPageId(file_page_id);
+        file_parent_path = dm_context.path_pool->getStableDiskDelegator().getDTFilePath(file_id);
         dmfile = DMFile::restore(
-            context.db_context.getFileProvider(),
+            dm_context.global_context.getFileProvider(),
             file_id,
             file_page_id,
             file_parent_path,
@@ -129,7 +132,7 @@ ColumnFilePersistedPtr ColumnFileBig::deserializeMetadata(
 }
 
 ColumnFilePersistedPtr ColumnFileBig::createFromCheckpoint(
-    DMContext & context, //
+    DMContext & dm_context, //
     const RowKeyRange & target_range,
     ReadBuffer & buf,
     UniversalPageStoragePtr temp_ps,
@@ -143,21 +146,21 @@ ColumnFilePersistedPtr ColumnFileBig::createFromCheckpoint(
     readIntBinary(valid_bytes, buf);
 
     auto remote_page_id = UniversalPageIdFormat::toFullPageId(
-        UniversalPageIdFormat::toFullPrefix(context.keyspace_id, StorageType::Data, context.physical_table_id),
+        UniversalPageIdFormat::toFullPrefix(dm_context.keyspace_id, StorageType::Data, dm_context.physical_table_id),
         file_page_id);
     auto remote_data_location = temp_ps->getCheckpointLocation(remote_page_id);
     auto data_key_view = S3::S3FilenameView::fromKey(*(remote_data_location->data_file_id)).asDataFile();
     auto file_oid = data_key_view.getDMFileOID();
     auto data_key = data_key_view.toFullKey();
-    auto delegator = context.path_pool->getStableDiskDelegator();
-    auto new_local_page_id = context.storage_pool->newDataPageIdForDTFile(delegator, __PRETTY_FUNCTION__);
+    auto delegator = dm_context.path_pool->getStableDiskDelegator();
+    auto new_local_page_id = dm_context.storage_pool->newDataPageIdForDTFile(delegator, __PRETTY_FUNCTION__);
     PS::V3::CheckpointLocation loc{
         .data_file_id = std::make_shared<String>(data_key),
         .offset_in_file = 0,
         .size_in_file = 0,
     };
     wbs.data.putRemoteExternal(new_local_page_id, loc);
-    auto remote_data_store = context.db_context.getSharedContextDisagg()->remote_data_store;
+    auto remote_data_store = dm_context.global_context.getSharedContextDisagg()->remote_data_store;
     auto prepared = remote_data_store->prepareDMFile(file_oid, new_local_page_id);
     auto dmfile = prepared->restore(DMFile::ReadMetaMode::all());
     wbs.writeLogAndData();
@@ -172,10 +175,13 @@ void ColumnFileBigReader::initStream()
     if (file_stream)
         return;
 
-    DMFileBlockInputStreamBuilder builder(context.db_context);
-    file_stream
-        = builder.setTracingID(context.tracing_id)
-              .build(column_file.getFile(), *col_defs, RowKeyRanges{column_file.segment_range}, context.scan_context);
+    DMFileBlockInputStreamBuilder builder(dm_context.global_context);
+    file_stream = builder.setTracingID(dm_context.tracing_id)
+                      .build(
+                          column_file.getFile(),
+                          *col_defs,
+                          RowKeyRanges{column_file.segment_range},
+                          dm_context.scan_context);
 
     header = file_stream->getHeader();
     // If we only need to read pk and version columns, then cache columns data in memory.
@@ -383,7 +389,7 @@ size_t ColumnFileBigReader::skipNextBlock()
 ColumnFileReaderPtr ColumnFileBigReader::createNewReader(const ColumnDefinesPtr & new_col_defs)
 {
     // Currently we don't reuse the cache data.
-    return std::make_shared<ColumnFileBigReader>(context, column_file, new_col_defs);
+    return std::make_shared<ColumnFileBigReader>(dm_context, column_file, new_col_defs);
 }
 
 } // namespace DM

dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileBig.h

@@ -114,7 +114,7 @@ class ColumnFileBig : public ColumnFilePersisted
 class ColumnFileBigReader : public ColumnFileReader
 {
 private:
-    const DMContext & context;
+    const DMContext & dm_context;
     const ColumnFileBig & column_file;
     const ColumnDefinesPtr col_defs;
 
@@ -155,7 +155,7 @@ class ColumnFileBigReader : public ColumnFileReader
         const DMContext & context_,
         const ColumnFileBig & column_file_,
         const ColumnDefinesPtr & col_defs_)
-        : context(context_)
+        : dm_context(context_)
         , column_file(column_file_)
         , col_defs(col_defs_)
     {

dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFilePersisted.cpp

@@ -15,6 +15,7 @@
 #include <IO/CompressedReadBuffer.h>
 #include <IO/CompressedWriteBuffer.h>
 #include <IO/MemoryReadWriteBuffer.h>
+#include <IO/ReadBufferFromString.h>
 #include <Storages/DeltaMerge/ColumnFile/ColumnFileBig.h>
 #include <Storages/DeltaMerge/ColumnFile/ColumnFileDeleteRange.h>
 #include <Storages/DeltaMerge/ColumnFile/ColumnFileInMemory.h>

dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFilePersisted.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <IO/CompressedStream.h>
 #include <IO/MemoryReadWriteBuffer.h>
 #include <Storages/DeltaMerge/ColumnFile/ColumnFile.h>
 

dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileSchema.cpp

@@ -15,6 +15,7 @@
 #include <Common/TiFlashMetrics.h>
 #include <Interpreters/Context.h>
 #include <Storages/DeltaMerge/ColumnFile/ColumnFileSchema.h>
+#include <Storages/DeltaMerge/DMContext.h>
 
 namespace DB
 {
@@ -129,9 +130,9 @@ ColumnFileSchemaPtr SharedBlockSchemas::getOrCreate(const Block & block)
     return schema;
 }
 
-std::shared_ptr<DB::DM::SharedBlockSchemas> getSharedBlockSchemas(const DMContext & context)
+std::shared_ptr<DB::DM::SharedBlockSchemas> getSharedBlockSchemas(const DMContext & dm_context)
 {
-    return context.db_context.getSharedBlockSchemas();
+    return dm_context.global_context.getSharedBlockSchemas();
 }
 } // namespace DM
-} // namespace DB
+} // namespace DB

dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileSchema.h

@@ -17,6 +17,7 @@
 #include <Core/Block.h>
 #include <Interpreters/Context_fwd.h>
 #include <Storages/BackgroundProcessingPool.h>
+#include <Storages/DeltaMerge/DMContext_fwd.h>
 #include <Storages/DeltaMerge/DeltaMergeDefines.h>
 #include <Storages/DeltaMerge/DeltaMergeHelpers.h>
 #include <Storages/DeltaMerge/Remote/Serializer_fwd.h>

dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileSetReader.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include <Storages/DeltaMerge/ColumnFile/ColumnFileSetSnapshot.h>
+#include <Storages/DeltaMerge/DMContext.h>
 #include <Storages/DeltaMerge/SkippableBlockInputStream.h>
 
 namespace DB

dbms/src/Storages/DeltaMerge/ColumnFile/ColumnFileTiny.cpp

@@ -277,13 +277,13 @@ ColumnFileTinyPtr ColumnFileTiny::writeColumnFile(
 }
 
 PageIdU64 ColumnFileTiny::writeColumnFileData(
-    const DMContext & context,
+    const DMContext & dm_context,
     const Block & block,
     size_t offset,
     size_t limit,
     WriteBatches & wbs)
 {
-    auto page_id = context.storage_pool->newLogPageId();
+    auto page_id = dm_context.storage_pool->newLogPageId();
 
     MemoryWriteBuffer write_buf;
     PageFieldSizes col_data_sizes;
@@ -296,8 +296,8 @@ PageIdU64 ColumnFileTiny::writeColumnFileData(
             col.type,
             offset,
             limit,
-            context.db_context.getSettingsRef().dt_compression_method,
-            context.db_context.getSettingsRef().dt_compression_level);
+            dm_context.global_context.getSettingsRef().dt_compression_method,
+            dm_context.global_context.getSettingsRef().dt_compression_level);
         size_t serialized_size = write_buf.count() - last_buf_size;
         RUNTIME_CHECK_MSG(
             serialized_size != 0,

dbms/src/Storages/DeltaMerge/DMContext.cpp

@@ -19,11 +19,11 @@ namespace DB::DM
 {
 WriteLimiterPtr DMContext::getWriteLimiter() const
 {
-    return db_context.getWriteLimiter();
+    return global_context.getWriteLimiter();
 }
 ReadLimiterPtr DMContext::getReadLimiter() const
 {
-    return db_context.getReadLimiter();
+    return global_context.getReadLimiter();
 }
 
 } // namespace DB::DM