KIKIMR-20521: Add UseVDisksBalancing feature flag (#523)

Co-authored-by: robdrynkin <robdrynkin@nebius.com>

12 files changed, 134 insertions, 0 deletions

diff --git a/ydb/core/base/blobstorage.h b/ydb/core/base/blobstorage.h
index 0a5a56bff7..34c07e6fea 100644
--- a/ydb/core/base/blobstorage.h
+++ b/ydb/core/base/blobstorage.h
@@ -706,6 +706,7 @@ struct TEvBlobStorage {
         EvWriteMetadata,
         EvPermitGarbageCollection,
         EvReplInvoke,
+        EvStartBalancing,
 
         EvYardInitResult = EvPut + 9 * 512,                     /// 268 636 672
         EvLogResult,
diff --git a/ydb/core/blobstorage/nodewarden/node_warden_vdisk.cpp b/ydb/core/blobstorage/nodewarden/node_warden_vdisk.cpp
index 5c25c5e328..781dd6070f 100644
--- a/ydb/core/blobstorage/nodewarden/node_warden_vdisk.cpp
+++ b/ydb/core/blobstorage/nodewarden/node_warden_vdisk.cpp
@@ -174,6 +174,7 @@ namespace NKikimr::NStorage {
         vdiskConfig->EnableVDiskCooldownTimeout = Cfg->EnableVDiskCooldownTimeout;
         vdiskConfig->ReplPausedAtStart = Cfg->VDiskReplPausedAtStart;
         vdiskConfig->EnableVPatch = EnableVPatch;
+        vdiskConfig->FeatureFlags = Cfg->FeatureFlags;
 
         // issue initial report to whiteboard before creating actor to avoid races
         Send(WhiteboardId, new NNodeWhiteboard::TEvWhiteboard::TEvVDiskStateUpdate(vdiskId, groupInfo->GetStoragePoolName(),
diff --git a/ydb/core/blobstorage/vdisk/balance/balancing_actor.cpp b/ydb/core/blobstorage/vdisk/balance/balancing_actor.cpp
new file mode 100644
index 0000000000..aaae0b8b2a
--- /dev/null
+++ b/ydb/core/blobstorage/vdisk/balance/balancing_actor.cpp
@@ -0,0 +1,29 @@
+#include "balancing_actor.h"
+#include "defs.h"
+
+
+namespace NKikimr {
+
+    class TBalancingActor : public TActorBootstrapped<TBalancingActor> {
+    private:
+        std::shared_ptr<TBalancingCtx> Ctx;
+    public:
+        void Bootstrap() {
+            Become(&TThis::StateFunc);
+        }
+
+        STRICT_STFUNC(StateFunc,
+            CFunc(NActors::TEvents::TEvPoison::EventType, Die)
+        );
+
+        TBalancingActor(std::shared_ptr<TBalancingCtx> &ctx)
+            : TActorBootstrapped<TBalancingActor>()
+            , Ctx(ctx)
+        {
+        }
+    };
+
+    IActor* CreateBalancingActor(std::shared_ptr<TBalancingCtx> ctx) {
+        return new TBalancingActor(ctx);
+    }
+}
diff --git a/ydb/core/blobstorage/vdisk/balance/balancing_actor.h b/ydb/core/blobstorage/vdisk/balance/balancing_actor.h
new file mode 100644
index 0000000000..c39572fb93
--- /dev/null
+++ b/ydb/core/blobstorage/vdisk/balance/balancing_actor.h
@@ -0,0 +1,8 @@
+#pragma once
+
+#include "defs.h"
+
+
+namespace NKikimr {
+    IActor* CreateBalancingActor(std::shared_ptr<TBalancingCtx> ctx);
+} // NKikimr
diff --git a/ydb/core/blobstorage/vdisk/balance/defs.h b/ydb/core/blobstorage/vdisk/balance/defs.h
new file mode 100644
index 0000000000..0a8bc35ad5
--- /dev/null
+++ b/ydb/core/blobstorage/vdisk/balance/defs.h
@@ -0,0 +1,41 @@
+#pragma once
+
+#include <ydb/core/blobstorage/vdisk/common/vdisk_context.h>
+#include <ydb/core/blobstorage/vdisk/common/vdisk_pdiskctx.h>
+#include <ydb/core/blobstorage/vdisk/hulldb/hull_ds_all_snap.h>
+#include <ydb/core/blobstorage/vdisk/common/vdisk_hulllogctx.h>
+
+
+namespace NKikimr {
+    struct TBalancingCtx {
+        TIntrusivePtr<TVDiskContext> VCtx;
+        TPDiskCtxPtr PDiskCtx;
+        TActorId SkeletonId;
+        NMonGroup::TBalancingGroup MonGroup;
+
+        NKikimr::THullDsSnap Snap;
+
+        TIntrusivePtr<TVDiskConfig> VDiskCfg;
+        TIntrusivePtr<TBlobStorageGroupInfo> GInfo;
+
+        TBalancingCtx(
+            TIntrusivePtr<TVDiskContext> vCtx,
+            TPDiskCtxPtr pDiskCtx,
+            TActorId skeletonId,
+            NKikimr::THullDsSnap snap,
+            TIntrusivePtr<TVDiskConfig> vDiskCfg,
+            TIntrusivePtr<TBlobStorageGroupInfo> gInfo
+        )
+            : VCtx(std::move(vCtx))
+            , PDiskCtx(std::move(pDiskCtx))
+            , SkeletonId(skeletonId)
+            , MonGroup(VCtx->VDiskCounters, "subsystem", "balancing")
+            , Snap(std::move(snap))
+            , VDiskCfg(std::move(vDiskCfg))
+            , GInfo(std::move(gInfo))
+        {
+        }
+    };
+
+    struct TEvStartBalancing : TEventLocal<TEvStartBalancing, TEvBlobStorage::EvStartBalancing> {};
+} // NKikimr
diff --git a/ydb/core/blobstorage/vdisk/balance/ya.make b/ydb/core/blobstorage/vdisk/balance/ya.make
new file mode 100644
index 0000000000..6fb73e1555
--- /dev/null
+++ b/ydb/core/blobstorage/vdisk/balance/ya.make
@@ -0,0 +1,20 @@
+LIBRARY()
+
+PEERDIR(
+    ydb/core/blobstorage/base
+    ydb/core/blobstorage/groupinfo
+    ydb/core/blobstorage/vdisk/common
+    ydb/core/blobstorage/vdisk/hulldb
+    ydb/core/blobstorage/vdisk/ingress
+    ydb/core/blobstorage/vdisk/repl
+)
+
+SRCS(
+    balancing_actor.cpp
+)
+
+END()
+
+RECURSE_FOR_TESTS(
+    ut
+)
diff --git a/ydb/core/blobstorage/vdisk/common/vdisk_config.h b/ydb/core/blobstorage/vdisk/common/vdisk_config.h
index f4d64e2de7..93efc58d34 100644
--- a/ydb/core/blobstorage/vdisk/common/vdisk_config.h
+++ b/ydb/core/blobstorage/vdisk/common/vdisk_config.h
@@ -210,6 +210,9 @@ namespace NKikimr {
         bool EnableVDiskCooldownTimeout;
         TControlWrapper EnableVPatch = true;
 
+        ///////////// FEATURE FLAGS ////////////////////////
+        NKikimrConfig::TFeatureFlags FeatureFlags;
+
         TVDiskConfig(const TBaseInfo &baseInfo);
         void Merge(const NKikimrBlobStorage::TVDiskConfig &update);
     private:
diff --git a/ydb/core/blobstorage/vdisk/common/vdisk_mongroups.h b/ydb/core/blobstorage/vdisk/common/vdisk_mongroups.h
index 97137576df..08f3f60a66 100644
--- a/ydb/core/blobstorage/vdisk/common/vdisk_mongroups.h
+++ b/ydb/core/blobstorage/vdisk/common/vdisk_mongroups.h
@@ -547,6 +547,16 @@ public:
             COUNTER_DEF(DefragBytesRewritten);
         };
 
+        ///////////////////////////////////////////////////////////////////////////////////
+        // TBalancingGroup
+        ///////////////////////////////////////////////////////////////////////////////////
+        class TBalancingGroup : public TBase {
+        public:
+            GROUP_CONSTRUCTOR(TBalancingGroup)
+            {
+            }
+        };
+
     } // NMonGroup
 } // NKikimr
 
diff --git a/ydb/core/blobstorage/vdisk/skeleton/blobstorage_skeleton.cpp b/ydb/core/blobstorage/vdisk/skeleton/blobstorage_skeleton.cpp
index a197fe1a89..79ce890c99 100644
--- a/ydb/core/blobstorage/vdisk/skeleton/blobstorage_skeleton.cpp
+++ b/ydb/core/blobstorage/vdisk/skeleton/blobstorage_skeleton.cpp
@@ -19,6 +19,7 @@
 #include <ydb/core/base/feature_flags.h>
 #include <ydb/core/blobstorage/groupinfo/blobstorage_groupinfo_iter.h>
 #include <ydb/core/blobstorage/vdisk/localrecovery/localrecovery_public.h>
+#include <ydb/core/blobstorage/vdisk/balance/balancing_actor.h>
 #include <ydb/core/blobstorage/vdisk/hullop/blobstorage_hull.h>
 #include <ydb/core/blobstorage/vdisk/hullop/blobstorage_hulllog.h>
 #include <ydb/core/blobstorage/vdisk/huge/blobstorage_hullhuge.h>
@@ -1711,6 +1712,7 @@ namespace NKikimr {
                 ReplDone = true;
             }
             UpdateReplState();
+            RunBalancing(ctx);
         }
 
         void SkeletonErrorState(const TActorContext &ctx,
@@ -2449,6 +2451,20 @@ namespace NKikimr {
             ev->Get()->Callback({}, "replication is not in progress");
         }
 
+        void RunBalancing(const TActorContext &ctx) {
+            if (!Config->FeatureFlags.GetUseVDisksBalancing()) {
+                return;
+            }
+            if (BalancingId) {
+                ActiveActors.Erase(BalancingId);
+            }
+            auto balancingCtx = std::make_shared<TBalancingCtx>(
+                VCtx, PDiskCtx, SelfId(), Hull->GetHullDs()->GetIndexSnapshot(), Config, GInfo
+            );
+            BalancingId = ctx.Register(CreateBalancingActor(balancingCtx));
+            ActiveActors.Insert(BalancingId, __FILE__, __LINE__, ctx, NKikimrServices::BLOBSTORAGE);
+        }
+
         // NOTES: we have 4 state functions, one of which is an error state (StateDatabaseError) and
         // others are good: StateLocalRecovery, StateSyncGuidRecovery, StateNormal
         // We switch between states in the following manner:
@@ -2622,6 +2638,7 @@ namespace NKikimr {
             hFunc(NPDisk::TEvChunkForgetResult, Handle)
             FFunc(TEvPrivate::EvCheckSnapshotExpiration, CheckSnapshotExpiration)
             hFunc(TEvReplInvoke, Handle)
+            CFunc(TEvStartBalancing::EventType, RunBalancing)
         )
 
         STRICT_STFUNC(StateDatabaseError,
@@ -2730,6 +2747,7 @@ namespace NKikimr {
         bool CommenceRepl = false;
         TActorId ScrubId;
         TActorId DefragId;
+        TActorId BalancingId;
         bool HasUnreadableBlobs = false;
         std::unique_ptr<TVDiskCompactionState> VDiskCompactionState;
         TMemorizableControlWrapper EnableVPatch;
diff --git a/ydb/core/blobstorage/vdisk/skeleton/ya.make b/ydb/core/blobstorage/vdisk/skeleton/ya.make
index b96a726183..ab6a4e82be 100644
--- a/ydb/core/blobstorage/vdisk/skeleton/ya.make
+++ b/ydb/core/blobstorage/vdisk/skeleton/ya.make
@@ -2,6 +2,7 @@ LIBRARY()
 
 PEERDIR(
     ydb/core/base
+    ydb/core/blobstorage/vdisk/balance
     ydb/core/blobstorage/vdisk/hulldb/base
     ydb/core/blobstorage/vdisk/hulldb/bulksst_add
     ydb/core/blobstorage/vdisk/synclog
diff --git a/ydb/core/blobstorage/vdisk/ya.make b/ydb/core/blobstorage/vdisk/ya.make
index 61d0aa1a40..6628dec312 100644
--- a/ydb/core/blobstorage/vdisk/ya.make
+++ b/ydb/core/blobstorage/vdisk/ya.make
@@ -30,6 +30,7 @@ END()
 
 RECURSE(
     anubis_osiris
+    balance
     common
     defrag
     handoff
diff --git a/ydb/core/protos/feature_flags.proto b/ydb/core/protos/feature_flags.proto
index 247fb014a1..eb99a085e6 100644
--- a/ydb/core/protos/feature_flags.proto
+++ b/ydb/core/protos/feature_flags.proto
@@ -123,4 +123,5 @@ message TFeatureFlags {
     optional bool EnableTablePgTypes = 108 [default = false];
     optional bool EnableLocalDBBtreeIndex = 109 [default = false];
     optional bool EnablePDiskHighHDDInFlight = 110 [default = false];
+    optional bool UseVDisksBalancing = 111 [default = false];
 }