Fix relationship between cmd/main and pkg/driver

cmd/main.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2019 The Kubernetes Authors.
+Copyright 2024 The Kubernetes Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -18,64 +18,121 @@ package main
 
 import (
 	"context"
+	"fmt"
+	"os"
+	"strings"
 	"time"
 
-	flag "github.com/spf13/pflag"
-
+	"github.com/kubernetes-sigs/aws-ebs-csi-driver/cmd/hooks"
+	"github.com/kubernetes-sigs/aws-ebs-csi-driver/pkg/cloud"
 	"github.com/kubernetes-sigs/aws-ebs-csi-driver/pkg/driver"
 	"github.com/kubernetes-sigs/aws-ebs-csi-driver/pkg/metrics"
+	flag "github.com/spf13/pflag"
+	"k8s.io/component-base/featuregate"
 	logsapi "k8s.io/component-base/logs/api/v1"
 	json "k8s.io/component-base/logs/json"
-
 	"k8s.io/klog/v2"
 )
 
+var (
+	osExit      = os.Exit
+	featureGate = featuregate.NewFeatureGate()
+)
+
 func main() {
 	fs := flag.NewFlagSet("aws-ebs-csi-driver", flag.ExitOnError)
-
 	if err := logsapi.RegisterLogFormat(logsapi.JSONLogFormat, json.Factory{}, logsapi.LoggingBetaOptions); err != nil {
 		klog.ErrorS(err, "failed to register JSON log format")
 	}
 
-	options := GetOptions(fs)
+	var (
+		version  = fs.Bool("version", false, "Print the version and exit.")
+		toStderr = fs.Bool("logtostderr", false, "log to standard error instead of files. DEPRECATED: will be removed in a future release.")
+		args     = os.Args[1:]
+		cmd      = string(driver.AllMode)
+		options  = driver.Options{}
+	)
+
+	options.AddFlags(fs)
+
+	c := logsapi.NewLoggingConfiguration()
+	err := logsapi.AddFeatureGates(featureGate)
+	if err != nil {
+		klog.ErrorS(err, "failed to add feature gates")
+	}
+	logsapi.AddFlags(c, fs)
+
+	if len(os.Args) > 1 && !strings.HasPrefix(os.Args[1], "-") {
+		cmd = os.Args[1]
+		args = os.Args[2:]
+	}
+
+	switch cmd {
+	case "pre-stop-hook":
+		clientset, clientErr := cloud.DefaultKubernetesAPIClient()
+		if clientErr != nil {
+			klog.ErrorS(err, "unable to communicate with k8s API")
+		} else {
+			err = hooks.PreStop(clientset)
+			if err != nil {
+				klog.ErrorS(err, "failed to execute PreStop lifecycle hook")
+				klog.FlushAndExit(klog.ExitFlushTimeout, 1)
+			}
+		}
+		klog.FlushAndExit(klog.ExitFlushTimeout, 0)
+	case string(driver.ControllerMode), string(driver.NodeMode), string(driver.AllMode):
+		options.Mode = driver.Mode(cmd)
+	default:
+		klog.Errorf("Unknown driver mode %s: Expected %s, %s, %s, or pre-stop-hook", cmd, driver.ControllerMode, driver.NodeMode, driver.AllMode)
+		klog.FlushAndExit(klog.ExitFlushTimeout, 0)
+	}
+
+	if err = fs.Parse(args); err != nil {
+		panic(err)
+	}
+
+	err = logsapi.ValidateAndApply(c, featureGate)
+	if err != nil {
+		klog.ErrorS(err, "failed to validate and apply logging configuration")
+	}
+
+	if *version {
+		versionInfo, versionErr := driver.GetVersionJSON()
+		if versionErr != nil {
+			klog.ErrorS(err, "failed to get version")
+			klog.FlushAndExit(klog.ExitFlushTimeout, 1)
+		}
+		fmt.Println(versionInfo)
+		osExit(0)
+	}
+
+	if *toStderr {
+		klog.SetOutput(os.Stderr)
+	}
 
 	// Start tracing as soon as possible
-	if options.ServerOptions.EnableOtelTracing {
-		exporter, err := driver.InitOtelTracing()
-		if err != nil {
+	if options.EnableOtelTracing {
+		exporter, exporterErr := driver.InitOtelTracing()
+		if exporterErr != nil {
 			klog.ErrorS(err, "failed to initialize otel tracing")
 			klog.FlushAndExit(klog.ExitFlushTimeout, 1)
 		}
 		// Exporter will flush traces on shutdown
 		defer func() {
 			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 			defer cancel()
-			if err := exporter.Shutdown(ctx); err != nil {
+			if shutdownErr := exporter.Shutdown(ctx); shutdownErr != nil {
 				klog.ErrorS(err, "could not shutdown otel exporter")
 			}
 		}()
 	}
 
-	if options.ServerOptions.HttpEndpoint != "" {
+	if options.HttpEndpoint != "" {
 		r := metrics.InitializeRecorder()
-		r.InitializeMetricsHandler(options.ServerOptions.HttpEndpoint, "/metrics")
+		r.InitializeMetricsHandler(options.HttpEndpoint, "/metrics")
 	}
 
-	drv, err := driver.NewDriver(
-		driver.WithEndpoint(options.ServerOptions.Endpoint),
-		driver.WithExtraTags(options.ControllerOptions.ExtraTags),
-		driver.WithExtraVolumeTags(options.ControllerOptions.ExtraVolumeTags),
-		driver.WithMode(options.DriverMode),
-		driver.WithVolumeAttachLimit(options.NodeOptions.VolumeAttachLimit),
-		driver.WithReservedVolumeAttachments(options.NodeOptions.ReservedVolumeAttachments),
-		driver.WithKubernetesClusterID(options.ControllerOptions.KubernetesClusterID),
-		driver.WithAwsSdkDebugLog(options.ControllerOptions.AwsSdkDebugLog),
-		driver.WithWarnOnInvalidTag(options.ControllerOptions.WarnOnInvalidTag),
-		driver.WithUserAgentExtra(options.ControllerOptions.UserAgentExtra),
-		driver.WithOtelTracing(options.ServerOptions.EnableOtelTracing),
-		driver.WithBatching(options.ControllerOptions.Batching),
-		driver.WithModifyVolumeRequestHandlerTimeout(options.ControllerOptions.ModifyVolumeRequestHandlerTimeout),
-	)
+	drv, err := driver.NewDriver(&options)
 	if err != nil {
 		klog.ErrorS(err, "failed to create driver")
 		klog.FlushAndExit(klog.ExitFlushTimeout, 1)