Add static node selectors by labels (#18913)

3 files changed, 314 insertions, 12 deletions

diff --git a/ydb/core/config/init/init_impl.h b/ydb/core/config/init/init_impl.h
index 792272c8d02..80dfd65f8b0 100644
--- a/ydb/core/config/init/init_impl.h
+++ b/ydb/core/config/init/init_impl.h
@@ -1249,6 +1249,7 @@ public:
         Labels["branch"] = GetBranch();
         Labels["rev"] = GetProgramCommitId();
         Labels["dynamic"] = ToString(CommonAppOptions.IsStaticNode() ? "false" : "true");
+        Labels["node_kind"] = CommonAppOptions.IsStaticNode() ? "static" : "dynamic";
 
         for (const auto& [name, value] : Labels) {
             auto *label = AppConfig.AddLabels();
@@ -1319,14 +1320,57 @@ public:
         }
     }
 
+    class TAppConfigFieldsPreserver {
+    public:
+        TAppConfigFieldsPreserver(NKikimrConfig::TAppConfig& appConfig) 
+            : AppConfig(appConfig)
+            , ConfigDirPath(appConfig.GetConfigDirPath())
+            , StoredConfigYaml(appConfig.GetStoredConfigYaml())
+            , StartupConfigYaml(appConfig.GetStartupConfigYaml())
+            , StartupStorageYaml(appConfig.GetStartupStorageYaml())
+        {}
+
+        ~TAppConfigFieldsPreserver() {
+            if (ConfigDirPath) {
+                AppConfig.SetConfigDirPath(*ConfigDirPath);
+            }
+            if (StoredConfigYaml) {
+                AppConfig.MutableStoredConfigYaml()->CopyFrom(*StoredConfigYaml);
+            }
+            if (StartupConfigYaml) {
+                AppConfig.SetStartupConfigYaml(*StartupConfigYaml);
+            }
+            if (StartupStorageYaml) {
+                AppConfig.SetStartupStorageYaml(*StartupStorageYaml);
+            }
+        }
+    private:
+        NKikimrConfig::TAppConfig& AppConfig;
+        std::optional<TString> ConfigDirPath;
+        std::optional<NKikimrBlobStorage::TYamlConfig> StoredConfigYaml;
+        std::optional<TString> StartupConfigYaml;
+        std::optional<TString> StartupStorageYaml;
+    };
+
     void InitStaticNode() {
         CommonAppOptions.ValidateStaticNodeConfig();
-
+        Labels["node_kind"] = "static";
         Labels["dynamic"] = "false";
+
+        if (!AppConfig.HasStartupConfigYaml()) {
+            return;
+        }
+
+        TAppConfigFieldsPreserver preserver(AppConfig);
+
+        NKikimrConfig::TAppConfig appConfig;
+        NYamlConfig::ResolveAndParseYamlConfig(AppConfig.GetStartupConfigYaml(), {}, Labels, appConfig);
+        ApplyConfigForNode(appConfig);
     }
 
     void InitDynamicNode() {
         Labels["dynamic"] = "true";
+        Labels["node_kind"] = "dynamic";
         RegisterDynamicNode(CommonAppOptions);
 
         Labels["node_id"] = ToString(NodeId);
diff --git a/ydb/core/config/init/init_ut.cpp b/ydb/core/config/init/init_ut.cpp
index 8344b4954ba..f3f4f79f111 100644
--- a/ydb/core/config/init/init_ut.cpp
+++ b/ydb/core/config/init/init_ut.cpp
@@ -1,7 +1,14 @@
 #include "init_impl.h"
+#include "mock.h"
 
 #include <library/cpp/testing/unittest/registar.h>
+#include <util/system/tempfile.h>
+#include <util/stream/output.h>
+#include <ydb/core/config/init/init.h>
+#include <ydb/core/config/validation/validators.h>
+#include <ydb/library/yaml_config/yaml_config.h>
 
+using namespace NKikimr;
 using namespace NKikimr::NConfig;
 
 Y_UNIT_TEST_SUITE(Init) {
@@ -48,3 +55,250 @@ Y_UNIT_TEST_SUITE(Init) {
         }
     }
 }
