Fixes panic when using multiple client with the same prometheus regis…

…try. Introduce by grafana#3175. Signed-off-by: Cyril Tovena <[email protected]>
cyriltovena · Jan 21, 2021 · 33e9b3b · 33e9b3b
1 parent ea488f9
commit 33e9b3b
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 24 deletions.
diff --git a/pkg/promtail/client/client.go b/pkg/promtail/client/client.go
@@ -41,9 +41,7 @@ const (
 	HostLabel    = "host"
 )
 
-var (
-	UserAgent = fmt.Sprintf("promtail/%s", version.Version)
-)
+var UserAgent = fmt.Sprintf("promtail/%s", version.Version)
 
 type metrics struct {
 	encodedBytes     *prometheus.CounterVec
@@ -111,21 +109,29 @@ func newMetrics(reg prometheus.Registerer) *metrics {
 	}
 
 	if reg != nil {
-		reg.MustRegister(
-			m.encodedBytes,
-			m.sentBytes,
-			m.droppedBytes,
-			m.sentEntries,
-			m.droppedEntries,
-			m.requestDuration,
-			m.batchRetries,
-			m.streamLag,
-		)
+		m.encodedBytes = mustRegisterOrGet(reg, m.encodedBytes).(*prometheus.CounterVec)
+		m.sentBytes = mustRegisterOrGet(reg, m.sentBytes).(*prometheus.CounterVec)
+		m.droppedBytes = mustRegisterOrGet(reg, m.droppedBytes).(*prometheus.CounterVec)
+		m.sentEntries = mustRegisterOrGet(reg, m.sentEntries).(*prometheus.CounterVec)
+		m.droppedEntries = mustRegisterOrGet(reg, m.droppedEntries).(*prometheus.CounterVec)
+		m.requestDuration = mustRegisterOrGet(reg, m.requestDuration).(*prometheus.HistogramVec)
+		m.batchRetries = mustRegisterOrGet(reg, m.batchRetries).(*prometheus.CounterVec)
+		m.streamLag = mustRegisterOrGet(reg, m.streamLag).(*metric.Gauges)
 	}
 
 	return &m
 }
 
+func mustRegisterOrGet(reg prometheus.Registerer, c prometheus.Collector) prometheus.Collector {
+	if err := reg.Register(c); err != nil {
+		if are, ok := err.(prometheus.AlreadyRegisteredError); ok {
+			return are.ExistingCollector.(prometheus.Collector)
+		}
+		panic(err)
+	}
+	return c
+}
+
 // Client pushes entries to Loki and can be stopped
 type Client interface {
 	api.EntryHandler

diff --git a/pkg/promtail/client/client_test.go b/pkg/promtail/client/client_test.go
@@ -24,16 +24,14 @@ import (
 	lokiflag "github.com/grafana/loki/pkg/util/flagext"
 )
 
-var (
-	logEntries = []api.Entry{
-		{Labels: model.LabelSet{}, Entry: logproto.Entry{Timestamp: time.Unix(1, 0).UTC(), Line: "line1"}},
-		{Labels: model.LabelSet{}, Entry: logproto.Entry{Timestamp: time.Unix(2, 0).UTC(), Line: "line2"}},
-		{Labels: model.LabelSet{}, Entry: logproto.Entry{Timestamp: time.Unix(3, 0).UTC(), Line: "line3"}},
-		{Labels: model.LabelSet{"__tenant_id__": "tenant-1"}, Entry: logproto.Entry{Timestamp: time.Unix(4, 0).UTC(), Line: "line4"}},
-		{Labels: model.LabelSet{"__tenant_id__": "tenant-1"}, Entry: logproto.Entry{Timestamp: time.Unix(5, 0).UTC(), Line: "line5"}},
-		{Labels: model.LabelSet{"__tenant_id__": "tenant-2"}, Entry: logproto.Entry{Timestamp: time.Unix(6, 0).UTC(), Line: "line6"}},
-	}
-)
+var logEntries = []api.Entry{
+	{Labels: model.LabelSet{}, Entry: logproto.Entry{Timestamp: time.Unix(1, 0).UTC(), Line: "line1"}},
+	{Labels: model.LabelSet{}, Entry: logproto.Entry{Timestamp: time.Unix(2, 0).UTC(), Line: "line2"}},
+	{Labels: model.LabelSet{}, Entry: logproto.Entry{Timestamp: time.Unix(3, 0).UTC(), Line: "line3"}},
+	{Labels: model.LabelSet{"__tenant_id__": "tenant-1"}, Entry: logproto.Entry{Timestamp: time.Unix(4, 0).UTC(), Line: "line4"}},
+	{Labels: model.LabelSet{"__tenant_id__": "tenant-1"}, Entry: logproto.Entry{Timestamp: time.Unix(5, 0).UTC(), Line: "line5"}},
+	{Labels: model.LabelSet{"__tenant_id__": "tenant-2"}, Entry: logproto.Entry{Timestamp: time.Unix(6, 0).UTC(), Line: "line6"}},
+}
 
 type receivedReq struct {
 	tenantID string

diff --git a/pkg/promtail/client/multi_test.go b/pkg/promtail/client/multi_test.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
 
 	"github.com/grafana/loki/pkg/logproto"
@@ -37,7 +38,7 @@ func TestNewMulti(t *testing.T) {
 		ExternalLabels: lokiflag.LabelSet{LabelSet: model.LabelSet{"hi": "there"}},
 	}
 
-	clients, err := NewMulti(nil, util.Logger, lokiflag.LabelSet{LabelSet: model.LabelSet{"order": "command"}}, cc1, cc2)
+	clients, err := NewMulti(prometheus.DefaultRegisterer, util.Logger, lokiflag.LabelSet{LabelSet: model.LabelSet{"order": "command"}}, cc1, cc2)
 	if err != nil {
 		t.Fatalf("expected err: nil got:%v", err)
 	}