From 20281ad988e05865752418d4619ca8408e6e22e3 Mon Sep 17 00:00:00 2001
From: Kirill Pleshivtsev <drbasic@yandex-team.ru>
Date: Mon, 20 Jan 2025 18:53:34 +0700
Subject: [PATCH] Configs for direct copy throttling

---
 cloud/blockstore/config/disk.proto            | 19 ++++++
 cloud/blockstore/config/storage.proto         |  4 --
 .../blockstore/libs/disk_agent/bootstrap.cpp  |  1 +
 .../libs/disk_agent/config_initializer.cpp    | 58 ++++++++++++++++---
 .../libs/disk_agent/config_initializer.h      |  5 +-
 cloud/blockstore/libs/storage/core/config.cpp |  1 -
 cloud/blockstore/libs/storage/core/config.h   |  1 -
 .../storage/disk_agent/disk_agent_actor.cpp   |  7 ---
 .../storage/disk_agent/disk_agent_actor.h     |  2 +-
 .../disk_agent/model/bandwidth_calculator.cpp | 24 +++++++-
 .../disk_agent/model/bandwidth_calculator.h   |  6 +-
 .../libs/storage/disk_agent/model/config.h    |  5 ++
 .../disk_agent/spdk_initializer_ut.cpp        | 10 +++-
 ...nrepl_migration_common_actor_migration.cpp |  4 +-
 example/0-setup.sh                            | 14 +++++
 example/nbs/nbs-throttling.json               | 10 ++++
 16 files changed, 139 insertions(+), 32 deletions(-)
 create mode 100644 example/nbs/nbs-throttling.json

diff --git a/cloud/blockstore/config/disk.proto b/cloud/blockstore/config/disk.proto
index ed488ec2c3d..8b715760fa9 100644
--- a/cloud/blockstore/config/disk.proto
+++ b/cloud/blockstore/config/disk.proto
@@ -150,6 +150,22 @@ message TStorageDiscoveryConfig
 
 ////////////////////////////////////////////////////////////////////////////////
 
+message TDiskAgentThrottlingConfig
+{
+    // Host limits.
+    optional string InfraThrottlingConfigPath = 1;
+    optional uint32 DefaultNetworkMbitThroughput = 2;
+
+    // Fraction of network throughput utilized for migrations and shadow
+    // disk fill.
+    optional double DirectCopyBandwidthFraction = 3;
+
+    // Maximum bandwidth for one device in MiB/s.
+    optional uint64 MaxDeviceBandwidthMiB = 4;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 message TDiskAgentConfig
 {
     optional bool Enabled = 1;
@@ -275,6 +291,9 @@ message TDiskAgentConfig
     }
 
     repeated TPathToSerialNumber PathToSerialNumberMapping = 37;
+
+    // Settings for traffic shaping.
+    optional TDiskAgentThrottlingConfig ThrottlingConfig = 38;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/cloud/blockstore/config/storage.proto b/cloud/blockstore/config/storage.proto
index d1a6fe2fe81..e032170f956 100644
--- a/cloud/blockstore/config/storage.proto
+++ b/cloud/blockstore/config/storage.proto
@@ -1089,8 +1089,4 @@ message TStorageServiceConfig
 
     // Enable buttons for device state changing, when they in error state.
     optional bool EnableToChangeErrorStatesFromDiskRegistryMonpage = 398;
-
-    // Fraction of network throughput utilized for migrations and shadow
-    // disk fill.
-    optional double DirectCopyBandwidthFraction = 399;
 }