+
+
+Y_UNIT_TEST_SUITE(StaticNodeSelectorsInit) {
+
+    TTempFileHandle CreateConfigFile() {
+        TString config = R"(
+metadata:
+  kind: MainConfig
+  version: 0
+  cluster: test_cluster
+config:
+  default_disk_type: NVME
+  self_management_config:
+    enabled: true
+
+  actor_system_config:
+    use_auto_config: true
+    node_type: COMPUTE
+    cpu_count: 10
+
+  host_configs:
+  - nvme:
+    - disk1
+    - disk2
+  hosts:
+    - host: localhost
+      port: 2135
+
+  log_config:
+    default_level: 5
+    entry:
+    - component: BS_CONTROLLER
+      level: 7
+
+allowed_labels:
+  test:
+    type: string
+
+selector_config:
+  - description: "Selector for static nodes"
+    selector:
+      test: abc
+    config:
+      actor_system_config:
+        use_auto_config: true
+        node_type: STORAGE
+        cpu_count: 100
+      log_config: !inherit
+        entry: !append
+        - component: CONSOLE_HANDSHAKE
+          level: 6
+  - description: "Selector by node_id"
+    selector:
+      node_id: 1
+      test: node_id
+    config:
+      actor_system_config:
+        use_auto_config: true
+        node_type: STORAGE
+        cpu_count: 1
+  - description: "Selector by node_host"
+    selector:
+      node_host: localhost
+      test: node_host
+    config:
+      actor_system_config:
+        use_auto_config: true
+        node_type: STORAGE
+        cpu_count: 50
+  - description: "Selector by node_kind"
+    selector:
+      node_kind: static
+      test: node_kind
+    config:
+      actor_system_config:
+        use_auto_config: true
+        node_type: STORAGE
+        cpu_count: 42
+)";
+        TTempFileHandle tempFile = TTempFileHandle::InCurrentDir("test_config", ".yaml");
+        TUnbufferedFileOutput fileOutput(tempFile.Name());
+        fileOutput.Write(config);
+        fileOutput.Finish();
+        return tempFile;
+    }
+
+    void PreFillArgs(std::vector<TString>& args, const TString& configPath) {
+        args.push_back("server");
+
+        args.push_back("--node");
+        args.push_back("static");
+
+        args.push_back("--ic-port");
+        args.push_back("2135");
+
+        args.push_back("--grpc-port");
+        args.push_back("9001");
+
+        args.push_back("--mon-port");
+        args.push_back("8765");
+
+        args.push_back("--yaml-config");
+        args.push_back(configPath);
+    }
+
+    NKikimrConfig::TAppConfig TransformConfig(const std::vector<TString>& args) {
+        auto errorCollector = NConfig::MakeDefaultErrorCollector();
+        auto protoConfigFileProvider = NConfig::MakeDefaultProtoConfigFileProvider();
+        auto configUpdateTracer = NConfig::MakeDefaultConfigUpdateTracer();
+        auto memLogInit = NConfig::MakeNoopMemLogInitializer();
+        auto nodeBrokerClient = NConfig::MakeNoopNodeBrokerClient();
+        auto dynConfigClient = NConfig::MakeNoopDynConfigClient();
+        auto logger = NConfig::MakeNoopInitLogger();
+
+        auto envMock = std::make_unique<NConfig::TEnvMock>();
+        envMock->SavedHostName = "localhost";
+        envMock->SavedFQDNHostName = "localhost";
+
+        NConfig::TInitialConfiguratorDependencies deps{
+            *errorCollector,
+            *protoConfigFileProvider,
+            *configUpdateTracer,
+            *memLogInit,
+            *nodeBrokerClient,
+            *dynConfigClient,
+            *envMock,
+            *logger,
+        };
+        auto initCfg = NConfig::MakeDefaultInitialConfigurator(deps);
+
+        std::vector<const char*> argv;
+
+        for (const auto& arg : args) {
+            argv.push_back(arg.data());
+        }
+
+        NLastGetopt::TOpts opts;
+        initCfg->RegisterCliOptions(opts);
+        protoConfigFileProvider->RegisterCliOptions(opts);
+
+        NLastGetopt::TOptsParseResult parseResult(&opts, argv.size(), argv.data());
+
+        initCfg->ValidateOptions(opts, parseResult);
+        initCfg->Parse(parseResult.GetFreeArgs(), nullptr);
+
+        NKikimrConfig::TAppConfig appConfig;
+        ui32 nodeId;
+        TKikimrScopeId scopeId;
+        TString tenantName;
+        TBasicKikimrServicesMask servicesMask;
+        TString clusterName;
+        NConfig::TConfigsDispatcherInitInfo configsDispatcherInitInfo;
+
+        initCfg->Apply(
+            appConfig,
+            nodeId,
+            scopeId,
+            tenantName,
+            servicesMask,
+            clusterName,
+            configsDispatcherInitInfo);
+
+        return appConfig;
+    }
+
+    Y_UNIT_TEST(TestStaticNodeSelectorForActorSystem) {
+        TTempFileHandle configFile = CreateConfigFile();
+        TVector<TString> args;
+        PreFillArgs(args, configFile.Name());
+        args.push_back("--label");
+        args.push_back("test=abc");
+        NKikimrConfig::TAppConfig appConfig = TransformConfig(args);
+
+        UNIT_ASSERT(appConfig.HasActorSystemConfig());
+        const auto& actorConfig = appConfig.GetActorSystemConfig();
+        UNIT_ASSERT_EQUAL(actorConfig.GetCpuCount(), 100);
+    }
+
+    Y_UNIT_TEST(TestStaticNodeSelectorWithAnotherLabel) {
+        TTempFileHandle configFile = CreateConfigFile();
+        TVector<TString> args;
+        PreFillArgs(args, configFile.Name());
+        args.push_back("--label");
+        args.push_back("test=abd");
+        NKikimrConfig::TAppConfig appConfig = TransformConfig(args);
+
+        UNIT_ASSERT(appConfig.HasActorSystemConfig());
+        const auto& actorConfig = appConfig.GetActorSystemConfig();
+        UNIT_ASSERT_VALUES_EQUAL(actorConfig.GetCpuCount(), 10);
+        
+    }
+
+    Y_UNIT_TEST(TestStaticNodeSelectorInheritance) {
+        TTempFileHandle configFile = CreateConfigFile();
+        TVector<TString> args;
+        PreFillArgs(args, configFile.Name());
+        args.push_back("--label");
+        args.push_back("test=abc");
+        NKikimrConfig::TAppConfig appConfig = TransformConfig(args);
+
+        UNIT_ASSERT(appConfig.HasLogConfig());
+        const auto& logConfig = appConfig.GetLogConfig();
+        UNIT_ASSERT_EQUAL(logConfig.GetDefaultLevel(), 5);
+        UNIT_ASSERT_EQUAL(logConfig.GetEntry(0).GetLevel(), 7);
+        UNIT_ASSERT_EQUAL(logConfig.GetEntry(1).GetLevel(), 6);
+    }
+
+    Y_UNIT_TEST(TestStaticNodeSelectorByNodeId) {
+        TTempFileHandle configFile = CreateConfigFile();
+        TVector<TString> args;
+        PreFillArgs(args, configFile.Name());
+        args.push_back("--label");
+        args.push_back("test=node_id");
+        NKikimrConfig::TAppConfig appConfig = TransformConfig(args);
+
+        UNIT_ASSERT(appConfig.HasActorSystemConfig());
+        const auto& actorConfig = appConfig.GetActorSystemConfig();
+        UNIT_ASSERT_EQUAL(actorConfig.GetCpuCount(), 1);
+    }
+
+    Y_UNIT_TEST(TestStaticNodeSelectorByNodeHost) {
+        TTempFileHandle configFile = CreateConfigFile();
+        TVector<TString> args;
+        PreFillArgs(args, configFile.Name());
+        args.push_back("--label");
+        args.push_back("test=node_host");
+
+        NKikimrConfig::TAppConfig appConfig = TransformConfig(args);
+
+        UNIT_ASSERT(appConfig.HasActorSystemConfig());
+        const auto& actorConfig = appConfig.GetActorSystemConfig();
+        UNIT_ASSERT_EQUAL(actorConfig.GetCpuCount(), 50);
+    }
+
+    Y_UNIT_TEST(TestStaticNodeSelectorByNodeKind) {
+        TTempFileHandle configFile = CreateConfigFile();
+        TVector<TString> args;
+        PreFillArgs(args, configFile.Name());
+        args.push_back("--label");
+        args.push_back("test=node_kind");
+        NKikimrConfig::TAppConfig appConfig = TransformConfig(args);
+
+        UNIT_ASSERT(appConfig.HasActorSystemConfig());
+        const auto& actorConfig = appConfig.GetActorSystemConfig();
+        UNIT_ASSERT_EQUAL(actorConfig.GetCpuCount(), 42);
+    }
+}
diff --git a/ydb/library/yaml_config/public/yaml_config.cpp b/ydb/library/yaml_config/public/yaml_config.cpp
index 2bdaf1b30f4..7a8fc9f28f4 100644
--- a/ydb/library/yaml_config/public/yaml_config.cpp
+++ b/ydb/library/yaml_config/public/yaml_config.cpp
@@ -359,17 +359,21 @@ TDocumentConfig Resolve(
     res.first.Resolve();
 
     auto rootMap = res.first.Root().Map();
-    auto config = rootMap.at("config");
-    if (rootMap.Has("selector_config")) {
-        auto selectorConfig = rootMap.at("selector_config").Sequence();
-
-        for (auto it = selectorConfig.begin(); it != selectorConfig.end(); ++it) {
-            auto selectorMap = it->Map();
-            auto desc = selectorMap.at("description").Scalar();
-            auto selectorNode = selectorMap.at("selector");
-            auto selector = ParseSelector(selectorNode);
-            if (Fit(selector, labels)) {
-                Apply(config, selectorMap.at("config"));
+    auto config = res.first.Root();
+
+    if (rootMap.Has("config")) {
+        config = rootMap.at("config");
+        if (rootMap.Has("selector_config")) {
+            auto selectorConfig = rootMap.at("selector_config").Sequence();
+
+            for (auto it = selectorConfig.begin(); it != selectorConfig.end(); ++it) {
+                auto selectorMap = it->Map();
+                auto desc = selectorMap.at("description").Scalar();
+                auto selectorNode = selectorMap.at("selector");
+                auto selector = ParseSelector(selectorNode);
+                if (Fit(selector, labels)) {
+                    Apply(config, selectorMap.at("config"));
+                }
             }
         }
     }