Storage Balancer KIKIMR-20636

ydb/core/mind/hive/balancer.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "hive_impl.h"
+#include "leader_tablet_info.h"
 
 namespace NKikimr {
 namespace NHive {
@@ -11,5 +12,8 @@ void BalanceNodes(std::vector<TNodeInfo*>& nodes, EResourceToBalance resourceTob
 template<NKikimrConfig::THiveConfig::EHiveTabletBalanceStrategy EHiveTabletBalanceStrategy>
 void BalanceTablets(std::vector<TTabletInfo*>& tablets, EResourceToBalance resourceToBalance);
 
+template <NKikimrConfig::THiveConfig::EHiveChannelBalanceStrategy>
+void BalanceChannels(std::vector<TLeaderTabletInfo::TChannel>& channels, NKikimrConfig::THiveConfig::EHiveStorageBalanceStrategy metricToBalance);
+
 }
 }

ydb/core/mind/hive/hive.cpp

@@ -38,6 +38,7 @@ TString EBalancerTypeName(EBalancerType value) {
         case EBalancerType::Emergency: return "Emergency";
         case EBalancerType::SpreadNeighbours: return "Spread";
         case EBalancerType::Manual: return "Manual";
+        case EBalancerType::Storage: return "Storage";
     }
 }
 

ydb/core/mind/hive/hive.h

@@ -85,8 +85,9 @@ enum class EBalancerType {
     ScatterNetwork,
     Emergency,
     SpreadNeighbours,
+    Storage,
 
-    Last = SpreadNeighbours,
+    Last = Storage,
 };
 
 constexpr std::size_t EBalancerTypeSize = static_cast<std::size_t>(EBalancerType::Last) + 1;
@@ -261,6 +262,12 @@ struct TBalancerSettings {
     std::optional<TFullObjectId> FilterObjectId;
 };
 
+struct TStorageBalancerSettings {
+    ui64 NumReassigns;
+    ui64 MaxInFlight;
+    TString StoragePool;
+};
+
 struct TBalancerStats {
     ui64 TotalRuns = 0;
     ui64 TotalMovements = 0;

ydb/core/mind/hive/hive_events.h

@@ -27,6 +27,8 @@ struct TEvPrivate {
         EvProcessIncomingEvent,
         EvRefreshStorageInfo,
         EvLogTabletMoves,
+        EvStartStorageBalancer,
+        EvRestartCancelled,
         EvEnd
     };
 
@@ -90,6 +92,18 @@ struct TEvPrivate {
     struct TEvRefreshStorageInfo : TEventLocal<TEvRefreshStorageInfo, EvRefreshStorageInfo> {};
 
     struct TEvLogTabletMoves : TEventLocal<TEvLogTabletMoves, EvLogTabletMoves> {};
+
+    struct TEvStartStorageBalancer : TEventLocal<TEvStartStorageBalancer, EvStartStorageBalancer> {
+        TStorageBalancerSettings Settings;
+
+        TEvStartStorageBalancer(TStorageBalancerSettings settings) : Settings(settings) {}
+    };
+
+    struct TEvRestartCancelled : TEventLocal<TEvRestartCancelled, EvRestartCancelled> {
+        TFullTabletId TabletId;
+
+        TEvRestartCancelled(TFullTabletId tabletId) : TabletId(tabletId) {}
+    };
 };
 
 } // NHive

ydb/core/mind/hive/hive_impl.cpp

@@ -330,6 +330,12 @@ void THive::Handle(TEvPrivate::TEvBalancerOut::TPtr&) {
     BLOG_D("Handle BalancerOut");
 }
 
