From 34d95179faa597f82c74aaeaff0e6485122f223b Mon Sep 17 00:00:00 2001
From: Robert Fratto <robertfratto@gmail.com>
Date: Fri, 23 Feb 2024 15:38:04 -0500
Subject: [PATCH] loki.write: fix duplicate metrics collector registration
 (#6511)

Fix an issue where re-creating a write client led to a duplicate
metrics collector registration panic.

Fixes #6510
---
 CHANGELOG.md                                  | 10 ++++++---
 .../common/loki/client/internal/metrics.go    |  7 ++++--
 component/common/loki/client/manager_test.go  | 22 +++++++++++++++++++
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 72a8d2f70e13..4f5a6c3edf60 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,10 @@ internal API changes are not present.
 Main (unreleased)
 -----------------
 
+### Bugfixes
+
+- Fix an issue where changing the configuration of `loki.write` would cause a panic. (@rfratto)
+
 v0.40.0-rc.1 (2024-02-23)
 -------------------------
 
@@ -27,7 +31,7 @@ v0.40.0-rc.1 (2024-02-23)
 
 ### Features
 
-- Modules have been redesigned to split the import logic from the instantiation. 
+- Modules have been redesigned to split the import logic from the instantiation.
   You can now define custom components via the `declare` config block and import modules via `import.git`, `import.http`, `import.string`, `import.file`. (@wildum)
 
 - A new `discovery.process` component for discovering Linux OS processes on the current host. (@korniltsev)
@@ -40,7 +44,7 @@ v0.40.0-rc.1 (2024-02-23)
 - Expose track_timestamps_staleness on Prometheus scraping, to fix the issue where container metrics live for 5 minutes after the container disappears. (@ptodev)
 
 - Introduce the `remotecfg` service that enables loading configuration from a
-  remote endpoint. (@tpaschalis) 
+  remote endpoint. (@tpaschalis)
 
 - Add `otelcol.connector.host_info` component to gather usage metrics for cloud users. (@rlankfo, @jcreixell)
 
@@ -102,7 +106,7 @@ v0.40.0-rc.1 (2024-02-23)
 - Fix an issue where agent logs are emitted before the logging format
   is correctly determined. (@hainenber)
 
-- Fix divide-by-zero issue when sharding targets. (@hainenber) 
+- Fix divide-by-zero issue when sharding targets. (@hainenber)
 
 - Fix bug where custom headers were not actually being set in loki client. (@captncraig)
 
diff --git a/component/common/loki/client/internal/metrics.go b/component/common/loki/client/internal/metrics.go
index 3abd7572eac2..f169fde39807 100644
--- a/component/common/loki/client/internal/metrics.go
+++ b/component/common/loki/client/internal/metrics.go
@@ -1,6 +1,9 @@
 package internal
 
-import "github.com/prometheus/client_golang/prometheus"
+import (
+	"github.com/grafana/agent/pkg/util"
+	"github.com/prometheus/client_golang/prometheus"
+)
 
 type MarkerMetrics struct {
 	lastMarkedSegment *prometheus.GaugeVec
@@ -19,7 +22,7 @@ func NewMarkerMetrics(reg prometheus.Registerer) *MarkerMetrics {
 		),
 	}
 	if reg != nil {
-		reg.MustRegister(m.lastMarkedSegment)
+		m.lastMarkedSegment = util.MustRegisterOrGet(reg, m.lastMarkedSegment).(*prometheus.GaugeVec)
 	}
 	return m
 }
diff --git a/component/common/loki/client/manager_test.go b/component/common/loki/client/manager_test.go
index 97ba4c1b46cb..6dae1370d476 100644
--- a/component/common/loki/client/manager_test.go
+++ b/component/common/loki/client/manager_test.go
@@ -36,6 +36,28 @@ var (
 	metrics    = NewMetrics(prometheus.DefaultRegisterer)
 )
 
+// TestManager_NoDuplicateMetricsPanic ensures that creating two managers does
+// not lead to duplicate metrics registration.
+func TestManager_NoDuplicateMetricsPanic(t *testing.T) {
+	var (
+		host, _ = url.Parse("http://localhost:3100")
+
+		reg     = prometheus.NewRegistry()
+		metrics = NewMetrics(reg)
+	)
+
+	require.NotPanics(t, func() {
+		for i := 0; i < 2; i++ {
+			_, err := NewManager(metrics, log.NewLogfmtLogger(os.Stdout), testLimitsConfig, reg, wal.Config{
+				WatchConfig: wal.DefaultWatchConfig,
+			}, NilNotifier, Config{
+				URL: flagext.URLValue{URL: host},
+			})
+			require.NoError(t, err)
+		}
+	})
+}
+
 func TestManager_ErrorCreatingWhenNoClientConfigsProvided(t *testing.T) {
 	for _, walEnabled := range []bool{true, false} {
 		t.Run(fmt.Sprintf("wal-enabled = %t", walEnabled), func(t *testing.T) {