Resource labels

CyberROFL · Jan 30, 2025 · 6b5142f · 6b5142f
1 parent d043784
commit 6b5142f
Show file tree

Hide file tree

Showing 8 changed files with 64 additions and 2 deletions.
diff --git a/ydb/core/grpc_services/rpc_rate_limiter_api.cpp b/ydb/core/grpc_services/rpc_rate_limiter_api.cpp
@@ -277,6 +277,7 @@ static void CopyProps(const Ydb::RateLimiter::Resource& src, NKikimrKesus::TStre
         auto copyMetric = [] (const Ydb::RateLimiter::MeteringConfig::Metric& srcMetric, NKikimrKesus::TAccountingConfig::TMetric& metric) {
             metric.SetEnabled(srcMetric.enabled());
             metric.SetBillingPeriodSec(srcMetric.billing_period_sec());
+            *metric.MutableLabels() = srcMetric.labels();
 
             /* overwrite if we have new fields */
             /* TODO: support arbitrary fields in metering core */
@@ -345,6 +346,7 @@ static void CopyProps(const NKikimrKesus::TStreamingQuoterResource& src, Ydb::Ra
         auto copyMetric = [] (const NKikimrKesus::TAccountingConfig::TMetric& srcMetric, Ydb::RateLimiter::MeteringConfig::Metric& metric) {
             metric.set_enabled(srcMetric.GetEnabled());
             metric.set_billing_period_sec(srcMetric.GetBillingPeriodSec());
+            *metric.mutable_labels() = srcMetric.GetLabels();
 
             /* TODO: support arbitrary fields in metering core */
             auto& metricFields = *metric.mutable_metric_fields()->mutable_fields();

diff --git a/ydb/core/kesus/tablet/quoter_resource_tree.cpp b/ydb/core/kesus/tablet/quoter_resource_tree.cpp
@@ -804,6 +804,7 @@ void THierarchicalDRRQuoterResourceTree::CalcParametersForAccounting() {
         if (!cfg->GetResourceId() && parent) { cfg->SetResourceId(parent->GetResourceId()); }
         if (!cfg->GetSourceId() && parent) { cfg->SetSourceId(parent->GetSourceId()); }
         if (cfg->GetTags().empty() && parent) { *cfg->MutableTags() = parent->GetTags(); }
+        if (cfg->GetLabels().empty() && parent) { *cfg->MutableLabels() = parent->GetLabels(); }
     };
     calcMetricsParams(accCfg->MutableProvisioned(), accCfgParent ? &accCfgParent->GetProvisioned() : nullptr);
     calcMetricsParams(accCfg->MutableOnDemand(), accCfgParent ? &accCfgParent->GetOnDemand() : nullptr);

diff --git a/ydb/core/kesus/tablet/rate_accounting.cpp b/ydb/core/kesus/tablet/rate_accounting.cpp
@@ -170,7 +170,8 @@ class TBillingMetric {
                 .FolderId(Cfg.GetFolderId())
                 .ResourceId(Cfg.GetResourceId())
                 .SourceId(Cfg.GetSourceId())
-                .Tags(ToJsonMap(Cfg.GetTags()));
+                .Tags(ToJsonMap(Cfg.GetTags()))
+                .Labels(ToJsonMap(Cfg.GetLabels()));
 
             LWPROBE(ResourceBillSend,
                 QuoterPath,

diff --git a/ydb/core/kesus/tablet/tablet_ut.cpp b/ydb/core/kesus/tablet/tablet_ut.cpp
@@ -2152,6 +2152,53 @@ Y_UNIT_TEST_SUITE(TKesusTest) {
         }
     }
 
+    Y_UNIT_TEST(TestQuoterAccountLabels) {
+        TTestContext ctx;
+        ctx.Setup();
+
+        TString billRecord;
+        ctx.Runtime->SetObserverFunc([&billRecord](TAutoPtr<IEventHandle>& ev) {
+            if (ev->GetTypeRewrite() == NMetering::TEvMetering::EvWriteMeteringJson) {
+                billRecord = ev->Get<NMetering::TEvMetering::TEvWriteMeteringJson>()->MeteringJson;
+            }
+            return TTestActorRuntime::EEventAction::PROCESS;
+        });
+
+        NKikimrKesus::TStreamingQuoterResource cfg;
+        cfg.SetResourcePath("/Root");
+        cfg.MutableHierarchicalDRRResourceConfig()->SetMaxUnitsPerSecond(100.0);
+        cfg.MutableHierarchicalDRRResourceConfig()->SetPrefetchCoefficient(300.0);
+        cfg.MutableAccountingConfig()->SetEnabled(true);
+        cfg.MutableAccountingConfig()->SetReportPeriodMs(1000);
+        cfg.MutableAccountingConfig()->SetAccountPeriodMs(1000);
+        cfg.MutableAccountingConfig()->SetCollectPeriodSec(2);
+        cfg.MutableAccountingConfig()->MutableOnDemand()->SetEnabled(true);
+        cfg.MutableAccountingConfig()->MutableOnDemand()->SetBillingPeriodSec(2);
+        cfg.MutableAccountingConfig()->MutableOnDemand()->MutableLabels()->insert({"k1", "v1"});
+        cfg.MutableAccountingConfig()->MutableOnDemand()->MutableLabels()->insert({"k2", "v2"});
+        ctx.AddQuoterResource(cfg);
+
+        auto edge = ctx.Runtime->AllocateEdgeActor();
+        auto client = ctx.Runtime->AllocateEdgeActor();
+        const NKikimrKesus::TEvSubscribeOnResourcesResult subscribeResult = ctx.SubscribeOnResource(client, edge, "/Root", false, 0);
+
+        TInstant start = ctx.Runtime->GetCurrentTime();
+        TDuration interval = TConsumptionHistory::Interval();
+        ctx.AccountResources(client, edge, subscribeResult.GetResults(0).GetResourceId(), start, interval, {50.0});
+
+        if (billRecord.empty()) {
+            TDispatchOptions opts;
+            opts.FinalEvents.emplace_back([&billRecord](IEventHandle&) -> bool {
+                return !billRecord.empty();
+            });
+            ctx.Runtime->DispatchEvents(opts);
+        }
+
+        const TString expectedBillRecord = R"({"usage":{"start":0,"quantity":50,"finish":1,"unit":"request_unit","type":"delta"},"tags":{},"id":"Root-ondemand-0-1","cloud_id":"","source_wt":5,"source_id":"","resource_id":"","schema":"","labels":{"k2":"v2","k1":"v1"},"folder_id":"","version":""})";
+
+        UNIT_ASSERT_NO_DIFF(billRecord, expectedBillRecord + "\n");
+    }
+
     Y_UNIT_TEST(TestPassesUpdatedPropsToSession) {
         TTestContext ctx;
         ctx.Setup();

diff --git a/ydb/core/metering/bill_record.cpp b/ydb/core/metering/bill_record.cpp
@@ -22,7 +22,7 @@ TString TBillRecord::TUsage::ToString() const {
 }
 
 NJson::TJsonMap TBillRecord::ToJson() const {
-    return NJson::TJsonMap{
+    auto json = NJson::TJsonMap{
         {"version", Version_},
         {"id", Id_},
         {"schema", Schema_},
@@ -34,6 +34,12 @@ NJson::TJsonMap TBillRecord::ToJson() const {
         {"tags", Tags_},
         {"usage", Usage_.ToJson()},
     };
+
+    if (Labels_.IsMap() && Labels_.GetMap()) {
+        json["labels"] = Labels_;
+    }
+
+    return json;
 }
 
 TString TBillRecord::ToString() const {

diff --git a/ydb/core/metering/bill_record.h b/ydb/core/metering/bill_record.h
@@ -55,6 +55,7 @@ struct TBillRecord {
     BILL_RECORD_FIELD_DEFAULT(TString, SourceId, "sless-docapi-ydb-ss");
     BILL_RECORD_FIELD(TInstant, SourceWt);
     BILL_RECORD_FIELD_DEFAULT(NJson::TJsonMap, Tags, {});
+    BILL_RECORD_FIELD_DEFAULT(NJson::TJsonMap, Labels, {});
     BILL_RECORD_FIELD(TUsage, Usage);
 
     NJson::TJsonMap ToJson() const;

diff --git a/ydb/core/protos/kesus.proto b/ydb/core/protos/kesus.proto
@@ -298,6 +298,7 @@ message TAccountingConfig {
         string ResourceId = 8;
         string SourceId = 9;
         map<string, string> Tags = 10;
+        map<string, string> Labels = 11;
     }
 
     // Consumption within provisioned limit.

diff --git a/ydb/public/api/protos/ydb_rate_limiter.proto b/ydb/public/api/protos/ydb_rate_limiter.proto
@@ -63,6 +63,9 @@ message MeteringConfig {
         // Default value is inherited from parent or equals 60 seconds for root.
         uint64 billing_period_sec = 2;
 
+        // User-defined labels.
+        map<string, string> labels = 3 [(map_key).length.le = 256, (length).le = 10240, (size).le = 100];
+
         // Billing metric JSON fields (inherited from parent if not set)
         google.protobuf.Struct metric_fields = 10;
     }