From 5f8f3de317b29168b1225f8e77c49788bba11318 Mon Sep 17 00:00:00 2001
From: Jacob Aronoff <jacobaronoff45@gmail.com>
Date: Wed, 17 Aug 2022 15:10:08 -0400
Subject: [PATCH] Added fix for re-assigning targets

---
 cmd/otel-allocator/allocation/allocator.go    | 16 +++++++++
 .../allocation/allocator_test.go              | 36 +++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/cmd/otel-allocator/allocation/allocator.go b/cmd/otel-allocator/allocation/allocator.go
index 41ed6104c9..d6622cc296 100644
--- a/cmd/otel-allocator/allocation/allocator.go
+++ b/cmd/otel-allocator/allocation/allocator.go
@@ -104,6 +104,19 @@ func (allocator *Allocator) SetWaitingTargets(targets []TargetItem) {
 	//allocator.log.Info(fmt.Sprintf("len(targetsWaiting): %d\nlen(targets): %d", len(allocator.targetsWaiting), len(targets)))
 }
 
+// allCollectorsPresent checks if all of the collectors provided are in the allocator's map
+func (allocator *Allocator) allCollectorsPresent(collectors []string) bool {
+	if len(collectors) != len(allocator.collectors) {
+		return false
+	}
+	for _, s := range collectors {
+		if _, ok := allocator.collectors[s]; !ok {
+			return false
+		}
+	}
+	return true
+}
+
 // SetCollectors sets the set of collectors with key=collectorName, value=Collector object.
 // SetCollectors is called when Collectors are added or removed
 func (allocator *Allocator) SetCollectors(collectors []string) {
@@ -114,6 +127,9 @@ func (allocator *Allocator) SetCollectors(collectors []string) {
 	if len(collectors) == 0 {
 		log.Info("No collector instances present")
 		return
+	} else if allocator.allCollectorsPresent(collectors) {
+		log.Info("No changes to the collectors found")
+		return
 	}
 	for k := range allocator.collectors {
 		delete(allocator.collectors, k)
diff --git a/cmd/otel-allocator/allocation/allocator_test.go b/cmd/otel-allocator/allocation/allocator_test.go
index 596ff3b628..72109bfa1f 100644
--- a/cmd/otel-allocator/allocation/allocator_test.go
+++ b/cmd/otel-allocator/allocation/allocator_test.go
@@ -83,6 +83,42 @@ func TestAddingAndRemovingTargets(t *testing.T) {
 	}
 }
 
+func TestNoCollectorReassignment(t *testing.T) {
+	s := NewAllocator(logger)
+
+	cols := []string{"col-1", "col-2", "col-3"}
+	s.SetCollectors(cols)
+	labels := model.LabelSet{}
+
+	excpectedColLen := len(cols)
+	assert.Len(t, s.collectors, excpectedColLen)
+
+	for _, i := range cols {
+		assert.NotNil(t, s.collectors[i])
+	}
+	initTargets := []string{"prometheus:1000", "prometheus:1001", "prometheus:1002", "prometheus:1003", "prometheus:1004", "prometheus:1005"}
+	var targetList []TargetItem
+	for _, i := range initTargets {
+		targetList = append(targetList, TargetItem{JobName: "sample-name", TargetURL: i, Label: labels})
+	}
+	// test that targets and collectors are added properly
+	s.SetWaitingTargets(targetList)
+	s.AllocateTargets()
+
+	// verify
+	expectedTargetLen := len(initTargets)
+	targetItems := s.TargetItems()
+	assert.Len(t, targetItems, expectedTargetLen)
+
+	// assign new set of collectors with the same names
+	newCols := []string{"col-1", "col-2", "col-3"}
+	s.SetCollectors(newCols)
+
+	newTargetItems := s.TargetItems()
+	assert.Equal(t, targetItems, newTargetItems)
+
+}
+
 // Tests that two targets with the same target url and job name but different label set are both added
 func TestAllocationCollision(t *testing.T) {
 	// prepare allocator with initial targets and collectors