Test for info collector in Bridge Mode CMS (#20631)

Author: stanislav-shchetinin

ydb/core/cms/cms_ut.cpp

@@ -2353,6 +2353,49 @@ Y_UNIT_TEST_SUITE(TCmsTest) {
         // Wall-E soft maintainance task can continue
         env.CheckWalleCheckTask("task-1", TStatus::ALLOW, env.GetNodeId(2));
     }
+
+    void ChangePileMap(TEvInterconnect::TEvNodesInfo::TPtr* ev) {
+        UNIT_ASSERT((*ev)->Get()->PileMap);
+        auto nodes = MakeIntrusive<TIntrusiveVector<TEvInterconnect::TNodeInfo>>((*ev)->Get()->Nodes);
+        auto pileMap = std::make_shared<TEvInterconnect::TEvNodesInfo::TPileMap>(*(*ev)->Get()->PileMap);
+        for (const auto& node : *nodes) {
+            pileMap->at(node.NodeId % pileMap->size()).push_back(node.NodeId);
+        }
+
+        auto newEv = IEventHandle::Downcast<TEvInterconnect::TEvNodesInfo>(
+            new IEventHandle((*ev)->Recipient, (*ev)->Sender, new TEvInterconnect::TEvNodesInfo(nodes, pileMap))
+        );
+        ev->Swap(newEv);
+    }
+
+    Y_UNIT_TEST(BridgeModeCollectInfo)
+    {
+        TTestEnvOpts opts(8);
+        TCmsTestEnv env(opts.WithBridgeMode());
+
+        auto observerFunc = [&](TAutoPtr<IEventHandle>& ev) {
+            if (ev->GetTypeRewrite() == TEvInterconnect::EvNodesInfo) {
+                auto *x = reinterpret_cast<TEvInterconnect::TEvNodesInfo::TPtr*>(&ev);
+                ChangePileMap(x);
+            }
+            return TTestActorRuntime::EEventAction::PROCESS;
+        };
+        env.SetObserverFunc(observerFunc);
+
+        env.Register(CreateInfoCollector(env.GetSender(), TDuration::Minutes(1)));
+
+        TAutoPtr<IEventHandle> handle;
+        auto reply = env.GrabEdgeEventRethrow<TCms::TEvPrivate::TEvClusterInfo>(handle);
+        UNIT_ASSERT(reply);
+        const auto &info = *reply->Info;
+        UNIT_ASSERT(info.IsBridgeMode);
+        UNIT_ASSERT(info.NodeIdToPileId.size() == 8);
+
+        for (const auto [nodeId, pileId] : info.NodeIdToPileId) {
+            UNIT_ASSERT(nodeId % opts.PileCount == pileId);
+        }
+
+    }
 }
 
 }

ydb/core/cms/cms_ut_common.cpp

@@ -565,6 +565,12 @@ static void SetupServices(TTestBasicRuntime &runtime, const TTestEnvOpts &option
     runtime.GetAppData().DynamicNameserviceConfig = dnsConfig;
     runtime.GetAppData().DisableCheckingSysNodesCms = true;
     runtime.GetAppData().BootstrapConfig = TFakeNodeWhiteboardService::BootstrapConfig;
+    
+    if (options.IsBridgeMode) {
+        for (ui32 pileId = 0; pileId < options.PileCount; ++pileId) {
+            runtime.GetAppData().BridgeConfig->AddPiles()->SetName("r" + ToString(pileId));
+        }
+    }
 
     NKikimrCms::TCmsConfig cmsConfig;
     cmsConfig.MutableSentinelConfig()->SetEnable(options.EnableSentinel);
@@ -622,11 +628,25 @@ TCmsTestEnv::TCmsTestEnv(const TTestEnvOpts &options)
     mallocInfo.SetParam("FillMemoryOnAllocation", "false");
     SetupLogging();
 
-    for (ui32 nodeIndex = 0; nodeIndex < GetNodeCount(); ++nodeIndex) {
-        if (options.NRings > 1) {
-            SetupCustomStateStorage(*this, options.NToSelect, options.NRings, options.RingSize);
-        } else {
-            SetupStateStorage(*this, nodeIndex);
+    if (options.IsBridgeMode) {
+        TVector<TStateStorageInfo::TRingGroup> ringGroups = {{.State = PRIMARY, .NToSelect = options.NToSelect}};
+        std::fill_n(
+            std::back_inserter(ringGroups),
+            options.PileCount - 1,
+            TStateStorageInfo::TRingGroup{.State = SYNCHRONIZED, .NToSelect = options.NToSelect}
+        );
+        auto setuper = CreateCustomStateStorageSetupper(ringGroups, options.NRings * options.RingSize);
+
+        for (ui32 nodeIndex = 0; nodeIndex < GetNodeCount(); ++nodeIndex) {
+            setuper(*this, nodeIndex);
+        }
+    } else {
+        for (ui32 nodeIndex = 0; nodeIndex < GetNodeCount(); ++nodeIndex) {
+            if (options.NRings > 1) {
+                SetupCustomStateStorage(*this, options.NToSelect, options.NRings, options.RingSize);
+            } else {
+                SetupStateStorage(*this, nodeIndex);
+            }
         }
     }
     SetupServices(*this, options);

ydb/core/cms/cms_ut_common.h

@@ -85,13 +85,15 @@ struct TTestEnvOpts {
     ui32 NRings;
     ui32 RingSize;
     ui32 DataCenterCount;
+    ui32 PileCount;
     TNodeTenantsMap Tenants;
     bool UseMirror3dcErasure;
     bool AdvanceCurrentTime;
     bool EnableSentinel;
     bool EnableCMSRequestPriorities;
     bool EnableSingleCompositeActionGroup;
     bool EnableDynamicGroups;
+    bool IsBridgeMode;
 
     using TNodeLocationCallback = std::function<TNodeLocation(ui32)>;
     TNodeLocationCallback NodeLocationCallback;
@@ -107,13 +109,15 @@ struct TTestEnvOpts {
         , NRings(1)
         , RingSize(nodeCount)
         , DataCenterCount(1)
+        , PileCount(0)
         , Tenants(tenants)
         , UseMirror3dcErasure(false)
         , AdvanceCurrentTime(false)
         , EnableSentinel(false)
         , EnableCMSRequestPriorities(true)
         , EnableSingleCompositeActionGroup(true)
         , EnableDynamicGroups(false)
+        , IsBridgeMode(false)
     {
     }
 
@@ -141,6 +145,12 @@ struct TTestEnvOpts {
         EnableDynamicGroups = true;
         return *this;
     }
+
+    TTestEnvOpts& WithBridgeMode(ui32 pileCount = 2) {
+        IsBridgeMode = true;
+        PileCount = pileCount;
+        return *this;
+    }
 };
 
 class TCmsTestEnv : public TTestBasicRuntime {