+
+void THive::Handle(TEvPrivate::TEvStartStorageBalancer::TPtr& ev) {
+    BLOG_D("Handle StartStorageBalancer");
+    StartHiveStorageBalancer(std::move(ev->Get()->Settings));
+}
+
 void THive::Handle(TEvHive::TEvBootTablet::TPtr& ev) {
     TTabletId tabletId = ev->Get()->Record.GetTabletID();
     TTabletInfo* tablet = FindTablet(tabletId);
@@ -2603,6 +2609,7 @@ TDuration THive::GetBalancerCooldown() const {
         case EBalancerType::ScatterMemory:
         case EBalancerType::ScatterNetwork:
         case EBalancerType::SpreadNeighbours:
+        case EBalancerType::Storage:
             return GetMinPeriodBetweenBalance();
         case EBalancerType::Emergency:
             return GetMinPeriodBetweenEmergencyBalance();
@@ -2813,6 +2820,7 @@ void THive::ProcessEvent(std::unique_ptr<IEventHandle> event) {
         hFunc(TEvHive::TEvUpdateTabletsObject, Handle);
         hFunc(TEvPrivate::TEvRefreshStorageInfo, Handle);
         hFunc(TEvPrivate::TEvLogTabletMoves, Handle);
+        hFunc(TEvPrivate::TEvStartStorageBalancer, Handle);
     }
 }
 
@@ -2910,6 +2918,7 @@ STFUNC(THive::StateWork) {
         fFunc(TEvHive::TEvUpdateTabletsObject::EventType, EnqueueIncomingEvent);
         fFunc(TEvPrivate::TEvRefreshStorageInfo::EventType, EnqueueIncomingEvent);
         fFunc(TEvPrivate::TEvLogTabletMoves::EventType, EnqueueIncomingEvent);
+        fFunc(TEvPrivate::TEvStartStorageBalancer::EventType, EnqueueIncomingEvent);
         hFunc(TEvPrivate::TEvProcessIncomingEvent, Handle);
     default:
         if (!HandleDefaultEvents(ev, SelfId())) {

ydb/core/mind/hive/hive_impl.h

@@ -169,6 +169,7 @@ class THive : public TActor<THive>, public TTabletExecutedFlat, public THiveShar
     friend class TQueryMigrationWaitActor;
     friend class TReleaseTabletsWaitActor;
     friend class TDrainNodeWaitActor;
+    friend class THiveStorageBalancer;;
     friend struct TNodeInfo;
 
     friend class TTxInitScheme;
@@ -204,6 +205,7 @@ class THive : public TActor<THive>, public TTabletExecutedFlat, public THiveShar
     friend class TTxMonEvent_QueryMigration;
     friend class TTxMonEvent_RebalanceFromScratch;
     friend class TTxMonEvent_ObjectStats;
+    friend class TTxMonEvent_StorageRebalance;
     friend class TTxKillNode;
     friend class TTxLoadEverything;
     friend class TTxRestartTablet;
@@ -239,6 +241,7 @@ class THive : public TActor<THive>, public TTabletExecutedFlat, public THiveShar
     void StartHiveBalancer(TBalancerSettings&& settings);
     void StartHiveDrain(TNodeId nodeId, TDrainSettings settings);
     void StartHiveFill(TNodeId nodeId, const TActorId& initiator);
+    void StartHiveStorageBalancer(TStorageBalancerSettings settings);
     void CreateEvMonitoring(NMon::TEvRemoteHttpInfo::TPtr& ev, const TActorContext& ctx);
     NJson::TJsonValue GetBalancerProgressJson();
     ITransaction* CreateDeleteTablet(TEvHive::TEvDeleteTablet::TPtr& ev);
@@ -548,6 +551,7 @@ class THive : public TActor<THive>, public TTabletExecutedFlat, public THiveShar
     void Handle(TEvHive::TEvUpdateTabletsObject::TPtr& ev);
     void Handle(TEvPrivate::TEvRefreshStorageInfo::TPtr& ev);
     void Handle(TEvPrivate::TEvLogTabletMoves::TPtr& ev);
+    void Handle(TEvPrivate::TEvStartStorageBalancer::TPtr& ev);
     void Handle(TEvPrivate::TEvProcessIncomingEvent::TPtr& ev);
 
 protected:
@@ -894,6 +898,14 @@ class THive : public TActor<THive>, public TTabletExecutedFlat, public THiveShar
         return CurrentConfig.GetBootStrategy();
     }
 
+    NKikimrConfig::THiveConfig::EHiveChannelBalanceStrategy GetChannelBalanceStrategy() const {
+        return CurrentConfig.GetChannelBalanceStrategy();
+    }
+
+    ui64 GetMaxChannelHistorySize() const {
+        return CurrentConfig.GetMaxChannelHistorySize();
+    }
+
     static void ActualizeRestartStatistics(google::protobuf::RepeatedField<google::protobuf::uint64>& restartTimestamps, ui64 barrier);
     static bool IsSystemTablet(TTabletTypes::EType type);
 

ydb/core/mind/hive/hive_ut.cpp

@@ -2778,6 +2778,77 @@ Y_UNIT_TEST_SUITE(THiveTest) {
         UNIT_ASSERT_VALUES_EQUAL(getGroup(tabletId), goodGroup);
     }
 
+    Y_UNIT_TEST(TestStorageBalancer) {
+        static constexpr ui64 NUM_TABLETS = 4;
+        TTestBasicRuntime runtime(1, false);
+        Setup(runtime, true, 2, [](TAppPrepare& app) {
+            app.HiveConfig.SetMinPeriodBetweenReassign(0);
+        });
+        const ui64 hiveTablet = MakeDefaultHiveID(0);
+        const ui64 testerTablet = MakeDefaultHiveID(1);
+        CreateTestBootstrapper(runtime, CreateTestTabletInfo(hiveTablet, TTabletTypes::Hive), &CreateDefaultHive);
+
+        TTabletTypes::EType tabletType = TTabletTypes::Dummy;
+        TVector<ui64> tablets;
+        for (ui64 i = 0; i < NUM_TABLETS; ++i) {
+            THolder<TEvHive::TEvCreateTablet> ev(new TEvHive::TEvCreateTablet(testerTablet, 100500 + i, tabletType, BINDED_CHANNELS));
+            ev->Record.SetObjectId(i);
+            ui64 tabletId = SendCreateTestTablet(runtime, hiveTablet, testerTablet, std::move(ev), 0, true);
+            tablets.emplace_back(tabletId);
+            MakeSureTabletIsUp(runtime, tabletId, 0);
+        }
+        ui64 tabletBase = tablets.front();
+
+        TActorId sender = runtime.AllocateEdgeActor();
+        auto getGroup = [&runtime, sender, hiveTablet](ui64 tabletId) {
+            runtime.SendToPipe(hiveTablet, sender, new TEvHive::TEvRequestHiveInfo({
+                .TabletId = tabletId,
+                .ReturnChannelHistory = true,
+            }));
+            TAutoPtr<IEventHandle> handle;
+            TEvHive::TEvResponseHiveInfo* response = runtime.GrabEdgeEventRethrow<TEvHive::TEvResponseHiveInfo>(handle);
+
+            const auto& tablet = response->Record.GetTablets().Get(0);
+            const auto& channel = tablet.GetTabletChannels().Get(0);
+            const auto& history = channel.GetHistory();
+            return history.Get(history.size() - 1).GetGroup();
+        };
+
+        std::unordered_map<ui64, std::vector<ui64>> groupToTablets;
+        for (auto tablet : tablets) {
+            groupToTablets[getGroup(tablet)].push_back(tablet);
+        }
+        ui64 tabletA;
+        ui64 tabletB;
+        for (const auto& [group, tablets] : groupToTablets) {
+            if (tablets.size() >= 2) {
+                tabletA = tablets[0];
+                tabletB = tablets[1];
+            }
+        }
+        TChannelsBindings channels = BINDED_CHANNELS;
+        for (auto& bind : channels) {
+            bind.SetSize(200'000'000);
+        }
+        for (auto tablet : {tabletA, tabletB}) {
+            TAutoPtr<TEvHive::TEvCreateTablet> updateTablet(new TEvHive::TEvCreateTablet(testerTablet, 100500 + (tablet - tabletBase), tabletType, channels));
+            SendCreateTestTablet(runtime, hiveTablet, testerTablet, updateTablet, 0, true);
+        }
+        runtime.SendToPipe(hiveTablet, sender, new NHive::TEvPrivate::TEvStartStorageBalancer({
+            .NumReassigns = 100,
+            .MaxInFlight = 1,
+            .StoragePool = "def1",
+        }));
+
+        {
+            TDispatchOptions options;
+            options.FinalEvents.emplace_back(NHive::TEvPrivate::EvRestartComplete, 4); // should actually be less than 4
+            runtime.DispatchEvents(options, TDuration::Seconds(10));
+        }
+
+        UNIT_ASSERT_VALUES_UNEQUAL(getGroup(tabletA), getGroup(tabletB));
+    }
+
 //    Y_UNIT_TEST(TestCreateTabletAndChangeProfiles) {
 //        TTestBasicRuntime runtime(1, false);
 //        Setup(runtime, true);

ydb/core/mind/hive/leader_tablet_info.cpp

@@ -239,6 +239,12 @@ const NKikimrBlobStorage::TEvControllerSelectGroupsResult::TGroupParameters* TLe
                 });
                 break;
             }
