From 2cb011f59588ea5be42b0c5e51c8000ca99b5346 Mon Sep 17 00:00:00 2001
From: Julien Salleyron <julien.salleyron@gmail.com>
Date: Mon, 18 Jul 2022 10:36:11 +0200
Subject: [PATCH] Fix service up gauge for Prometheus metrics

Co-authored-by: Tom Moulard <tom.moulard@traefik.io>
---
 pkg/metrics/prometheus.go      | 17 ++++++------
 pkg/metrics/prometheus_test.go | 48 +++++++++++++++++++++++++++++++++-
 2 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/pkg/metrics/prometheus.go b/pkg/metrics/prometheus.go
index 019279b0b..1f23dcf2d 100644
--- a/pkg/metrics/prometheus.go
+++ b/pkg/metrics/prometheus.go
@@ -310,7 +310,7 @@ func OnConfigurationUpdate(conf dynamic.Configuration, entryPoints []string) {
 func newPrometheusState() *prometheusState {
 	return &prometheusState{
 		dynamicConfig: newDynamicConfig(),
-		deletedURLs:   make(map[string]string),
+		deletedURLs:   make(map[string][]string),
 	}
 }
 
@@ -327,7 +327,7 @@ type prometheusState struct {
 	deletedEP       []string
 	deletedRouters  []string
 	deletedServices []string
-	deletedURLs     map[string]string
+	deletedURLs     map[string][]string
 }
 
 func (ps *prometheusState) SetDynamicConfig(dynamicConfig *dynamicConfig) {
@@ -350,11 +350,10 @@ func (ps *prometheusState) SetDynamicConfig(dynamicConfig *dynamicConfig) {
 		actualService, ok := dynamicConfig.services[service]
 		if !ok {
 			ps.deletedServices = append(ps.deletedServices, service)
-			continue
 		}
 		for url := range serV {
 			if _, ok := actualService[url]; !ok {
-				ps.deletedURLs[service] = url
+				ps.deletedURLs[service] = append(ps.deletedURLs[service], url)
 			}
 		}
 	}
@@ -401,16 +400,18 @@ func (ps *prometheusState) Collect(ch chan<- stdprometheus.Metric) {
 		}
 	}
 
-	for service, url := range ps.deletedURLs {
-		if !ps.dynamicConfig.hasServerURL(service, url) {
-			ps.DeletePartialMatch(map[string]string{"service": service, "url": url})
+	for service, urls := range ps.deletedURLs {
+		for _, url := range urls {
+			if !ps.dynamicConfig.hasServerURL(service, url) {
+				ps.DeletePartialMatch(map[string]string{"service": service, "url": url})
+			}
 		}
 	}
 
 	ps.deletedEP = nil
 	ps.deletedRouters = nil
 	ps.deletedServices = nil
-	ps.deletedURLs = make(map[string]string)
+	ps.deletedURLs = make(map[string][]string)
 }
 
 // DeletePartialMatch deletes all metrics where the variable labels contain all of those passed in as labels.
diff --git a/pkg/metrics/prometheus_test.go b/pkg/metrics/prometheus_test.go
index 02688dfc3..30627c3bd 100644
--- a/pkg/metrics/prometheus_test.go
+++ b/pkg/metrics/prometheus_test.go
@@ -364,6 +364,7 @@ func TestPrometheusMetricRemoval(t *testing.T) {
 				th.WithService("bar@providerName", th.WithServers(
 					th.WithServer("http://localhost:9000"),
 					th.WithServer("http://localhost:9999"),
+					th.WithServer("http://localhost:9998"),
 				)),
 				th.WithService("service1", th.WithServers(th.WithServer("http://localhost:9000"))),
 			),
@@ -403,6 +404,10 @@ func TestPrometheusMetricRemoval(t *testing.T) {
 		ServiceServerUpGauge().
 		With("service", "bar@providerName", "url", "http://localhost:9999").
 		Set(1)
+	prometheusRegistry.
+		ServiceServerUpGauge().
+		With("service", "bar@providerName", "url", "http://localhost:9998").
+		Set(1)
 
 	assertMetricsExist(t, mustScrape(), entryPointReqsTotalName, routerReqsTotalName, serviceReqsTotalName, serviceServerUpName)
 	assertMetricsAbsent(t, mustScrape(), entryPointReqsTotalName, routerReqsTotalName, serviceReqsTotalName, serviceServerUpName)
@@ -432,6 +437,47 @@ func TestPrometheusMetricRemoval(t *testing.T) {
 	assertMetricsExist(t, mustScrape(), entryPointReqsTotalName, serviceReqsTotalName, serviceServerUpName, routerReqsTotalName)
 }
 
+func TestPrometheusMetricRemoveEndpointForRecoveredService(t *testing.T) {
+	promState = newPrometheusState()
+	promRegistry = prometheus.NewRegistry()
+	t.Cleanup(promState.reset)
+
+	prometheusRegistry := RegisterPrometheus(context.Background(), &types.Prometheus{AddServicesLabels: true})
+	defer promRegistry.Unregister(promState)
+
+	conf1 := dynamic.Configuration{
+		HTTP: th.BuildConfiguration(
+			th.WithLoadBalancerServices(
+				th.WithService("service1", th.WithServers(th.WithServer("http://localhost:9000"))),
+			),
+		),
+	}
+
+	conf2 := dynamic.Configuration{
+		HTTP: th.BuildConfiguration(),
+	}
+
+	conf3 := dynamic.Configuration{
+		HTTP: th.BuildConfiguration(
+			th.WithLoadBalancerServices(
+				th.WithService("service1", th.WithServers(th.WithServer("http://localhost:9001"))),
+			),
+		),
+	}
+
+	OnConfigurationUpdate(conf1, []string{})
+	OnConfigurationUpdate(conf2, []string{})
+	OnConfigurationUpdate(conf3, []string{})
+
+	prometheusRegistry.
+		ServiceServerUpGauge().
+		With("service", "service1", "url", "http://localhost:9000").
+		Add(1)
+
+	assertMetricsExist(t, mustScrape(), serviceServerUpName)
+	assertMetricsAbsent(t, mustScrape(), serviceServerUpName)
+}
+
 func TestPrometheusRemovedMetricsReset(t *testing.T) {
 	t.Cleanup(promState.reset)
 
@@ -488,7 +534,7 @@ func (ps *prometheusState) reset() {
 	ps.deletedEP = nil
 	ps.deletedRouters = nil
 	ps.deletedServices = nil
-	ps.deletedURLs = make(map[string]string)
+	ps.deletedURLs = make(map[string][]string)
 }
 
 // Tracking and gathering the metrics happens concurrently.