PR from branch users/t1mursadykov/KIKIMR-13670

ref:e53a46f5debe708f270270953ae73ce498bfe0b3

12 files changed, 433 insertions, 15 deletions

diff --git a/ydb/core/base/statestorage.h b/ydb/core/base/statestorage.h
index 9b3e3b6b1b7..3a3c06c3b81 100644
--- a/ydb/core/base/statestorage.h
+++ b/ydb/core/base/statestorage.h
@@ -25,6 +25,7 @@ struct TEvStateStorage {
         EvResolveSchemeBoard, // subset (by hash)
         EvListSchemeBoard, // all
         EvUpdateGroupConfig,
+        EvListStateStorage,
 
         // replies (local, from proxy)
         EvInfo = EvLookup + 512,
@@ -33,6 +34,7 @@ struct TEvStateStorage {
         EvResponseReplicasDumps,
         EvDeleteResult,
         EvListSchemeBoardResult,
+        EvListStateStorageResult,
 
         // replicas interface
         EvReplicaLookup = EvLock + 2 * 512,
@@ -373,6 +375,8 @@ struct TEvStateStorage {
     struct TEvReplicaBoardInfo;
     struct TEvListSchemeBoard;
     struct TEvListSchemeBoardResult;
+    struct TEvListStateStorage;
+    struct TEvListStateStorageResult;
     struct TEvUpdateGroupConfig;
     struct TEvReplicaProbeSubscribe;
     struct TEvReplicaProbeUnsubscribe;
diff --git a/ydb/core/base/statestorage_impl.h b/ydb/core/base/statestorage_impl.h
index 2a2aa22f6a7..8951dc6c4be 100644
--- a/ydb/core/base/statestorage_impl.h
+++ b/ydb/core/base/statestorage_impl.h
@@ -165,6 +165,17 @@ struct TEvStateStorage::TEvListSchemeBoardResult : public TEventLocal<TEvListSch
     {}
 };
 
+struct TEvStateStorage::TEvListStateStorage : public TEventLocal<TEvListStateStorage, EvListStateStorage> {
+};
+
+struct TEvStateStorage::TEvListStateStorageResult : public TEventLocal<TEvListStateStorageResult, EvListStateStorageResult> {
+    TIntrusiveConstPtr<TStateStorageInfo> Info;
+
+    TEvListStateStorageResult(const TIntrusiveConstPtr<TStateStorageInfo> &info)
+        : Info(info)
+    {}
+};
+
 struct TEvStateStorage::TEvReplicaLookup : public TEventPB<TEvStateStorage::TEvReplicaLookup, NKikimrStateStorage::TEvLookup, TEvStateStorage::EvReplicaLookup>{
     struct TActualityCounter : public TRefCounted<TActualityCounter, TAtomicCounter> {};
     using TActualityCounterPtr = TIntrusivePtr<TActualityCounter>;
diff --git a/ydb/core/base/statestorage_proxy.cpp b/ydb/core/base/statestorage_proxy.cpp
index 13d396cb0f1..38fc1661dc8 100644
--- a/ydb/core/base/statestorage_proxy.cpp
+++ b/ydb/core/base/statestorage_proxy.cpp
@@ -812,6 +812,15 @@ class TStateStorageProxy : public TActor<TStateStorageProxy> {
         Send(ev->Sender, new TEvStateStorage::TEvListSchemeBoardResult(SchemeBoardInfo), 0, ev->Cookie);
     }
 
+    void Handle(TEvStateStorage::TEvListStateStorage::TPtr &ev) {
+        if (!Info) {
+            Send(ev->Sender, new TEvStateStorage::TEvListStateStorageResult(nullptr), 0, ev->Cookie);
+            return;
+        }
+
+        Send(ev->Sender, new TEvStateStorage::TEvListStateStorageResult(Info), 0, ev->Cookie);
+    }
+
     void Handle(TEvStateStorage::TEvReplicaProbeSubscribe::TPtr &ev) {
         const auto *msg = ev->Get();
 
@@ -972,6 +981,7 @@ public:
             hFunc(TEvStateStorage::TEvResolveBoard, Handle);
             hFunc(TEvStateStorage::TEvResolveSchemeBoard, Handle);
             hFunc(TEvStateStorage::TEvListSchemeBoard, Handle);
+            hFunc(TEvStateStorage::TEvListStateStorage, Handle);
             hFunc(TEvStateStorage::TEvUpdateGroupConfig, Handle);
             hFunc(TEvStateStorage::TEvReplicaProbeSubscribe, Handle);
             hFunc(TEvStateStorage::TEvReplicaProbeUnsubscribe, Handle);
diff --git a/ydb/core/cms/cms.cpp b/ydb/core/cms/cms.cpp
index 0a205024e3c..45452c9a21c 100644
--- a/ydb/core/cms/cms.cpp
+++ b/ydb/core/cms/cms.cpp
@@ -5,6 +5,8 @@
 
 #include <ydb/core/actorlib_impl/long_timer.h>
 #include <ydb/core/base/appdata.h>
+#include <ydb/core/base/statestorage.h>
+#include <ydb/core/base/statestorage_impl.h>
 #include <ydb/core/cms/console/config_helpers.h>
 #include <ydb/core/base/ticket_parser.h>
 #include <ydb/core/tablet_flat/tablet_flat_executed.h>
@@ -166,6 +168,16 @@ void TCms::ProcessInitQueue(const TActorContext &ctx)
     }
 }
 
+void TCms::RequestStateStorageConfig(const TActorContext &ctx) {
+    const auto& domains = *AppData(ctx)->DomainsInfo;
+    ui32 domainUid = domains.Domains.begin()->second->DomainUid;
+    const ui32 stateStorageGroup = domains.GetDefaultStateStorageGroup(domainUid);
+
+    const TActorId proxy = MakeStateStorageProxyID(stateStorageGroup);
+
+    ctx.Send(proxy, new TEvStateStorage::TEvListStateStorage(), IEventHandle::FlagTrackDelivery);
+}
+
 void TCms::SubscribeForConfig(const TActorContext &ctx)
 {
     ctx.Register(NConsole::CreateConfigSubscriber(TabletID(),
@@ -452,6 +464,11 @@ bool TCms::CheckActionShutdownNode(const NKikimrCms::TAction &action,
         return false;
     }
 
+    // node is not locked
+    if (!CheckActionShutdownStateStorage(action, opts, node, error)) {
+        return false;
+    }
+
     return true;
 }
 
@@ -504,6 +521,75 @@ bool TCms::CheckActionShutdownHost(const TAction &action,
     return true;
 }
 
+bool TCms::CheckActionShutdownStateStorage(
+                         const TAction& action,
+                         const TActionOptions& opts,
+                         const TNodeInfo& node,
+                         TErrorInfo& error) const 
+{
+    // TODO (t1mursadykov): отслеживание времени отключенных стейт стораджей
+    if (opts.AvailabilityMode == MODE_FORCE_RESTART) {
+        return true;
+    }
+    
+    TInstant defaultDeadline = TActivationContext::Now() + State->Config.DefaultRetryTime;
+
+    if (!StateStorageInfo) {
+        error.Code = TStatus::DISALLOW_TEMP;
+        error.Reason = "Did not received state storage configuration";
+        error.Deadline = defaultDeadline;
+        return false;
+    }
+
+    if (!StateStorageNodes.contains(node.NodeId)) {
+        return true;
+    }
+
+    THashSet<ui32> injuredRings;
+    TDuration duration = TDuration::MicroSeconds(action.GetDuration()) + opts.PermissionDuration;
+    TErrorInfo err;
+    TStringStream brokenNodesMsg;
+    for (auto& i : StateStorageNodes) {
+        if (node.NodeId == i) {
+            continue;
+        }
+        if (ClusterInfo->Node(i).IsLocked(err, State->Config.DefaultRetryTime, 
+                                               TActivationContext::Now(), duration) || 
+            ClusterInfo->Node(i).IsDown(err, defaultDeadline)) {  
+
+            injuredRings.insert(NodeToRing.at(i));
+            brokenNodesMsg << " " << i;
+        }
+    }
+
+    if (injuredRings.size() == 0) {
+        return true;
+    }
+
+    if ((opts.AvailabilityMode == MODE_MAX_AVAILABILITY && injuredRings.size() > 1) ||
+        (opts.AvailabilityMode == MODE_KEEP_AVAILABLE && injuredRings.size() > 2)) {
+        error.Code = TStatus::DISALLOW_TEMP;
+        error.Reason = TStringBuilder() << "Too many broken state storage rings: " << injuredRings.size() <<
+                                        " Down state storage nodes:" << brokenNodesMsg.Str();
+        error.Deadline = defaultDeadline;
+        return false;
+    }
+
+    if (injuredRings.contains(NodeToRing.at(node.NodeId))) {
+        return true;
+    }
+
+    if (opts.AvailabilityMode == MODE_MAX_AVAILABILITY || injuredRings.size() > 2) {
+        error.Code = TStatus::DISALLOW_TEMP;
+        error.Reason = TStringBuilder() << "There are down state storage nodes in other rings. "
+                                         "Down state storage nodes: " << brokenNodesMsg.Str();
+        error.Deadline = defaultDeadline;
+        return false;
+    }
+
+    return true;
+}
+
 bool TCms::TryToLockNode(const TAction& action,
                          const TActionOptions& opts,
                          const TNodeInfo& node,
@@ -1330,6 +1416,27 @@ void TCms::OnBSCPipeDestroyed(const TActorContext &ctx)
         ctx.Send(State->Sentinel, new TEvSentinel::TEvBSCPipeDisconnected);
 }
 
+void TCms::Handle(TEvStateStorage::TEvListStateStorageResult::TPtr& ev, const TActorContext &ctx) {
+    auto& info = ev->Get()->Info;
+    if (!info) {
+        LOG_NOTICE_S(ctx, NKikimrServices::CMS,
+                     "Couldn't collect group info");
+        return;
+    }
+
+    StateStorageInfo = info;
+    
+    // index in array will be used as ring id for simplicity
+    for (ui32 ring = 0; ring < info->Rings.size(); ++ring) {
+        for (auto& replica : info->Rings[ring].Replicas) {
+            ui32 nodeId = replica.NodeId();
+
+            NodeToRing[nodeId] = ring;
+            StateStorageNodes.insert(nodeId);
+        }
+    }
+}
+
 void TCms::Handle(TEvPrivate::TEvClusterInfo::TPtr &ev, const TActorContext &ctx)
 {
     if (!ev->Get()->Success) {
diff --git a/ydb/core/cms/cms_impl.h b/ydb/core/cms/cms_impl.h
index 48138884907..1cf57487800 100644
--- a/ydb/core/cms/cms_impl.h
+++ b/ydb/core/cms/cms_impl.h
@@ -10,6 +10,7 @@
 #include <library/cpp/actors/core/hfunc.h>
 #include <library/cpp/actors/core/interconnect.h>
 #include <ydb/core/base/tablet_pipe.h>
+#include <ydb/core/base/statestorage_impl.h>
 #include <ydb/core/cms/console/console.h>
 #include <ydb/core/tablet_flat/tablet_flat_executed.h>
 #include <ydb/core/engine/minikql/flat_local_tx_factory.h>
@@ -236,6 +237,7 @@ private:
     STFUNC(StateWork)
     {
         switch (ev->GetTypeRewrite()) {
+            HFunc(TEvStateStorage::TEvListStateStorageResult, Handle);
             HFunc(TEvPrivate::TEvClusterInfo, Handle);
             HFunc(TEvPrivate::TEvLogAndSend, Handle);
             FFunc(TEvPrivate::EvUpdateClusterInfo, EnqueueRequest);
@@ -285,6 +287,7 @@ private:
                  const TActorContext &ctx) override;
     void ProcessInitQueue(const TActorContext &ctx);
 
+    void RequestStateStorageConfig(const TActorContext &ctx);
     void SubscribeForConfig(const TActorContext &ctx);
     void AdjustInfo(TClusterInfoPtr &info, const TActorContext &ctx) const;
     bool CheckPermissionRequest(const NKikimrCms::TPermissionRequest &request,
@@ -330,6 +333,10 @@ private:
                         const TVDiskInfo &vdisk,
                         TDuration duration,
                         TErrorInfo &error) const;
+    bool CheckActionShutdownStateStorage(const NKikimrCms::TAction &action,
+                       const TActionOptions &options,
+                       const TNodeInfo &node, 
+                       TErrorInfo& error) const;
     void AcceptPermissions(NKikimrCms::TPermissionResponse &resp, const TString &requestId,
                            const TString &owner, const TActorContext &ctx, bool check = false);
     void ScheduleUpdateClusterInfo(const TActorContext &ctx, bool now = false);
@@ -383,6 +390,7 @@ private:
 
     void OnBSCPipeDestroyed(const TActorContext &ctx);
 
+    void Handle(TEvStateStorage::TEvListStateStorageResult::TPtr& ev, const TActorContext &ctx);
     void Handle(TEvPrivate::TEvClusterInfo::TPtr &ev, const TActorContext &ctx);
     void Handle(TEvPrivate::TEvLogAndSend::TPtr &ev, const TActorContext &ctx);
     void Handle(TEvPrivate::TEvUpdateClusterInfo::TPtr &ev, const TActorContext &ctx);
@@ -426,6 +434,11 @@ private:
     TSchedulerCookieHolder WalleCleanupTimerCookieHolder;
     TSchedulerCookieHolder LogCleanupTimerCookieHolder;
 
+    // State Storage
+    TIntrusiveConstPtr<TStateStorageInfo> StateStorageInfo;
+    THashMap<ui32, ui32> NodeToRing;
+    THashSet<ui32> StateStorageNodes;
+
 public:
     TCms(const TActorId &tablet, TTabletStorageInfo *info)
         : TActor(&TThis::StateInit)
diff --git a/ydb/core/cms/cms_tx_load_state.cpp b/ydb/core/cms/cms_tx_load_state.cpp
index eb144dc0def..441c5d52972 100644
--- a/ydb/core/cms/cms_tx_load_state.cpp
+++ b/ydb/core/cms/cms_tx_load_state.cpp
@@ -195,6 +195,7 @@ public:
         Self->ScheduleLogCleanup(ctx);
         Self->ScheduleUpdateClusterInfo(ctx, true);
         Self->ProcessInitQueue(ctx);
+        Self->RequestStateStorageConfig(ctx);
     }
 };
 
diff --git a/ydb/core/cms/cms_ut.cpp b/ydb/core/cms/cms_ut.cpp
index 7be5390fdeb..2709267cf80 100644
--- a/ydb/core/cms/cms_ut.cpp
+++ b/ydb/core/cms/cms_ut.cpp
@@ -1059,6 +1059,160 @@ Y_UNIT_TEST_SUITE(TCmsTest) {
         env.CheckPermissionRequest("user", false, true, false, true, TStatus::DISALLOW_TEMP,
                                    MakeAction(TAction::SHUTDOWN_HOST, env.GetNodeId(1), 60000000));
     }
+
+    Y_UNIT_TEST(StateStorageTwoRings)
+    {   
+        TTestEnvOpts options;
+        options.NodeCount = 5;
+        options.VDisks = 0;
+        options.NRings = 2;
+        options.RingSize = 2;
+        options.NToSelect = 2;
+    
+        TCmsTestEnv env(options);
+        
+        env.CheckPermissionRequest("user", false, false, false, true, TStatus::ALLOW,
+                                   MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(0), 60000000, "storage"));
+        
+        env.CheckPermissionRequest("user", false, false, false, true, TStatus::DISALLOW_TEMP,
+                                    MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(2), 60000000, "storage"));
+    }
+
+    Y_UNIT_TEST(StateStorageNodesFromOneRing)
+    {   
+        TTestEnvOpts options;
+        options.NodeCount = 5;
+        options.VDisks = 0;
+        options.NRings = 2;
+        options.RingSize = 2;
+        options.NToSelect = 2;
+    
+        TCmsTestEnv env(options);
+
+        env.CheckPermissionRequest("user", false, false, false, true, TStatus::ALLOW,
+                                   MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(0), 60000000, "storage"));
+        
+        env.CheckPermissionRequest("user", false, false, false, true, TStatus::ALLOW,
+                                    MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(1), 60000000, "storage"));
+    }
+
+    Y_UNIT_TEST(StateStorageAvailabilityMode)
+    {   
+        TTestEnvOpts options;
+        options.NodeCount = 5;
+        options.VDisks = 0;
+        options.NRings = 2;
+        options.RingSize = 2;
+        options.NToSelect = 2;
+    
+        TCmsTestEnv env(options);
+        
+        TFakeNodeWhiteboardService::Info[env.GetNodeId(1)].Connected = false;
+        env.RestartCms();
+
+        env.CheckPermissionRequest("user", false, true, false, true, MODE_KEEP_AVAILABLE, TStatus::ALLOW,
+                                   MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(3), 60000000, "storage"));
+        
+        env.CheckPermissionRequest("user", false, true, false, true, MODE_MAX_AVAILABILITY, TStatus::DISALLOW_TEMP,
+                                    MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(3), 60000000, "storage"));
+    }
+
+    Y_UNIT_TEST(StateStorageTwoBrokenRings)
+    {   
+        TTestEnvOpts options;
+        options.NodeCount = 7;
+        options.VDisks = 0;
+        options.NRings = 3;
+        options.RingSize = 2;
+        options.NToSelect = 2;
+    
+        TCmsTestEnv env(options);
+
+        TFakeNodeWhiteboardService::Info[env.GetNodeId(0)].Connected = false;
+        TFakeNodeWhiteboardService::Info[env.GetNodeId(2)].Connected = false;
+
+        env.RestartCms();
+
+        env.CheckPermissionRequest("user", false, true, false, true, MODE_KEEP_AVAILABLE, TStatus::ALLOW,
+                                   MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(1), 60000000, "storage"));
+        
+        env.CheckPermissionRequest("user", false, true, false, true, MODE_MAX_AVAILABILITY, TStatus::DISALLOW_TEMP,
+                                    MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(1), 60000000, "storage"));
+    }
+
+    Y_UNIT_TEST(StateStorageRollingRestart) 
+    {
+        TTestEnvOpts options;
+        options.NodeCount = 20;
+        options.VDisks = 0;
+        options.NRings = 6;
+        options.RingSize = 3;
+        options.NToSelect = 5;
+    
+        TCmsTestEnv env(options);
+
+        TIntrusiveConstPtr<TStateStorageInfo> info = env.GetStateStorageInfo();
+
+        THashMap<ui32, ui32> NodeToRing;
+        THashSet<ui32> StateStorageNodes;
+
+        for (ui32 ring = 0; ring < info->Rings.size(); ++ring) {
+            for (auto& replica : info->Rings[ring].Replicas) {
+                ui32 nodeId = replica.NodeId();
+
+                NodeToRing[nodeId] = ring;
+                StateStorageNodes.insert(nodeId);
+            }
+        }
+        
+        THashSet<TString> restarted;
+
+        while(restarted.size() < env.GetNodeCount()) {
+
+            TAutoPtr<NCms::TEvCms::TEvPermissionRequest> event = new NCms::TEvCms::TEvPermissionRequest;
+            event->Record.SetUser("user");
+            event->Record.SetPartialPermissionAllowed(true);
+            event->Record.SetDryRun(false);
+            event->Record.SetSchedule(false);
+
+            for (ui32 i = 0; i < env.GetNodeCount(); ++i) {
+                if (!restarted.contains(TStringBuilder() << env.GetNodeId(i))) {
+                    AddActions(event, MakeAction(TAction::RESTART_SERVICES, env.GetNodeId(i), 0, "storage"));
+                }
+            }
+
+            NKikimrCms::TPermissionResponse res;
+
+            // In the last request comes the permission 
+            // for all nodes of the same ring
+            if (env.GetNodeCount() - restarted.size() == options.RingSize) {
+                res = env.CheckPermissionRequest(event, TStatus::ALLOW);
+            } else {
+                res = env.CheckPermissionRequest(event, TStatus::ALLOW_PARTIAL);
+            }
+            
+            ui32 permRing = env.GetNodeCount() + 1;
+            for (auto& perm : res.GetPermissions()) {
+                auto &action = perm.GetAction();
+                restarted.insert(action.GetHost());
+                env.CheckDonePermission("user", perm.GetId());
+
+                auto nodeId = std::stoi(action.GetHost());
+                if (!StateStorageNodes.contains(nodeId)) {
+                    continue;
+                }
+
+                // Check that all state storages in permissions 
+                // from the same ring
+                ui32 curRing = NodeToRing.at(nodeId);
+                if (permRing >= options.NRings) { // we have not met a state storage yet
+                    permRing = curRing;
+                }
+
+                UNIT_ASSERT_VALUES_EQUAL(permRing, curRing);
+            }
+        }
+    }
 }
 
 } // NCmsTest
