From ed62b86961453c1f1e2d9cbd3da43e81cc63b311 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miko=C5=82aj=20=C5=9Awi=C4=85tek?=
 <mail+sumo@mikolajswiatek.com>
Date: Wed, 16 Aug 2023 18:49:15 +0200
Subject: [PATCH] Fix VolumeClaimTemplate in Collector CRD

---
 cmd/otel-allocator/config/config.go | 41 ++++++++++++++++++++++-------
 cmd/otel-allocator/config/flags.go  | 37 +++++++++++++++-----------
 cmd/otel-allocator/main.go          |  4 ++-
 3 files changed, 57 insertions(+), 25 deletions(-)

diff --git a/cmd/otel-allocator/config/config.go b/cmd/otel-allocator/config/config.go
index 0d410b5c14..85f9ddee3f 100644
--- a/cmd/otel-allocator/config/config.go
+++ b/cmd/otel-allocator/config/config.go
@@ -75,14 +75,18 @@ func LoadFromFile(file string, target *Config) error {
 	return unmarshal(target, file)
 }
 
-func LoadFromCLI(target *Config) error {
+func LoadFromCLI(target *Config, flagSet *pflag.FlagSet) error {
+	var err error
 	// set the rest of the config attributes based on command-line flag values
 	target.RootLogger = zap.New(zap.UseFlagOptions(&zapCmdLineOpts))
 	klog.SetLogger(target.RootLogger)
 	ctrl.SetLogger(target.RootLogger)
 
-	target.KubeConfigFilePath = *kubeConfigPathFlag
-	clusterConfig, err := clientcmd.BuildConfigFromFlags("", *kubeConfigPathFlag)
+	target.KubeConfigFilePath, err = flagSet.GetString(kubeConfigPathFlagName)
+	if err != nil {
+		return err
+	}
+	clusterConfig, err := clientcmd.BuildConfigFromFlags("", target.KubeConfigFilePath)
 	if err != nil {
 		pathError := &fs.PathError{}
 		if ok := errors.As(err, &pathError); !ok {
@@ -93,10 +97,18 @@ func LoadFromCLI(target *Config) error {
 			return err
 		}
 	}
-	target.ListenAddr = *listenAddrFlag
-	target.PrometheusCR.Enabled = *prometheusCREnabledFlag
 	target.ClusterConfig = clusterConfig
 
+	target.ListenAddr, err = flagSet.GetString(target.ListenAddr)
+	if err != nil {
+		return err
+	}
+
+	target.PrometheusCR.Enabled, err = flagSet.GetBool(prometheusCREnabledFlagName)
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -121,21 +133,32 @@ func CreateDefaultConfig() Config {
 }
 
 func Load() (*Config, string, error) {
-	pflag.Parse()
+	var err error
+
+	flagSet := getFlagSet(pflag.ExitOnError)
+	err = flagSet.Parse(os.Args)
+	if err != nil {
+		return nil, "", err
+	}
+
 	config := CreateDefaultConfig()
 
 	// load the config from the config file
-	err := LoadFromFile(*configFilePathFlag, &config)
+	configFilePath, err := getConfigFilePath(flagSet)
+	if err != nil {
+		return nil, "", err
+	}
+	err = LoadFromFile(configFilePath, &config)
 	if err != nil {
 		return nil, "", err
 	}
 
-	err = LoadFromCLI(&config)
+	err = LoadFromCLI(&config, flagSet)
 	if err != nil {
 		return nil, "", err
 	}
 
-	return &config, *configFilePathFlag, nil
+	return &config, configFilePath, nil
 }
 
 // ValidateConfig validates the cli and file configs together.
diff --git a/cmd/otel-allocator/config/flags.go b/cmd/otel-allocator/config/flags.go
index c788bb223b..2f965457ae 100644
--- a/cmd/otel-allocator/config/flags.go
+++ b/cmd/otel-allocator/config/flags.go
@@ -23,23 +23,30 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 )
 
-// Flag definitions.
-var (
-	configFilePathFlag      *string
-	listenAddrFlag          *string
-	prometheusCREnabledFlag *bool
-	kubeConfigPathFlag      *string
-	zapCmdLineOpts          zap.Options
+// Flag names.
+const (
+	targetAllocatorName         = "target-allocator"
+	configFilePathFlagName      = "config-file"
+	listenAddrFlagName          = "listen-addr"
+	prometheusCREnabledFlagName = "enable-prometheus-cr-watcher"
+	kubeConfigPathFlagName      = "kubeconfig-path"
 )
 
-func initFlags() {
-	configFilePathFlag = pflag.String("config-file", DefaultConfigFilePath, "The path to the config file.")
-	listenAddrFlag = pflag.String("listen-addr", ":8080", "The address where this service serves.")
-	prometheusCREnabledFlag = pflag.Bool("enable-prometheus-cr-watcher", false, "Enable Prometheus CRs as target sources")
-	kubeConfigPathFlag = pflag.String("kubeconfig-path", filepath.Join(homedir.HomeDir(), ".kube", "config"), "absolute path to the KubeconfigPath file")
-	zapCmdLineOpts.BindFlags(flag.CommandLine)
+// We can't bind this flag to our FlagSet, so we need to handle it separately.
+var zapCmdLineOpts zap.Options
+
+func getFlagSet(errorHandling pflag.ErrorHandling) *pflag.FlagSet {
+	flagSet := pflag.NewFlagSet(targetAllocatorName, errorHandling)
+	flagSet.String(configFilePathFlagName, DefaultConfigFilePath, "The path to the config file.")
+	flagSet.String(listenAddrFlagName, ":8080", "The address where this service serves.")
+	flagSet.Bool(prometheusCREnabledFlagName, false, "Enable Prometheus CRs as target sources")
+	flagSet.String(kubeConfigPathFlagName, filepath.Join(homedir.HomeDir(), ".kube", "config"), "absolute path to the KubeconfigPath file")
+	zapFlagSet := flag.NewFlagSet("", flag.ErrorHandling(errorHandling))
+	zapCmdLineOpts.BindFlags(zapFlagSet)
+	flagSet.AddGoFlagSet(zapFlagSet)
+	return flagSet
 }
 
-func init() {
-	initFlags()
+func getConfigFilePath(flagSet *pflag.FlagSet) (string, error) {
+	return flagSet.GetString(configFilePathFlagName)
 }
diff --git a/cmd/otel-allocator/main.go b/cmd/otel-allocator/main.go
index 6b237c62c3..ce9f66a641 100644
--- a/cmd/otel-allocator/main.go
+++ b/cmd/otel-allocator/main.go
@@ -16,6 +16,7 @@ package main
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"os/signal"
 	"syscall"
@@ -66,9 +67,10 @@ func main() {
 	)
 	cfg, configFilePath, err := config.Load()
 	if err != nil {
-		setupLog.Error(err, "Failed to parse parameters")
+		fmt.Printf("Failed to load config: %v", err)
 		os.Exit(1)
 	}
+	ctrl.SetLogger(cfg.RootLogger)
 
 	if validationErr := config.ValidateConfig(cfg); validationErr != nil {
 		setupLog.Error(validationErr, "Invalid configuration")