diff --git a/cloud/blockstore/libs/disk_agent/bootstrap.cpp b/cloud/blockstore/libs/disk_agent/bootstrap.cpp
index d2b3ca7ad1e..5c2c61a38f9 100644
--- a/cloud/blockstore/libs/disk_agent/bootstrap.cpp
+++ b/cloud/blockstore/libs/disk_agent/bootstrap.cpp
@@ -346,6 +346,7 @@ void TBootstrap::InitRdmaServer(NRdma::TRdmaConfig& config)
 
 bool TBootstrap::InitKikimrService()
 {
+    Configs->Log = Log;
     Configs->InitKikimrConfig();
     Configs->InitServerConfig();
     Configs->InitFeaturesConfig();
diff --git a/cloud/blockstore/libs/disk_agent/config_initializer.cpp b/cloud/blockstore/libs/disk_agent/config_initializer.cpp
index c833944c249..326053df34b 100644
--- a/cloud/blockstore/libs/disk_agent/config_initializer.cpp
+++ b/cloud/blockstore/libs/disk_agent/config_initializer.cpp
@@ -8,12 +8,12 @@
 #include <cloud/blockstore/libs/spdk/iface/config.h>
 #include <cloud/blockstore/libs/storage/core/config.h>
 #include <cloud/blockstore/libs/storage/disk_registry_proxy/model/config.h>
-
 #include <cloud/storage/core/libs/common/proto_helpers.h>
 #include <cloud/storage/core/libs/features/features_config.h>
 #include <cloud/storage/core/libs/kikimr/actorsystem.h>
 #include <cloud/storage/core/libs/version/version.h>
 
+#include <library/cpp/json/json_reader.h>
 #include <library/cpp/protobuf/util/pb_io.h>
 
 #include <util/datetime/base.h>
@@ -23,8 +23,54 @@
 
 namespace NCloud::NBlockStore::NServer {
 
+namespace {
+
 ////////////////////////////////////////////////////////////////////////////////
 
+std::optional<NJson::TJsonValue> ReadJsonFile(
+    TLog& Log,
+    const TString& filename)
+{
+    if (filename.empty()) {
+        return {};
+    }
+
+    try {
+        TFileInput in(filename);
+        return NJson::ReadJsonTree(&in, true);
+    } catch (...) {
+        STORAGE_ERROR(
+            "Failed to read file: " << filename.Quote() << " with error: "
+                                    << CurrentExceptionMessage().c_str());
+        return {};
+    }
+}
+
+ui32 ReadNetworkMbitThroughput(
+    TLog& Log,
+    const NProto::TDiskAgentConfig& diskAgentConfig)
+{
+    const auto& config = diskAgentConfig.GetThrottlingConfig();
+    ui32 networkThroughput = config.GetDefaultNetworkMbitThroughput();
+
+    if (auto json = ReadJsonFile(Log, config.GetInfraThrottlingConfigPath())) {
+        try {
+            if (auto* value = json->GetValueByPath("interfaces.[0].eth0.speed"))
+            {
+                networkThroughput = FromString<ui64>(value->GetStringSafe());
+            }
+        } catch (...) {
+            STORAGE_ERROR(
+                "Failed to read NetworkMbitThroughput. Error: "
+                << CurrentExceptionMessage().c_str());
+        }
+    }
+
+    return networkThroughput;
+}
+
+}   // namespace
+
 void TConfigInitializer::ApplyCMSConfigs(NKikimrConfig::TAppConfig cmsConfig)
 {
     if (cmsConfig.HasBlobStorageConfig()) {
@@ -167,10 +213,11 @@ void TConfigInitializer::InitDiskAgentConfig()
     SetupDiskAgentConfig(diskAgentConfig);
     ApplySpdkEnvConfig(diskAgentConfig.GetSpdkEnvConfig());
 
+    const ui32 networkMbitThroughput = ReadNetworkMbitThroughput(Log, diskAgentConfig);
     DiskAgentConfig = std::make_shared<NStorage::TDiskAgentConfig>(
         std::move(diskAgentConfig),
         Rack,
-        NetworkMbitThroughput);
+        networkMbitThroughput);
 }
 
 void TConfigInitializer::InitDiskRegistryProxyConfig()
@@ -184,10 +231,6 @@ void TConfigInitializer::InitDiskRegistryProxyConfig()
         std::move(config));
 }
 
-void TConfigInitializer::InitNetworkThroughput() {
-
-}
-
 void TConfigInitializer::InitServerConfig()
 {
     NProto::TServerAppConfig appConfig;
@@ -432,10 +475,11 @@ void TConfigInitializer::ApplyDiskAgentConfig(const TString& text)
             DiskAgentConfig->GetStorageDiscoveryConfig());
     }
 
+    const ui32 networkMbitThroughput = ReadNetworkMbitThroughput(Log, config);
     DiskAgentConfig = std::make_shared<NStorage::TDiskAgentConfig>(
         std::move(config),
         Rack,
-        NetworkMbitThroughput);
+        networkMbitThroughput);
 }
 
 void TConfigInitializer::ApplyDiskRegistryProxyConfig(const TString& text)
