support collect cross AZ network bytes for disagg tiflash (#10249)

close #10221

Signed-off-by: guo-shaoge <shaoge1994@163.com>

Author: guo-shaoge

contrib/client-c

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <Common/Exception.h>
 #include <common/defines.h>
 #include <common/types.h>
 
@@ -31,6 +32,16 @@ struct ConnectionProfileInfo
         InterZoneRemote = 2,
     };
     using ConnTypeVec = std::vector<ConnectionProfileInfo::ConnectionType>;
+
+    static ConnectionProfileInfo createForInterZone()
+    {
+        return ConnectionProfileInfo(ConnectionProfileInfo::ConnectionType::InterZoneRemote);
+    }
+    static ConnectionProfileInfo createForInnerZone()
+    {
+        return ConnectionProfileInfo(ConnectionProfileInfo::ConnectionType::InnerZoneRemote);
+    }
+
     static ALWAYS_INLINE ConnectionType inferConnectionType(bool is_local, bool same_zone)
     {
         if (is_local)
@@ -55,6 +66,12 @@ struct ConnectionProfileInfo
     {}
 
     String getTypeString() const { return String(magic_enum::enum_name(type)); }
+    void merge(const ConnectionProfileInfo & other)
+    {
+        RUNTIME_CHECK(type == other.type, magic_enum::enum_name(type), magic_enum::enum_name(other.type));
+        packets += other.packets;
+        bytes += other.bytes;
+    }
 
     Int64 packets = 0;
     Int64 bytes = 0;

dbms/src/Flash/Statistics/ConnectionProfileInfo.h

@@ -16,6 +16,8 @@
 #include <Flash/Statistics/ConnectionProfileInfo.h>
 #include <Flash/Statistics/TableScanImpl.h>
 #include <Interpreters/Join.h>
+#include <Storages/DeltaMerge/ReadThread/UnorderedInputStream.h>
+#include <Storages/DeltaMerge/Remote/RNSegmentInputStream.h>
 #include <Storages/DeltaMerge/ScanContext.h>
 
 namespace DB
@@ -68,12 +70,22 @@ void TableScanStatistics::updateTableScanDetail(const std::vector<ConnectionProf
             remote_table_scan_detail.inter_zone_conn_profile_info.packets += connection_profile_info.packets;
             remote_table_scan_detail.inter_zone_conn_profile_info.bytes += connection_profile_info.bytes;
         }
-        // Stats both send and receive bytes here, because remote execution summaries are not used now
+        // Whether it's for remote read (non-disaggregated mode) or communication between CN and WN (disaggregated mode),
+        // send bytes are always collected on the node that sends the request (i.e., the node where the table scan executor resides).
+        // Therefore, both send bytes and receive bytes are collected here.
         base.updateSendConnectionInfo(connection_profile_info);
         base.updateReceiveConnectionInfo(connection_profile_info);
     }
 }
 
+void TableScanStatistics::updateTableScanDetailForDisaggIfNecessary(const IProfilingBlockInputStream * stream)
+{
+    if (const auto * unordered_stream = dynamic_cast<const DM::UnorderedInputStream *>(stream))
+        updateTableScanDetail(unordered_stream->getConnectionProfileInfos());
+    else if (const auto * rn_segment_stream = dynamic_cast<const DM::Remote::RNSegmentInputStream *>(stream))
+        updateTableScanDetail(rn_segment_stream->getConnectionProfileInfos());
+}
+
 void TableScanStatistics::collectExtraRuntimeDetail()
 {
     switch (dag_context.getExecutionMode())
@@ -91,6 +103,8 @@ void TableScanStatistics::collectExtraRuntimeDetail()
             else if (const auto * local_stream = dynamic_cast<const IProfilingBlockInputStream *>(&stream);
                      local_stream)
             {
+                updateTableScanDetailForDisaggIfNecessary(local_stream);
+
                 /// local read input stream also is IProfilingBlockInputStream
                 const auto & prof = local_stream->getProfileInfo();
                 local_table_scan_detail.bytes += prof.bytes;

dbms/src/Flash/Statistics/TableScanImpl.cpp

@@ -70,5 +70,6 @@ class TableScanStatistics : public TableScanStatisticsBase
 
 private:
     void updateTableScanDetail(const std::vector<ConnectionProfileInfo> & connection_profile_infos);
+    void updateTableScanDetailForDisaggIfNecessary(const IProfilingBlockInputStream * stream);
 };
 } // namespace DB

dbms/src/Flash/Statistics/TableScanImpl.h

@@ -26,7 +26,8 @@ UnorderedSourceOp::UnorderedSourceOp(
     int extra_table_id_index_,
     const String & req_id,
     const RuntimeFilteList & runtime_filter_list_,
-    int max_wait_time_ms_)
+    int max_wait_time_ms_,
+    bool is_disagg_)
     : SourceOp(exec_context_, req_id)
     , task_pool(task_pool_)
     , ref_no(0)
