Add node name to Collector struct

Signed-off-by: Matej Gera <matejgera@gmail.com>
coralogix · Dec 11, 2023 · 9f22bec · 9f22bec
1 parent 540e0c4
commit 9f22bec
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 11 deletions.
diff --git a/cmd/otel-allocator/allocation/consistent_hashing.go b/cmd/otel-allocator/allocation/consistent_hashing.go
@@ -161,7 +161,7 @@ func (c *consistentHashingAllocator) handleCollectors(diff diff.Changes[*Collect
 	}
 	// Insert the new collectors
 	for _, i := range diff.Additions() {
-		c.collectors[i.Name] = NewCollector(i.Name)
+		c.collectors[i.Name] = NewCollector(i.Name, i.Node)
 		c.consistentHasher.Add(c.collectors[i.Name])
 	}
 

diff --git a/cmd/otel-allocator/allocation/least_weighted.go b/cmd/otel-allocator/allocation/least_weighted.go
@@ -191,7 +191,7 @@ func (allocator *leastWeightedAllocator) handleCollectors(diff diff.Changes[*Col
 	}
 	// Insert the new collectors
 	for _, i := range diff.Additions() {
-		allocator.collectors[i.Name] = NewCollector(i.Name)
+		allocator.collectors[i.Name] = NewCollector(i.Name, i.Node)
 	}
 	if allocateTargets {
 		for _, item := range allocator.targetItems {

diff --git a/cmd/otel-allocator/allocation/strategy.go b/cmd/otel-allocator/allocation/strategy.go
@@ -106,6 +106,7 @@ var _ consistent.Member = Collector{}
 // This struct can be extended with information like annotations and labels in the future.
 type Collector struct {
 	Name       string
+	Node       string
 	NumTargets int
 }
 
@@ -117,8 +118,8 @@ func (c Collector) String() string {
 	return c.Name
 }
 
-func NewCollector(name string) *Collector {
-	return &Collector{Name: name}
+func NewCollector(name, node string) *Collector {
+	return &Collector{Name: name, Node: node}
 }
 
 func init() {
@@ -130,4 +131,8 @@ func init() {
 	if err != nil {
 		panic(err)
 	}
+	err = Register(perNodeStrategyName, newPerNodeAllocator)
+	if err != nil {
+		panic(err)
+	}
 }
diff --git a/cmd/otel-allocator/allocation/strategy_test.go b/cmd/otel-allocator/allocation/strategy_test.go
@@ -89,11 +89,11 @@ func Benchmark_Setting(b *testing.B) {
 }
 
 func TestCollectorDiff(t *testing.T) {
-	collector0 := NewCollector("collector-0")
-	collector1 := NewCollector("collector-1")
-	collector2 := NewCollector("collector-2")
-	collector3 := NewCollector("collector-3")
-	collector4 := NewCollector("collector-4")
+	collector0 := NewCollector("collector-0", "")
+	collector1 := NewCollector("collector-1", "")
+	collector2 := NewCollector("collector-2", "")
+	collector3 := NewCollector("collector-3", "")
+	collector4 := NewCollector("collector-4", "")
 	type args struct {
 		current map[string]*Collector
 		new     map[string]*Collector

diff --git a/cmd/otel-allocator/collector/collector.go b/cmd/otel-allocator/collector/collector.go
@@ -77,7 +77,7 @@ func (k *Client) Watch(ctx context.Context, labelMap map[string]string, fn func(
 	for i := range pods.Items {
 		pod := pods.Items[i]
 		if pod.GetObjectMeta().GetDeletionTimestamp() == nil {
-			collectorMap[pod.Name] = allocation.NewCollector(pod.Name)
+			collectorMap[pod.Name] = allocation.NewCollector(pod.Name, pod.Spec.NodeName)
 		}
 	}
 
@@ -130,7 +130,7 @@ func runWatch(ctx context.Context, k *Client, c <-chan watch.Event, collectorMap
 
 			switch event.Type { //nolint:exhaustive
 			case watch.Added:
-				collectorMap[pod.Name] = allocation.NewCollector(pod.Name)
+				collectorMap[pod.Name] = allocation.NewCollector(pod.Name, pod.Spec.NodeName)
 			case watch.Deleted:
 				delete(collectorMap, pod.Name)
 			}