diff --git a/cloud/blockstore/libs/disk_agent/config_initializer.h b/cloud/blockstore/libs/disk_agent/config_initializer.h
index 40a72204a5a..ba71082f702 100644
--- a/cloud/blockstore/libs/disk_agent/config_initializer.h
+++ b/cloud/blockstore/libs/disk_agent/config_initializer.h
@@ -19,6 +19,8 @@
 #include <contrib/ydb/core/protos/blobstorage.pb.h>
 #include <contrib/ydb/core/protos/config.pb.h>
 
+#include <library/cpp/logger/log.h>
+
 namespace NCloud::NBlockStore::NServer {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -35,10 +37,10 @@ struct TConfigInitializer
     TDiagnosticsConfigPtr DiagnosticsConfig;
     NSpdk::TSpdkEnvConfigPtr SpdkEnvConfig;
     NFeatures::TFeaturesConfigPtr FeaturesConfig;
-    ui32 NetworkMbitThroughput = 0;
     NRdma::TRdmaConfigPtr RdmaConfig;
 
     TString Rack;
+    TLog Log;
 
     TConfigInitializer(TOptionsPtr options)
         : Options(std::move(options))
@@ -52,7 +54,6 @@ struct TConfigInitializer
     void InitStorageConfig();
     void InitDiskAgentConfig();
     void InitDiskRegistryProxyConfig();
-    void InitNetworkThroughput();
     void InitServerConfig();
     void InitSpdkEnvConfig();
     void InitFeaturesConfig();
diff --git a/cloud/blockstore/libs/storage/core/config.cpp b/cloud/blockstore/libs/storage/core/config.cpp
index 576a9ec04f0..fddceb9b504 100644
--- a/cloud/blockstore/libs/storage/core/config.cpp
+++ b/cloud/blockstore/libs/storage/core/config.cpp
@@ -491,7 +491,6 @@ TDuration MSeconds(ui32 value)
     xxx(VolumeProxyCacheRetryDuration,             TDuration, Seconds(15)     )\
                                                                                \
     xxx(UseDirectCopyRange,                             bool,      false         )\
-    xxx(DirectCopyBandwidthFraction,                    double,    0.5           )\
     xxx(MaxShadowDiskFillBandwidth,                     ui32,      512           )\
     xxx(MaxShadowDiskFillIoDepth,                       ui32,      1             )\
     xxx(BackgroundOperationsTotalBandwidth,             ui32,      1024          )\
diff --git a/cloud/blockstore/libs/storage/core/config.h b/cloud/blockstore/libs/storage/core/config.h
index b77f04cbcdf..56e10d2a106 100644
--- a/cloud/blockstore/libs/storage/core/config.h
+++ b/cloud/blockstore/libs/storage/core/config.h
@@ -570,7 +570,6 @@ class TStorageConfig
     TString GetCachedDiskAgentSessionsPath() const;
 
     bool GetUseDirectCopyRange() const;
-    double GetDirectCopyBandwidthFraction() const;
     ui32 GetMaxShadowDiskFillBandwidth() const;
     ui32 GetMaxShadowDiskFillIoDepth() const;
     ui32 GetBackgroundOperationsTotalBandwidth() const;
diff --git a/cloud/blockstore/libs/storage/disk_agent/disk_agent_actor.cpp b/cloud/blockstore/libs/storage/disk_agent/disk_agent_actor.cpp
index e255793009c..3655a29a887 100644
--- a/cloud/blockstore/libs/storage/disk_agent/disk_agent_actor.cpp
+++ b/cloud/blockstore/libs/storage/disk_agent/disk_agent_actor.cpp
@@ -21,10 +21,6 @@ namespace {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-ui64 GetNetworkBandwidth(ui32 networkMbitThroughput, double directCopyBandwidthFraction) {
-    return (static_cast<ui64>(networkMbitThroughput) * 1_MB / 8) *
-           directCopyBandwidthFraction;
-}
 
 }   // namespace
 
@@ -51,9 +47,6 @@ TDiskAgentActor::TDiskAgentActor(
     , Logging(std::move(logging))
     , RdmaServer(std::move(rdmaServer))
     , NvmeManager(std::move(nvmeManager))
-    , BandwidthCalculator(
-          GetNetworkBandwidth(AgentConfig->GetNetworkMbitThroughput(),
-          Config->GetDirectCopyBandwidthFraction()))
 {}
 
 TDiskAgentActor::~TDiskAgentActor()
diff --git a/cloud/blockstore/libs/storage/disk_agent/disk_agent_actor.h b/cloud/blockstore/libs/storage/disk_agent/disk_agent_actor.h
index 9db2533f689..5854307216c 100644
--- a/cloud/blockstore/libs/storage/disk_agent/disk_agent_actor.h
+++ b/cloud/blockstore/libs/storage/disk_agent/disk_agent_actor.h
@@ -72,7 +72,7 @@ class TDiskAgentActor final
     // Pending WaitReady requests
     TDeque<TPendingRequest> PendingRequests;
 
-    TBandwidthCalculator BandwidthCalculator;
+    TBandwidthCalculator BandwidthCalculator {*AgentConfig};
 
     ERegistrationState RegistrationState = ERegistrationState::NotStarted;
 
diff --git a/cloud/blockstore/libs/storage/disk_agent/model/bandwidth_calculator.cpp b/cloud/blockstore/libs/storage/disk_agent/model/bandwidth_calculator.cpp
index 4d440575a54..a8b90c2598f 100644
--- a/cloud/blockstore/libs/storage/disk_agent/model/bandwidth_calculator.cpp
+++ b/cloud/blockstore/libs/storage/disk_agent/model/bandwidth_calculator.cpp
@@ -1,13 +1,27 @@
 #include "bandwidth_calculator.h"
 
+#include <cloud/blockstore/libs/storage/disk_agent/model/config.h>
+
 namespace NCloud::NBlockStore::NStorage {
 
 ////////////////////////////////////////////////////////////////////////////////
 
+namespace {
+
 constexpr auto WindowDuration = TDuration::Seconds(1);
 
-TBandwidthCalculator::TBandwidthCalculator(ui64 maxTotalBandwidth)
-    : MaxTotalBandwidth(maxTotalBandwidth)
+ui64 GetNetworkBandwidth(const TDiskAgentConfig& config)
+{
+    return (static_cast<ui64>(config.GetNetworkMbitThroughput()) * 1_MB / 8) *
+           config.GetThrottlerConfig().GetDirectCopyBandwidthFraction();
+}
+
+}   // namespace
+
+TBandwidthCalculator::TBandwidthCalculator(const TDiskAgentConfig& config)
+    : MaxTotalBandwidth(GetNetworkBandwidth(config))
+    , MaxDeviceBandwidth(
+          config.GetThrottlerConfig().GetMaxDeviceBandwidthMiB() * 1_MB)
 {}
 
 TBandwidthCalculator::~TBandwidthCalculator() = default;
@@ -34,7 +48,11 @@ void TBandwidthCalculator::ClearHistory(TInstant deadline)
 
 ui64 TBandwidthCalculator::GetRecommendedBandwidth() const
 {
-    return MaxTotalBandwidth / DeviceLastRequest.size();
+    ui64 result = MaxTotalBandwidth / DeviceLastRequest.size();
+    if (MaxDeviceBandwidth) {
+        result = Min(MaxDeviceBandwidth, result);
+    }
+    return result;
 }
 
 }   // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/disk_agent/model/bandwidth_calculator.h b/cloud/blockstore/libs/storage/disk_agent/model/bandwidth_calculator.h
index 24142940b8f..c2c084f48b4 100644
--- a/cloud/blockstore/libs/storage/disk_agent/model/bandwidth_calculator.h
+++ b/cloud/blockstore/libs/storage/disk_agent/model/bandwidth_calculator.h
@@ -1,5 +1,7 @@
 #pragma once
 
+#include <cloud/blockstore/libs/storage/disk_agent/model/public.h>
+
 #include <util/datetime/base.h>
 #include <util/generic/hash.h>
 #include <util/generic/list.h>
@@ -11,10 +13,12 @@ namespace NCloud::NBlockStore::NStorage {
 class TBandwidthCalculator
 {
     const ui64 MaxTotalBandwidth;
+    const ui64 MaxDeviceBandwidth;
+
     THashMap<TString, TInstant> DeviceLastRequest;
 
 public:
-    explicit TBandwidthCalculator(ui64 maxTotalBandwidth);
+    explicit TBandwidthCalculator(const TDiskAgentConfig& config);
     ~TBandwidthCalculator();
 
     // Returns the recommended bandwidth for the next request.
diff --git a/cloud/blockstore/libs/storage/disk_agent/model/config.h b/cloud/blockstore/libs/storage/disk_agent/model/config.h
index 87b490fd9ff..8a7f7c7c18e 100644
--- a/cloud/blockstore/libs/storage/disk_agent/model/config.h
+++ b/cloud/blockstore/libs/storage/disk_agent/model/config.h
@@ -123,6 +123,11 @@ class TDiskAgentConfig
         return Config.GetPathToSerialNumberMapping();
     }
 
+    const NProto::TDiskAgentThrottlingConfig& GetThrottlerConfig() const
+    {
+        return Config.GetThrottlingConfig();
+    }
+
     ui32 GetNetworkMbitThroughput() const
     {
         return NetworkMbitThroughput;
diff --git a/cloud/blockstore/libs/storage/disk_agent/spdk_initializer_ut.cpp b/cloud/blockstore/libs/storage/disk_agent/spdk_initializer_ut.cpp
index c561b6d6a57..7e9f38dc475 100644
--- a/cloud/blockstore/libs/storage/disk_agent/spdk_initializer_ut.cpp
+++ b/cloud/blockstore/libs/storage/disk_agent/spdk_initializer_ut.cpp
@@ -122,9 +122,13 @@ TInitializeSpdkResult InitializeSpdkSync(
     NProto::TDiskAgentConfig config)
 {
     return InitializeSpdk(
-        std::make_shared<NStorage::TDiskAgentConfig>(std::move(config), "rack"),
-        std::make_shared<TTestEnv>(),
-        ICachingAllocatorPtr()).GetValueSync();
+               std::make_shared<NStorage::TDiskAgentConfig>(
+                   std::move(config),
+                   "rack",
+                   25000),
+               std::make_shared<TTestEnv>(),
+               ICachingAllocatorPtr())
+        .GetValueSync();
 }
 
 }   // namespace
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor_migration.cpp b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor_migration.cpp
index fe14f512ddf..3504cd82600 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor_migration.cpp
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor_migration.cpp
@@ -308,10 +308,10 @@ void TNonreplicatedPartitionMigrationCommonActor::HandleRangeMigrated(
     LOG_DEBUG(
         ctx,
         TBlockStoreComponents::PARTITION,
-        "[%s] Range %s migrated. Recommended bandwidth: %lu",
+        "[%s] Range %s migrated. Recommended bandwidth: %.2f MiB",
         DiskId.c_str(),
         DescribeRange(msg->Range).c_str(),
-        msg->RecommendedBandwidth);
+        static_cast<double>(msg->RecommendedBandwidth) / 1_MB);
 
     if (msg->AllZeroes) {
         ChangedRangesMap.MarkNotChanged(msg->Range);
diff --git a/example/0-setup.sh b/example/0-setup.sh
index 3b67a81a6d5..09971cb5700 100755
--- a/example/0-setup.sh
+++ b/example/0-setup.sh
@@ -77,6 +77,13 @@ FileDevices: {
     BlockSize: 4096
 }
 
+ThrottlingConfig: {
+    InfraThrottlingConfigPath: "nbs/nbs-throttling.json"
+    DefaultNetworkMbitThroughput: 100
+    DirectCopyBandwidthFraction: 0.5
+    MaxDeviceBandwidthMiB: 15
+}
+
 EOF
 
 cat > $BIN_DIR/nbs/nbs-location-$1.txt <<EOF
@@ -118,6 +125,13 @@ FileDevices: {
     BlockSize: 512
     PoolName: "local-ssd"
 }
+
+ThrottlingConfig: {
+    InfraThrottlingConfigPath: "nbs/nbs-throttling.json"
+    DefaultNetworkMbitThroughput: 100
+    DirectCopyBandwidthFraction: 0.5
+}
+
 EOF
 
 cat > $BIN_DIR/nbs/nbs-location-0.txt <<EOF
diff --git a/example/nbs/nbs-throttling.json b/example/nbs/nbs-throttling.json
new file mode 100644
index 00000000000..7a79230a763
--- /dev/null
+++ b/example/nbs/nbs-throttling.json
@@ -0,0 +1,10 @@
+{
+    "compute_cores_num": 16,
+    "interfaces": [
+        {
+            "eth0": {
+                "speed": "1000"
+            }
+        }
+    ]
+}