diff --git a/ydb/core/cms/cms_ut_common.cpp b/ydb/core/cms/cms_ut_common.cpp
index b458fe672ce..7883923eb98 100644
--- a/ydb/core/cms/cms_ut_common.cpp
+++ b/ydb/core/cms/cms_ut_common.cpp
@@ -342,8 +342,6 @@ static void SetupServices(TTestActorRuntime &runtime,
     }
 
     for (ui32 nodeIndex = 0; nodeIndex < runtime.GetNodeCount(); ++nodeIndex) {
-        SetupStateStorage(runtime, nodeIndex);
-
         TString staticConfig("AvailabilityDomains: 0 "
                              "PDisks { NodeID: $Node1 PDiskID: 0 PDiskGuid: 1 Path: \"pdisk0.dat\" }"
                              "VDisks { VDiskID { GroupID: 0 GroupGeneration: 1 Ring: 0 Domain: 0 VDisk: 0 }"
@@ -437,11 +435,9 @@ static void SetupServices(TTestActorRuntime &runtime,
 
 } // anonymous namespace
 
-TCmsTestEnv::TCmsTestEnv(ui32 nodeCount,
-                         ui32 pdisks,
-                         const TNodeTenantsMap &tenants)
-    : TTestBasicRuntime(nodeCount, false)
-    , CmsId(MakeCmsID(0))
+TCmsTestEnv::TCmsTestEnv(const TTestEnvOpts &options) 
+        : TTestBasicRuntime(options.NodeCount, false)
+        , CmsId(MakeCmsID(0)) 
 {
     TFakeNodeWhiteboardService::Config.MutableResponse()->SetSuccess(true);
     TFakeNodeWhiteboardService::Config.MutableResponse()->ClearStatus();
@@ -449,7 +445,7 @@ TCmsTestEnv::TCmsTestEnv(ui32 nodeCount,
     status.SetSuccess(true);
     auto *config = status.MutableBaseConfig();
 
-    GenerateExtendedInfo(*this, config, pdisks, 4, tenants);
+    GenerateExtendedInfo(*this, config, options.VDisks, 4, options.Tenants);
 
     // Set observer to pass fake base blobstorage config.
     auto redirectConfigRequest = [](TTestActorRuntimeBase&,
@@ -467,7 +463,15 @@ TCmsTestEnv::TCmsTestEnv(ui32 nodeCount,
     TMallocInfo mallocInfo = MallocInfo();
     mallocInfo.SetParam("FillMemoryOnAllocation", "false");
     SetupLogging();
-    SetupServices(*this, tenants);
+
+    for (ui32 nodeIndex = 0; nodeIndex < GetNodeCount(); ++nodeIndex) {
+        if (options.NRings > 1) {
+            SetupCustomStateStorage(*this, options.NToSelect, options.NRings, options.RingSize, 0);
+        } else {
+            SetupStateStorage(*this, nodeIndex);
+        }
+    }
+    SetupServices(*this, options.Tenants);
 
     Sender = AllocateEdgeActor();
 
@@ -477,12 +481,30 @@ TCmsTestEnv::TCmsTestEnv(ui32 nodeCount,
     SetCmsConfig(cmsConfig);
 }
 
+
 TCmsTestEnv::TCmsTestEnv(ui32 nodeCount,
+                         ui32 vdisks,
                          const TNodeTenantsMap &tenants)
-    : TCmsTestEnv(nodeCount, 0, tenants)
+    : TCmsTestEnv(TTestEnvOpts(nodeCount, vdisks, tenants))
 {
 }
 
+TCmsTestEnv::TCmsTestEnv(ui32 nodeCount,
+                         const TNodeTenantsMap &tenants)
+    : TCmsTestEnv(TTestEnvOpts(nodeCount, 0, tenants))
+{
+}
+
+TIntrusiveConstPtr<NKikimr::TStateStorageInfo> TCmsTestEnv::GetStateStorageInfo() {
+    ui32 StateStorageGroup = 0;
+    const TActorId proxy = MakeStateStorageProxyID(StateStorageGroup);
+    Send(new IEventHandle(proxy, Sender, new TEvStateStorage::TEvListStateStorage()));
+
+    auto reply = GrabEdgeEventRethrow<TEvStateStorage::TEvListStateStorageResult>(Sender);
+    const auto &rec = reply->Get()->Info;
+    return rec;
+}
+
 void TCmsTestEnv::SetupLogging()
 {
     NActors::NLog::EPriority priority = ENABLE_DETAILED_CMS_LOG ? NLog::PRI_DEBUG : NLog::PRI_ERROR;
diff --git a/ydb/core/cms/cms_ut_common.h b/ydb/core/cms/cms_ut_common.h
index f1f9c1a74db..588617d9b88 100644
--- a/ydb/core/cms/cms_ut_common.h
+++ b/ydb/core/cms/cms_ut_common.h
@@ -3,6 +3,7 @@
 #include "ut_helpers.h"
 
 #include <ydb/core/base/counters.h>
+#include <ydb/core/base/statestorage.h>
 #include <ydb/core/node_whiteboard/node_whiteboard.h>
 #include <ydb/core/mind/tenant_pool.h>
 #include <ydb/core/testlib/basics/helpers.h>
@@ -65,6 +66,30 @@ public:
     void Handle(TEvWhiteboard::TEvSystemStateRequest::TPtr &ev, const TActorContext &ctx);
 };
 
+struct TTestEnvOpts {
+    ui32 NodeCount;
+    ui32 VDisks;
+
+    ui32 NToSelect;
+    ui32 NRings;
+    ui32 RingSize;
+    TNodeTenantsMap Tenants;
+
+    TTestEnvOpts() = default;
+
+    TTestEnvOpts(ui32 nodeCount, 
+                ui32 vdisks = 1,
+                const TNodeTenantsMap &tenants = TNodeTenantsMap())
+            : NodeCount(nodeCount)
+            , VDisks(vdisks)
+            , NToSelect(1)
+            , NRings(1)
+            , RingSize(3)
+            , Tenants(tenants) 
+    {
+    }
+};
+
 class TCmsTestEnv : public TTestBasicRuntime {
 public:
     TCmsTestEnv(ui32 nodeCount,
@@ -72,6 +97,7 @@ public:
                 const TNodeTenantsMap &tenants = TNodeTenantsMap());
     TCmsTestEnv(ui32 nodeCount,
                 const TNodeTenantsMap &tenants);
+    TCmsTestEnv(const TTestEnvOpts &options);
 
     TActorId GetSender() { return Sender; }
 
@@ -89,6 +115,8 @@ public:
                    ui32 tenantRatioLimit,
                    ui32 clusterLimit,
                    ui32 clusterRatioLimit);
+    
+    TIntrusiveConstPtr<NKikimr::TStateStorageInfo> GetStateStorageInfo();
 
     NKikimrCms::TClusterState
     RequestState(const NKikimrCms::TClusterStateRequest &request = {},
@@ -165,6 +193,10 @@ public:
                                       code, actions...);
     }
 
+    NKikimrCms::TPermissionResponse
+    CheckPermissionRequest(TAutoPtr<NCms::TEvCms::TEvPermissionRequest> req,
+                           NKikimrCms::TStatus::ECode code);
+
     NKikimrCms::TManagePermissionResponse
     CheckManagePermissionRequest(const TString &user,
                                  NKikimrCms::TManagePermissionRequest::ECommand cmd,
@@ -335,9 +367,6 @@ public:
 private:
     void SetupLogging();
 
-    NKikimrCms::TPermissionResponse
-    CheckPermissionRequest(TAutoPtr<NCms::TEvCms::TEvPermissionRequest> req,
-                           NKikimrCms::TStatus::ECode code);
     NKikimrCms::TManagePermissionResponse
     CheckManagePermissionRequest(TAutoPtr<NCms::TEvCms::TEvManagePermissionRequest> req,
                                  NKikimrCms::TStatus::ECode code);
diff --git a/ydb/core/cms/walle_create_task_adapter.cpp b/ydb/core/cms/walle_create_task_adapter.cpp
index 56989e2db03..e6363ccdf69 100644
--- a/ydb/core/cms/walle_create_task_adapter.cpp
+++ b/ydb/core/cms/walle_create_task_adapter.cpp
@@ -134,7 +134,7 @@ private:
             }
         }
 
-                TAutoPtr<TEvCms::TEvPermissionRequest> request = new TEvCms::TEvPermissionRequest;
+        TAutoPtr<TEvCms::TEvPermissionRequest> request = new TEvCms::TEvPermissionRequest;
         request->Record.SetUser(WALLE_CMS_USER);
         request->Record.SetSchedule(true);
         request->Record.SetDryRun(task.GetDryRun());
diff --git a/ydb/core/testlib/basics/helpers.h b/ydb/core/testlib/basics/helpers.h
index 7a70063f7a3..c71e7d84f69 100644
--- a/ydb/core/testlib/basics/helpers.h
+++ b/ydb/core/testlib/basics/helpers.h
@@ -40,6 +40,8 @@ namespace NFake {
 
     void SetupStateStorage(TTestActorRuntime& runtime, ui32 nodeIndex,
                            ui64 stateStorageGroup = 0, bool replicasOnFirstNode = false);
+    void SetupCustomStateStorage(TTestActorRuntime &runtime, ui32 NToSelect, 
+                           ui32 nrings, ui32 ringSize, ui64 stateStorageGroup); 
     void SetupBSNodeWarden(TTestActorRuntime& runtime, ui32 nodeIndex, TIntrusivePtr<TNodeWardenConfig> nodeWardenConfig);
     void SetupTabletResolver(TTestActorRuntime& runtime, ui32 nodeIndex);
     void SetupTabletPipePeNodeCaches(TTestActorRuntime& runtime, ui32 nodeIndex);
diff --git a/ydb/core/testlib/basics/services.cpp b/ydb/core/testlib/basics/services.cpp
index b4585ad96c6..c67024a6cbc 100644
--- a/ydb/core/testlib/basics/services.cpp
+++ b/ydb/core/testlib/basics/services.cpp
@@ -153,7 +153,7 @@ namespace NPDisk {
     template<size_t N>
     static TIntrusivePtr<TStateStorageInfo> GenerateStateStorageInfo(const TActorId (&replicas)[N], ui64 stateStorageGroup)
     {
-        TIntrusivePtr<TStateStorageInfo> info(new TStateStorageInfo());
+        auto info = MakeIntrusive<TStateStorageInfo>();
         info->StateStorageGroup = stateStorageGroup;
         info->NToSelect = N;
         info->Rings.resize(N);
@@ -164,6 +164,26 @@ namespace NPDisk {
         return info;
     }
 
+    static TIntrusivePtr<TStateStorageInfo> GenerateStateStorageInfo(const TVector<TActorId> &replicas, ui32 NToSelect, ui32 nrings, ui32 ringSize)
+    {   
+        Y_VERIFY(replicas.size() >= nrings * ringSize);
+        Y_VERIFY(NToSelect <= nrings);
+
+        auto info = MakeIntrusive<TStateStorageInfo>();
+        info->StateStorageGroup = 0;
+        info->NToSelect = NToSelect;
+        info->Rings.resize(nrings);
+            
+        ui32 inode = 0;
+        for (size_t i = 0; i < nrings; ++i) {
+            for (size_t j = 0; j < ringSize; ++j) {
+                info->Rings[i].Replicas.push_back(replicas[inode++]);
+            }
+        }
+
+        return info;
+    }
+
     static TActorId MakeBoardReplicaID(
         const ui32 node,
         const ui64 stateStorageGroup,
@@ -175,6 +195,51 @@ namespace NPDisk {
         return TActorId(node, TStringBuf(x, 12));
     }
 
+    void SetupCustomStateStorage(
+        TTestActorRuntime &runtime,
+        ui32 NToSelect, 
+        ui32 nrings, 
+        ui32 ringSize,
+        ui64 stateStorageGroup)
+    {   
+        TVector<TActorId> ssreplicas;
+        for (size_t i = 0; i < nrings * ringSize; ++i) {
+            ssreplicas.push_back(MakeStateStorageReplicaID(runtime.GetNodeId(i), stateStorageGroup, i));
+        }
+
+        TVector<TActorId> breplicas;
+        for (size_t i = 0; i < nrings * ringSize; ++i) {
+            breplicas.push_back(MakeBoardReplicaID(runtime.GetNodeId(i), stateStorageGroup, i));
+        }
+
+        TVector<TActorId> sbreplicas;
+        for (size_t i = 0; i < nrings * ringSize; ++i) {
+            sbreplicas.push_back(MakeSchemeBoardReplicaID(runtime.GetNodeId(i), stateStorageGroup, i));
+        }
+
+        const TActorId ssproxy = MakeStateStorageProxyID(stateStorageGroup);
+
+        auto ssInfo = GenerateStateStorageInfo(ssreplicas, NToSelect, nrings, ringSize);
+        auto sbInfo = GenerateStateStorageInfo(sbreplicas, NToSelect, nrings, ringSize);
+        auto bInfo = GenerateStateStorageInfo(breplicas, NToSelect, nrings, ringSize);
+
+        
+        for (ui32 ssIndex = 0; ssIndex < nrings * ringSize; ++ssIndex) {
+            runtime.AddLocalService(ssreplicas[ssIndex],
+                TActorSetupCmd(CreateStateStorageReplica(ssInfo.Get(), ssIndex), TMailboxType::Revolving, 0), ssIndex);
+            runtime.AddLocalService(sbreplicas[ssIndex],
+                TActorSetupCmd(CreateSchemeBoardReplica(sbInfo.Get(), ssIndex), TMailboxType::Revolving, 0), ssIndex);
+            runtime.AddLocalService(breplicas[ssIndex],
+                TActorSetupCmd(CreateStateStorageBoardReplica(bInfo.Get(), ssIndex), TMailboxType::Revolving, 0), ssIndex);
+        }
+
+        for (ui32 nodeIndex = 0; nodeIndex < runtime.GetNodeCount(); ++nodeIndex) {
+            runtime.AddLocalService(ssproxy,
+                    TActorSetupCmd(CreateStateStorageProxy(ssInfo.Get(), bInfo.Get(), sbInfo.Get()), TMailboxType::Revolving, 0), nodeIndex);
+        }
+    }
+
+    
     void SetupStateStorage(TTestActorRuntime& runtime, ui32 nodeIndex, ui64 stateStorageGroup, bool firstNode)
     {
         const TActorId ssreplicas[3] = {