+            case NKikimrHive::TEvReassignTablet::HIVE_REASSIGN_REASON_BALANCE: {
+                return storagePool->FindFreeAllocationUnit([&params](const TStorageGroupInfo& newGroup) -> bool {
+                    return newGroup.IsMatchesParameters(*params);
+                });
+                break;
+            }
             case NKikimrHive::TEvReassignTablet::HIVE_REASSIGN_REASON_SPACE: {
                 NKikimrConfig::THiveConfig::EHiveStorageBalanceStrategy balanceStrategy = Hive.CurrentConfig.GetStorageBalanceStrategy();
                 Hive.CurrentConfig.SetStorageBalanceStrategy(NKikimrConfig::THiveConfig::HIVE_STORAGE_BALANCE_STRATEGY_SIZE);

ydb/core/mind/hive/leader_tablet_info.h

@@ -26,6 +26,25 @@ struct TLeaderTabletInfo : TTabletInfo {
     static TString DEFAULT_STORAGE_POOL_NAME;
 
 public:
+    struct TChannel {
+        TTabletId TabletId;
+        ui32 ChannelId;
+        const TChannelBind* ChannelInfo;
+
+        double GetWeight(NKikimrConfig::THiveConfig::EHiveStorageBalanceStrategy metricToBalance) const {
+            Y_DEBUG_ABORT_UNLESS(ChannelInfo);
+            switch (metricToBalance) {
+                case NKikimrConfig::THiveConfig::HIVE_STORAGE_BALANCE_STRATEGY_IOPS:
+                    return ChannelInfo->GetIOPS();
+                case NKikimrConfig::THiveConfig::HIVE_STORAGE_BALANCE_STRATEGY_THROUGHPUT:
+                    return ChannelInfo->GetThroughput();
+                default:
+                case NKikimrConfig::THiveConfig::HIVE_STORAGE_BALANCE_STRATEGY_SIZE:
+                    return ChannelInfo->GetSize();
+            }
+        }
+    };
+
     TTabletId Id;
     ETabletState State;
     TTabletTypes::EType Type;
@@ -297,6 +316,14 @@ struct TLeaderTabletInfo : TTabletInfo {
         return BoundChannels.size();
     }
 
+    TChannel GetChannel(ui32 channelId) const {
+        TChannel channel{.TabletId = Id, .ChannelId = channelId, .ChannelInfo = nullptr};
+        if (channelId < BoundChannels.size()) {
+            channel.ChannelInfo = &BoundChannels[channelId];
+        }
+        return channel;
+    }
+
     void AcquireAllocationUnits();
     void ReleaseAllocationUnits();
     bool AcquireAllocationUnit(ui32 channelId);

ydb/core/mind/hive/monitoring.cpp

@@ -1401,7 +1401,8 @@ class TTxMonEvent_Landing : public TTransactionBase<THive> {
             EBalancerType::Emergency,
             EBalancerType::SpreadNeighbours,
             EBalancerType::Scatter,
-            EBalancerType::Manual
+            EBalancerType::Manual,
+            EBalancerType::Storage,
         }) {
             int balancer = static_cast<int>(type);
             out << "<tr id='balancer" << balancer << "'><td>" << EBalancerTypeName(type) << "</td><td></td><td></td><td></td><td></td><td></td></tr>";
@@ -2482,6 +2483,32 @@ class TTxMonEvent_Rebalance : public TTransactionBase<THive> {
     }
 };
 
+class TTxMonEvent_StorageRebalance : public TTransactionBase<THive> {
+public:
+    const TActorId Source;
+    TStorageBalancerSettings Settings;
+
+    TTxMonEvent_StorageRebalance(const TActorId& source, NMon::TEvRemoteHttpInfo::TPtr& ev, TSelf* hive)
+        : TBase(hive)
+        , Source(source)
+    {
+        Settings.NumReassigns = FromStringWithDefault(ev->Get()->Cgi().Get("reassigns"), 1000);
+        Settings.MaxInFlight = FromStringWithDefault(ev->Get()->Cgi().Get("inflight"), 1);
+        Settings.StoragePool = ev->Get()->Cgi().Get("pool");
+    }
+
+    TTxType GetTxType() const override { return NHive::TXTYPE_MON_REBALANCE; }
+
+    bool Execute(TTransactionContext&, const TActorContext&) override {
+        Self->StartHiveStorageBalancer(Settings);
+        return true;
+    }
+
+    void Complete(const TActorContext& ctx) override {
+        ctx.Send(Source, new NMon::TEvRemoteJsonInfoRes("{}"));
+    }
+};
+
 class TTxMonEvent_RebalanceFromScratch : public TTransactionBase<THive> {
 public:
     const TActorId Source;
@@ -3913,6 +3940,8 @@ class TTxMonEvent_Storage : public TTransactionBase<THive> {
     }
 };
 
+
+
 void THive::CreateEvMonitoring(NMon::TEvRemoteHttpInfo::TPtr& ev, const TActorContext& ctx) {
     if (!ReadyForConnections) {
         return Execute(new TTxMonEvent_NotReady(ev->Sender, this), ctx);
@@ -4038,6 +4067,9 @@ void THive::CreateEvMonitoring(NMon::TEvRemoteHttpInfo::TPtr& ev, const TActorCo
     if (page == "Storage") {
         return Execute(new TTxMonEvent_Storage(ev->Sender, ev, this), ctx);
     }
+    if (page == "StorageRebalance") {
+        return Execute(new TTxMonEvent_StorageRebalance(ev->Sender, ev, this), ctx);
+    }
     return Execute(new TTxMonEvent_Landing(ev->Sender, ev, this), ctx);
 }