From 12023c681c4be06dc32fc7d2d9c3f72fc3fd1238 Mon Sep 17 00:00:00 2001
From: Koenraad Verheyden <koenraad.verheyden@grafana.com>
Date: Tue, 28 May 2024 18:34:49 +0200
Subject: [PATCH] Add toggle to inject the tenant ID to generated metrics
 (#3638)

(cherry picked from commit 6206fcdda6419a5cede65a6d118304cd91b56250)
---
 CHANGELOG.md                                |  1 +
 docs/sources/tempo/configuration/_index.md  |  3 +++
 modules/generator/registry/config.go        |  2 ++
 modules/generator/registry/registry.go      |  4 ++++
 modules/generator/registry/registry_test.go | 19 +++++++++++++++++++
 5 files changed, 29 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fb7ccb46532..98540075f81 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 * [CHANGE] Update Alpine image version to 3.20 [#3710](https://github.com/grafana/tempo/pull/3710) (@joe-elliott)
 * [ENHANCEMENT] TraceQL - Add support for trace:id and span:id [#3670](https://github.com/grafana/tempo/pull/3670) (@ie-pham)
+* [ENHANCEMENT] Add toggle to inject the tenant ID to generated metrics [#3638](https://github.com/grafana/tempo/pull/3638) (@kvrhdn)
 * [BUGFIX] Fix TraceQL queries involving non boolean operations between statics and attributes. [#3698](https://github.com/grafana/tempo/pull/3698) (@joe-elliott)
 
 ## v2.5.0-rc.0
diff --git a/docs/sources/tempo/configuration/_index.md b/docs/sources/tempo/configuration/_index.md
index 990c207688b..eaf1bf2476e 100644
--- a/docs/sources/tempo/configuration/_index.md
+++ b/docs/sources/tempo/configuration/_index.md
@@ -384,6 +384,9 @@ metrics_generator:
         # A list of labels that will be added to all generated metrics.
         [external_labels: <map>]
 
+        # If set, the tenant ID will added as label with the given label name to all generated metrics.
+        [inject_tenant_id_as: <string>]
+
         # The maximum length of label names. Label names exceeding this limit will be truncated.
         [max_label_name_length: <int> | default = 1024]
 
diff --git a/modules/generator/registry/config.go b/modules/generator/registry/config.go
index 42c8c71c14a..b264b31ee89 100644
--- a/modules/generator/registry/config.go
+++ b/modules/generator/registry/config.go
@@ -18,6 +18,8 @@ type Config struct {
 	// ExternalLabels are added to any time series generated by this instance.
 	ExternalLabels map[string]string `yaml:"external_labels,omitempty"`
 
+	InjectTenantIDAs string `yaml:"inject_tenant_id_as,omitempty"`
+
 	// MaxLabelNameLength configures the maximum length of label names. Label names exceeding
 	// this limit will be truncated.
 	MaxLabelNameLength int `yaml:"max_label_name_length"`
diff --git a/modules/generator/registry/registry.go b/modules/generator/registry/registry.go
index 12b42d916f3..be9f295c14c 100644
--- a/modules/generator/registry/registry.go
+++ b/modules/generator/registry/registry.go
@@ -101,6 +101,10 @@ func New(cfg *Config, overrides Overrides, tenant string, appendable storage.App
 	hostname, _ := os.Hostname()
 	externalLabels["__metrics_gen_instance"] = hostname
 
+	if cfg.InjectTenantIDAs != "" {
+		externalLabels[cfg.InjectTenantIDAs] = tenant
+	}
+
 	r := &ManagedRegistry{
 		onShutdown: cancel,
 
diff --git a/modules/generator/registry/registry_test.go b/modules/generator/registry/registry_test.go
index a13c83450f8..caf945851c6 100644
--- a/modules/generator/registry/registry_test.go
+++ b/modules/generator/registry/registry_test.go
@@ -153,6 +153,25 @@ func TestManagedRegistry_externalLabels(t *testing.T) {
 	collectRegistryMetricsAndAssert(t, registry, appender, expectedSamples)
 }
 
+func TestManagedRegistry_injectTenantIDAs(t *testing.T) {
+	appender := &capturingAppender{}
+
+	cfg := &Config{
+		InjectTenantIDAs: "__tempo_tenant",
+	}
+	registry := New(cfg, &mockOverrides{}, "test", appender, log.NewNopLogger())
+	defer registry.Close()
+
+	counter := registry.NewCounter("my_counter")
+	counter.Inc(nil, 1.0)
+
+	expectedSamples := []sample{
+		newSample(map[string]string{"__name__": "my_counter", "__metrics_gen_instance": mustGetHostname(), "__tempo_tenant": "test"}, 0, 0),
+		newSample(map[string]string{"__name__": "my_counter", "__metrics_gen_instance": mustGetHostname(), "__tempo_tenant": "test"}, 0, 1),
+	}
+	collectRegistryMetricsAndAssert(t, registry, appender, expectedSamples)
+}
+
 func TestManagedRegistry_maxSeries(t *testing.T) {
 	appender := &capturingAppender{}