From 99dfd775b35194eef367afb1c80618cf4f976ce7 Mon Sep 17 00:00:00 2001
From: Evan Huus <eapache@gmail.com>
Date: Wed, 29 Nov 2017 09:18:30 -0500
Subject: [PATCH] Unregister metrics when closing broker

Wanted to do this for a while, but upstream didn't have an unregister
method.
---
 broker.go      | 7 +++++++
 broker_test.go | 8 ++++----
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/broker.go b/broker.go
index f57a69094..923b07faf 100644
--- a/broker.go
+++ b/broker.go
@@ -178,6 +178,13 @@ func (b *Broker) Close() error {
 	b.done = nil
 	b.responses = nil
 
+	if b.id >= 0 {
+		b.conf.MetricRegistry.Unregister(getMetricNameForBroker("incoming-byte-rate", b))
+		b.conf.MetricRegistry.Unregister(getMetricNameForBroker("request-rate", b))
+		b.conf.MetricRegistry.Unregister(getMetricNameForBroker("outgoing-byte-rate", b))
+		b.conf.MetricRegistry.Unregister(getMetricNameForBroker("response-rate", b))
+	}
+
 	if err == nil {
 		Logger.Printf("Closed connection to broker %s\n", b.addr)
 	} else {
diff --git a/broker_test.go b/broker_test.go
index fcbe627fa..cc73b4440 100644
--- a/broker_test.go
+++ b/broker_test.go
@@ -77,10 +77,6 @@ func TestSimpleBrokerCommunication(t *testing.T) {
 			t.Fatal(err)
 		}
 		tt.runner(t, broker)
-		err = broker.Close()
-		if err != nil {
-			t.Error(err)
-		}
 		// Wait up to 500 ms for the remote broker to process the request and
 		// notify us about the metrics
 		timeout := 500 * time.Millisecond
@@ -91,6 +87,10 @@ func TestSimpleBrokerCommunication(t *testing.T) {
 			t.Errorf("No request received for: %s after waiting for %v", tt.name, timeout)
 		}
 		mb.Close()
+		err = broker.Close()
+		if err != nil {
+			t.Error(err)
+		}
 	}
 
 }