@@ -35,6 +36,17 @@ UnorderedSourceOp::UnorderedSourceOp(
 {
     setHeader(AddExtraTableIDColumnTransformAction::buildHeader(columns_to_read_, extra_table_id_index_));
     ref_no = task_pool->increaseUnorderedInputStreamRefCount();
+
+    if (is_disagg_)
+    {
+        // One connection for inter zone, the other one for inner zone.
+        const size_t connections = 2;
+        io_profile_info = IOProfileInfo::createForRemote(profile_info_ptr, connections);
+    }
+    else
+    {
+        io_profile_info = IOProfileInfo::createForLocal(profile_info_ptr);
+    }
 }
 
 OperatorStatus UnorderedSourceOp::readImpl(Block & block)
@@ -98,4 +110,29 @@ void UnorderedSourceOp::operatePrefixImpl()
         }
     });
 }
+
+void UnorderedSourceOp::operateSuffixImpl()
+{
+    // If io_profile_info is local, it means this table scan only read local data.
+    // So no need to update connection info, because connection info indicate it's a remote table scan,
+    // like CoprocessorReader or read delta data from WN.
+    if (!io_profile_info->is_local)
+    {
+        std::call_once(task_pool->getRemoteConnectionInfoFlag(), [&]() {
+            auto & connection_infos = io_profile_info->connection_profile_infos;
+            RUNTIME_CHECK(connection_infos.size() == 2, connection_infos.size());
+
+            auto pool_connection_info_opt = task_pool->getRemoteConnectionInfo();
+            RUNTIME_CHECK(
+                pool_connection_info_opt.has_value() || (task_pool->getTotalReadTasks() == 0),
+                pool_connection_info_opt.has_value(),
+                task_pool->getTotalReadTasks());
+            if (pool_connection_info_opt)
+            {
+                connection_infos[0] = pool_connection_info_opt->first;
+                connection_infos[1] = pool_connection_info_opt->second;
+            }
+        });
+    }
+}
 } // namespace DB

dbms/src/Operators/UnorderedSourceOp.cpp

@@ -34,7 +34,8 @@ class UnorderedSourceOp : public SourceOp
         int extra_table_id_index_,
         const String & req_id,
         const RuntimeFilteList & runtime_filter_list_ = std::vector<RuntimeFilterPtr>{},
-        int max_wait_time_ms_ = 0);
+        int max_wait_time_ms_ = 0,
+        bool is_disagg_ = false);
 
     ~UnorderedSourceOp() override
     {
@@ -50,7 +51,7 @@ class UnorderedSourceOp : public SourceOp
 
     String getName() const override { return "UnorderedSourceOp"; }
 
-    IOProfileInfoPtr getIOProfileInfo() const override { return IOProfileInfo::createForLocal(profile_info_ptr); }
+    IOProfileInfoPtr getIOProfileInfo() const override { return io_profile_info; }
 
     // only for unit test
     // The logic order of unit test is error, it will build source_op firstly and register rf secondly.
@@ -64,6 +65,7 @@ class UnorderedSourceOp : public SourceOp
 
 protected:
     void operatePrefixImpl() override;
+    void operateSuffixImpl() override;
 
     OperatorStatus readImpl(Block & block) override;
 
@@ -76,5 +78,6 @@ class UnorderedSourceOp : public SourceOp
     int max_wait_time_ms;
 
     bool done = false;
+    IOProfileInfoPtr io_profile_info;
 };
 } // namespace DB

dbms/src/Operators/UnorderedSourceOp.h

@@ -37,14 +37,16 @@ class UnorderedInputStream : public IProfilingBlockInputStream
         int extra_table_id_index_,
         const String & req_id,
         const RuntimeFilteList & runtime_filter_list_ = std::vector<RuntimeFilterPtr>{},
-        int max_wait_time_ms_ = 0)
+        int max_wait_time_ms_ = 0,
+        bool is_disagg_ = false)
         : task_pool(task_pool_)
         , header(AddExtraTableIDColumnTransformAction::buildHeader(columns_to_read_, extra_table_id_index_))
         , log(Logger::get(req_id))
         , ref_no(0)
         , task_pool_added(false)
         , runtime_filter_list(runtime_filter_list_)
         , max_wait_time_ms(max_wait_time_ms_)
+        , is_disagg(is_disagg_)
     {
         ref_no = task_pool->increaseUnorderedInputStreamRefCount();
     }
@@ -75,6 +77,8 @@ class UnorderedInputStream : public IProfilingBlockInputStream
         max_wait_time_ms = max_wait_time_ms_;
     }
 
