From cabce21b705669cacd356c42ad90465a0d39ff6e Mon Sep 17 00:00:00 2001
From: Krasi Georgiev <krasi.root@gmail.com>
Date: Fri, 12 Jan 2018 13:10:59 +0000
Subject: [PATCH] delete empty targets sets to avoid memory leaks

---
 discovery/manager.go      | 10 +++++++-
 discovery/manager_test.go | 51 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/discovery/manager.go b/discovery/manager.go
index 8110963dd..4e583dcde 100644
--- a/discovery/manager.go
+++ b/discovery/manager.go
@@ -179,8 +179,16 @@ func (m *Manager) allGroups() map[string][]*targetgroup.Group {
 	m.actionCh <- func(ctx context.Context) {
 		tSetsAll := map[string][]*targetgroup.Group{}
 		for pkey, tsets := range m.targets {
+			del := true
 			for _, tg := range tsets {
-				tSetsAll[pkey.setName] = append(tSetsAll[pkey.setName], tg)
+				if len(tg.Targets) != 0 {
+					tSetsAll[pkey.setName] = append(tSetsAll[pkey.setName], tg)
+					del = false
+				}
+			}
+			// Delete the empty map for this target set to avoid memory leaks.
+			if del {
+				delete(m.targets, pkey)
 			}
 		}
 		tSets <- tSetsAll
diff --git a/discovery/manager_test.go b/discovery/manager_test.go
index f3ecf784e..b7ea2457f 100644
--- a/discovery/manager_test.go
+++ b/discovery/manager_test.go
@@ -649,6 +649,57 @@ func TestDiscoveryManagerSyncCalls(t *testing.T) {
 				},
 			},
 		},
+		{
+			title: "Single TP update with an empty group to check for memory leaks",
+			updates: map[string][]update{
+				"tp1": {
+					{
+						targetGroups: []targetgroup.Group{
+							{
+								Source:  "tp1_group1",
+								Targets: []model.LabelSet{{"__instance__": "1"}},
+							},
+							{
+								Source:  "tp1_group2",
+								Targets: []model.LabelSet{{"__instance__": "2"}},
+							},
+						},
+						interval: 30,
+					},
+					{
+						targetGroups: []targetgroup.Group{
+							{
+								Source:  "tp1_group1",
+								Targets: []model.LabelSet{{"__instance__": "3"}},
+							},
+							{
+								Source:  "tp1_group2",
+								Targets: nil,
+							},
+						},
+						interval: 10,
+					},
+				},
+			},
+			expectedTargets: [][]*targetgroup.Group{
+				{
+					{
+						Source:  "tp1_group1",
+						Targets: []model.LabelSet{{"__instance__": "1"}},
+					},
+					{
+						Source:  "tp1_group2",
+						Targets: []model.LabelSet{{"__instance__": "2"}},
+					},
+				},
+				{
+					{
+						Source:  "tp1_group1",
+						Targets: []model.LabelSet{{"__instance__": "3"}},
+					},
+				},
+			},
+		},
 	}
 
 	for testIndex, testCase := range testCases {