From e1bcc87627a7e6967105b23f66f8b6cca6060574 Mon Sep 17 00:00:00 2001
From: Antoine Toulme <atoulme@splunk.com>
Date: Sun, 19 Jan 2025 23:52:43 -0800
Subject: [PATCH] [chore] simplify allocation Option and registry

---
 cmd/otel-allocator/allocation/allocator.go |  2 +-
 cmd/otel-allocator/allocation/strategy.go  | 52 ++++++----------------
 cmd/otel-allocator/main.go                 |  8 +---
 3 files changed, 16 insertions(+), 46 deletions(-)

diff --git a/cmd/otel-allocator/allocation/allocator.go b/cmd/otel-allocator/allocation/allocator.go
index b0a9125ba9..4a29e231fe 100644
--- a/cmd/otel-allocator/allocation/allocator.go
+++ b/cmd/otel-allocator/allocation/allocator.go
@@ -34,7 +34,7 @@ import (
 
 var _ Allocator = &allocator{}
 
-func newAllocator(log logr.Logger, strategy Strategy, opts ...AllocationOption) Allocator {
+func newAllocator(log logr.Logger, strategy Strategy, opts ...Option) Allocator {
 	chAllocator := &allocator{
 		strategy:                      strategy,
 		collectors:                    make(map[string]*Collector),
diff --git a/cmd/otel-allocator/allocation/strategy.go b/cmd/otel-allocator/allocation/strategy.go
index 35394d0f8c..f12d63aad4 100644
--- a/cmd/otel-allocator/allocation/strategy.go
+++ b/cmd/otel-allocator/allocation/strategy.go
@@ -15,7 +15,6 @@
 package allocation
 
 import (
-	"errors"
 	"fmt"
 
 	"github.com/buraksezer/consistent"
@@ -26,12 +25,14 @@ import (
 	"github.com/open-telemetry/opentelemetry-operator/cmd/otel-allocator/target"
 )
 
-type AllocatorProvider func(log logr.Logger, opts ...AllocationOption) Allocator
+type AllocatorProvider func(log logr.Logger, opts ...Option) Allocator
 
 var (
-	strategies = map[string]Strategy{}
-
-	registry = map[string]AllocatorProvider{}
+	strategies = map[string]Strategy{
+		leastWeightedStrategyName:     newleastWeightedStrategy(),
+		consistentHashingStrategyName: newConsistentHashingStrategy(),
+		perNodeStrategyName:           newPerNodeStrategy(),
+	}
 
 	// TargetsPerCollector records how many targets have been assigned to each collector.
 	// It is currently the responsibility of the strategy to track this information.
@@ -57,19 +58,19 @@ var (
 	})
 )
 
-type AllocationOption func(Allocator)
+type Option func(Allocator)
 
 type Filter interface {
 	Apply(map[string]*target.Item) map[string]*target.Item
 }
 
-func WithFilter(filter Filter) AllocationOption {
+func WithFilter(filter Filter) Option {
 	return func(allocator Allocator) {
 		allocator.SetFilter(filter)
 	}
 }
 
-func WithFallbackStrategy(fallbackStrategy string) AllocationOption {
+func WithFallbackStrategy(fallbackStrategy string) Option {
 	var strategy, ok = strategies[fallbackStrategy]
 	if fallbackStrategy != "" && !ok {
 		panic(fmt.Errorf("unregistered strategy used as fallback: %s", fallbackStrategy))
@@ -83,24 +84,16 @@ func RecordTargetsKept(targets map[string]*target.Item) {
 	TargetsRemaining.Set(float64(len(targets)))
 }
 
-func New(name string, log logr.Logger, opts ...AllocationOption) (Allocator, error) {
-	if p, ok := registry[name]; ok {
-		return p(log.WithValues("allocator", name), opts...), nil
+func New(name string, log logr.Logger, opts ...Option) (Allocator, error) {
+	if strategy, ok := strategies[name]; ok {
+		return newAllocator(log.WithValues("allocator", name), strategy, opts...), nil
 	}
 	return nil, fmt.Errorf("unregistered strategy: %s", name)
 }
 
-func Register(name string, provider AllocatorProvider) error {
-	if _, ok := registry[name]; ok {
-		return errors.New("already registered")
-	}
-	registry[name] = provider
-	return nil
-}
-
 func GetRegisteredAllocatorNames() []string {
 	var names []string
-	for s := range registry {
+	for s := range strategies {
 		names = append(names, s)
 	}
 	return names
@@ -123,7 +116,7 @@ type Strategy interface {
 	// SetCollectors call. Strategies which don't need this information can just ignore it.
 	SetCollectors(map[string]*Collector)
 	GetName() string
-	// Add fallback strategy for strategies whose main allocation method can sometimes leave targets unassigned
+	// SetFallbackStrategy adds fallback strategy for strategies whose main allocation method can sometimes leave targets unassigned
 	SetFallbackStrategy(Strategy)
 }
 
@@ -149,20 +142,3 @@ func (c Collector) String() string {
 func NewCollector(name, node string) *Collector {
 	return &Collector{Name: name, NodeName: node}
 }
-
-func init() {
-	strategies = map[string]Strategy{
-		leastWeightedStrategyName:     newleastWeightedStrategy(),
-		consistentHashingStrategyName: newConsistentHashingStrategy(),
-		perNodeStrategyName:           newPerNodeStrategy(),
-	}
-
-	for strategyName, strategy := range strategies {
-		err := Register(strategyName, func(log logr.Logger, opts ...AllocationOption) Allocator {
-			return newAllocator(log, strategy, opts...)
-		})
-		if err != nil {
-			panic(err)
-		}
-	}
-}
diff --git a/cmd/otel-allocator/main.go b/cmd/otel-allocator/main.go
index eff7502dcd..afada36328 100644
--- a/cmd/otel-allocator/main.go
+++ b/cmd/otel-allocator/main.go
@@ -81,13 +81,7 @@ func main() {
 	log := ctrl.Log.WithName("allocator")
 
 	allocatorPrehook = prehook.New(cfg.FilterStrategy, log)
-
-	var allocationOptions []allocation.AllocationOption
-	allocationOptions = append(allocationOptions, allocation.WithFilter(allocatorPrehook))
-	if cfg.AllocationFallbackStrategy != "" {
-		allocationOptions = append(allocationOptions, allocation.WithFallbackStrategy(cfg.AllocationFallbackStrategy))
-	}
-	allocator, err = allocation.New(cfg.AllocationStrategy, log, allocationOptions...)
+	allocator, err = allocation.New(cfg.AllocationStrategy, log, allocation.WithFilter(allocatorPrehook), allocation.WithFallbackStrategy(cfg.AllocationFallbackStrategy))
 	if err != nil {
 		setupLog.Error(err, "Unable to initialize allocation strategy")
 		os.Exit(1)