+    std::vector<ConnectionProfileInfo> getConnectionProfileInfos() const { return connection_profile_infos; }
+
 protected:
     Block readImpl() override
     {
@@ -116,6 +120,23 @@ class UnorderedInputStream : public IProfilingBlockInputStream
 
     void readSuffixImpl() override
     {
+        if (is_disagg)
+        {
+            std::call_once(task_pool->getRemoteConnectionInfoFlag(), [&]() {
+                auto pool_connection_info_opt = task_pool->getRemoteConnectionInfo();
+                RUNTIME_CHECK(
+                    pool_connection_info_opt.has_value() || (task_pool->getTotalReadTasks() == 0),
+                    pool_connection_info_opt.has_value(),
+                    task_pool->getTotalReadTasks());
+                RUNTIME_CHECK(connection_profile_infos.empty(), connection_profile_infos.size());
+
+                if (pool_connection_info_opt)
+                {
+                    connection_profile_infos.push_back(pool_connection_info_opt->first);
+                    connection_profile_infos.push_back(pool_connection_info_opt->second);
+                }
+            });
+        }
         LOG_DEBUG(
             log,
             "Finish read from storage, pool_id={} ref_no={} rows={}",
@@ -156,6 +177,8 @@ class UnorderedInputStream : public IProfilingBlockInputStream
     std::vector<RuntimeFilterPtr> runtime_filter_list;
     int max_wait_time_ms;
 
+    std::vector<ConnectionProfileInfo> connection_profile_infos;
+    const bool is_disagg;
     friend class tests::DeltaMergeStoreRWTest;
 };
 } // namespace DB::DM

dbms/src/Storages/DeltaMerge/ReadThread/UnorderedInputStream.h

@@ -89,6 +89,20 @@ Block RNSegmentInputStream::readImpl(FilterPtr & res_filter, bool return_filter)
 
         if (!res)
         {
+            static constexpr auto inter_type = ConnectionProfileInfo::ConnectionType::InterZoneRemote;
+            static constexpr auto inner_type = ConnectionProfileInfo::ConnectionType::InnerZoneRemote;
+
+            RUNTIME_CHECK(connection_profile_infos.size() == 2, connection_profile_infos.size());
+            RUNTIME_CHECK(current_seg_task->extra_remote_info.has_value());
+            if (current_seg_task->extra_remote_info->connection_profile_info.type == inter_type)
+                connection_profile_infos[INTER_ZONE_INDEX].merge(
+                    current_seg_task->extra_remote_info->connection_profile_info);
+            else if (current_seg_task->extra_remote_info->connection_profile_info.type == inner_type)
+                connection_profile_infos[INNER_ZONE_INDEX].merge(
+                    current_seg_task->extra_remote_info->connection_profile_info);
+            else
+                throw Exception("unexpected connection type");
+
             // Current stream is drained, try read from next stream.
             current_seg_task = nullptr;
             continue;

dbms/src/Storages/DeltaMerge/Remote/RNSegmentInputStream.cpp

@@ -34,6 +34,8 @@ class RNSegmentInputStream : public IProfilingBlockInputStream
 
     Block getHeader() const override { return action.getHeader(); }
 
+    std::vector<ConnectionProfileInfo> getConnectionProfileInfos() const { return connection_profile_infos; }
+
 protected:
     Block readImpl() override
     {
@@ -56,7 +58,11 @@ class RNSegmentInputStream : public IProfilingBlockInputStream
         : log(Logger::get(options.debug_tag))
         , workers(options.workers)
         , action(options.columns_to_read, options.extra_table_id_index)
-    {}
+    {
+        connection_profile_infos.resize(2);
+        connection_profile_infos[INTER_ZONE_INDEX] = ConnectionProfileInfo::createForInterZone();
+        connection_profile_infos[INNER_ZONE_INDEX] = ConnectionProfileInfo::createForInnerZone();
+    }
 
     static BlockInputStreamPtr create(const Options & options)
     {
@@ -74,6 +80,10 @@ class RNSegmentInputStream : public IProfilingBlockInputStream
 
     double duration_wait_ready_task_sec = 0;
     double duration_read_sec = 0;
+
+    static constexpr size_t INTER_ZONE_INDEX = 0;
+    static constexpr size_t INNER_ZONE_INDEX = 1;
+    std::vector<ConnectionProfileInfo> connection_profile_infos;
 };
 
 } // namespace DB::DM::Remote

dbms/src/Storages/DeltaMerge/Remote/RNSegmentInputStream.h

@@ -137,6 +137,21 @@ OperatorStatus RNSegmentSourceOp::executeIOImpl()
     }
     else
     {
+        static constexpr auto inter_type = ConnectionProfileInfo::ConnectionType::InterZoneRemote;
+        static constexpr auto inner_type = ConnectionProfileInfo::ConnectionType::InnerZoneRemote;
+
+        auto & connection_profile_infos = io_profile_info->connection_profile_infos;
+        RUNTIME_CHECK(connection_profile_infos.size() == 2, connection_profile_infos.size());
+        RUNTIME_CHECK(current_seg_task->extra_remote_info.has_value());
+        if (current_seg_task->extra_remote_info->connection_profile_info.type == inter_type)
+            connection_profile_infos[INTER_ZONE_INDEX].merge(
+                current_seg_task->extra_remote_info->connection_profile_info);
+        else if (current_seg_task->extra_remote_info->connection_profile_info.type == inner_type)
+            connection_profile_infos[INNER_ZONE_INDEX].merge(
+                current_seg_task->extra_remote_info->connection_profile_info);
+        else
+            throw Exception("unexpected connection type");
+
         // Current stream is drained, try to get next ready task.
         current_seg_task = nullptr;
         return startGettingNextReadyTask();