main.go

/*
 Copyright 2023 The Kapacity Authors.

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
*/

package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	rotatelogs "github.com/lestrrat-go/file-rotatelogs"
	promapi "github.com/prometheus/client_golang/api"
	"go.uber.org/zap/zapcore"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/client-go/discovery"
	"k8s.io/client-go/dynamic"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/rest"
	scaleclient "k8s.io/client-go/scale"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/klog/v2"
	resourceclient "k8s.io/metrics/pkg/client/clientset/versioned/typed/metrics/v1beta1"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/cache"
	"sigs.k8s.io/controller-runtime/pkg/client"
	ctrlconfigv1alpha1 "sigs.k8s.io/controller-runtime/pkg/config/v1alpha1"
	"sigs.k8s.io/controller-runtime/pkg/event"
	"sigs.k8s.io/controller-runtime/pkg/healthz"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	"sigs.k8s.io/controller-runtime/pkg/manager"

	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
	// to ensure that exec-entrypoint and run can make use of them.
	_ "k8s.io/client-go/plugin/pkg/client/auth"

	autoscalingv1alpha1 "github.com/traas-stack/kapacity/apis/autoscaling/v1alpha1"
	"github.com/traas-stack/kapacity/controllers/autoscaling"
	internalgrpc "github.com/traas-stack/kapacity/internal/grpc"
	"github.com/traas-stack/kapacity/internal/webhook"
	metricprovider "github.com/traas-stack/kapacity/pkg/metric/provider"
	"github.com/traas-stack/kapacity/pkg/metric/provider/metricsapi"
	"github.com/traas-stack/kapacity/pkg/metric/provider/prometheus"
	metricservice "github.com/traas-stack/kapacity/pkg/metric/service"
	"github.com/traas-stack/kapacity/pkg/portrait/algorithm/externaljob/jobcontroller"
	"github.com/traas-stack/kapacity/pkg/portrait/algorithm/externaljob/resultfetcher"
	portraitgenerator "github.com/traas-stack/kapacity/pkg/portrait/generator"
	"github.com/traas-stack/kapacity/pkg/portrait/generator/reactive"
	portraitprovider "github.com/traas-stack/kapacity/pkg/portrait/provider"
	"github.com/traas-stack/kapacity/pkg/scale"
	"github.com/traas-stack/kapacity/pkg/util"
	//+kubebuilder:scaffold:imports
)

const (
	defaultQPS   = float64(20)
	defaultBurst = 30
)

var (
	scheme   = runtime.NewScheme()
	setupLog = ctrl.Log.WithName("setup")
)

func init() {
	utilruntime.ReallyCrash = false

	utilruntime.Must(clientgoscheme.AddToScheme(scheme))

	utilruntime.Must(autoscalingv1alpha1.AddToScheme(scheme))
	//+kubebuilder:scaffold:scheme
}

type promConfig struct {
	Address                              string
	AuthInCluster                        bool
	AuthConfigFile                       string
	MetricsConfigFile                    string
	MetricsRelistInterval, MetricsMaxAge time.Duration
}

type algorithmJobConfig struct {
	Namespace                        string
	DefaultServiceAccount            string
	DefaultMetricsServerAddr         string
	DefaultImageHorizontalPredictive string
}

func main() {
	var (
		logPath             string
		logFileMaxAge       time.Duration
		logFileRotationTime time.Duration

		metricsAddr          string
		probeAddr            string
		enableLeaderElection bool
		kubeAPIQPS           float64
		kubeAPIBurst         int
		reconcileConcurrency int
		objectLabelSelector  string
		objectFieldSelector  string

		enableAdmissionWebhookServer bool

		grpcServerAddr string

		metricProviderType string
		promConfig         promConfig

		algorithmJobConfig algorithmJobConfig
	)
	flag.StringVar(&logPath, "log-path", "", "The path to write log files. Omit to disable logging to file.")
	flag.DurationVar(&logFileMaxAge, "log-file-max-age", 7*24*time.Hour, "The max age of the log files.")
	flag.DurationVar(&logFileRotationTime, "log-file-rotation-time", 24*time.Hour, "The rotation time of the log file.")
	flag.StringVar(&metricsAddr, "metrics-bind-address", ":8080", "The address the metric endpoint binds to.")
	flag.StringVar(&probeAddr, "health-probe-bind-address", ":8081", "The address the probe endpoint binds to.")
	flag.BoolVar(&enableLeaderElection, "leader-elect", false,
		"Enable leader election for controller manager. "+
			"Enabling this will ensure there is only one active controller manager.")
	flag.Float64Var(&kubeAPIQPS, "kube-api-qps", defaultQPS, "QPS to use while talking with kubernetes apiserver."+
		fmt.Sprintf("The number must be >= 0. If 0 will use defaultQPS: %f.", defaultQPS))
	flag.IntVar(&kubeAPIBurst, "kube-api-burst", defaultBurst, "Burst to use while talking with kubernetes apiserver."+
		fmt.Sprintf("The number must be >= 0. If 0 will use defaultBurst: %d.", defaultBurst))
	flag.IntVar(&reconcileConcurrency, "reconcile-concurrency", 1, "The reconciliation concurrency of each controller.")
	flag.StringVar(&objectLabelSelector, "object-label-selector", "", "The label selector to restrict controllers' list watch for objects.")
	flag.StringVar(&objectFieldSelector, "object-field-selector", "", "The field selector to restrict controllers' list watch for objects.")
	flag.BoolVar(&enableAdmissionWebhookServer, "serve-admission-webhooks", true, "Enable admission webhook servers.")
	flag.StringVar(&grpcServerAddr, "grpc-server-bind-address", ":9090",
		"The address the gRPC server binds to. Set 0 to disable the gRPC server.")
	flag.StringVar(&metricProviderType, "metric-provider", "prometheus",
		"The name of metric provider. Valid options are prometheus or metrics-api. Defaults to prometheus.")
	flag.StringVar(&promConfig.Address, "prometheus-address", "", "The address of the Prometheus to connect to.")
	flag.BoolVar(&promConfig.AuthInCluster, "prometheus-auth-incluster", false,
		"Use auth details from the in-cluster kubeconfig when connecting to prometheus.")
	flag.StringVar(&promConfig.AuthConfigFile, "prometheus-auth-config", "",
		"The kubeconfig file used to configure auth when connecting to Prometheus.")
	flag.StringVar(&promConfig.MetricsConfigFile, "prometheus-metrics-config", "",
		"The configuration file containing details of how to transform between Prometheus metrics and metrics API resources.")
	flag.DurationVar(&promConfig.MetricsRelistInterval, "prometheus-metrics-relist-interval", 10*time.Minute,
		"The interval at which to re-list the set of all available metrics from Prometheus.")
	flag.DurationVar(&promConfig.MetricsMaxAge, "prometheus-metrics-max-age", 0,
		"The period for which to query the set of available metrics from Prometheus. If not set, it defaults to prometheus-metrics-relist-interval.")
	flag.StringVar(&algorithmJobConfig.Namespace, "algorithm-job-namespace", "", "The namespace where algorithm jobs should run in.")
	flag.StringVar(&algorithmJobConfig.DefaultServiceAccount, "algorithm-job-default-service-account", "",
		"The default service account for algorithm jobs.")
	flag.StringVar(&algorithmJobConfig.DefaultMetricsServerAddr, "algorithm-job-default-metrics-server-addr", "",
		"The default metrics server address to inject to env of algorithm jobs.")
	flag.StringVar(&algorithmJobConfig.DefaultImageHorizontalPredictive, "algorithm-job-default-image-horizontal-predictive", "",
		"The default image for algorithm jobs of horizontal predictive portrait.")
	opts := zap.Options{
		Development: true,
		TimeEncoder: zapcore.RFC3339TimeEncoder,
	}
	opts.BindFlags(flag.CommandLine)
	flag.Parse()

	if logPath != "" {
		logFileName := logPath + "/kapacity-manager.log"
		rotateLogWriter, err := rotatelogs.New(
			logFileName+".%Y-%m-%d-%H",
			rotatelogs.WithLinkName(logFileName),
			rotatelogs.WithMaxAge(logFileMaxAge),
			rotatelogs.WithRotationTime(logFileRotationTime),
		)
		if err != nil {
			setupLog.Error(err, "failed to new rotate log writer")
			os.Exit(1)
		}
		opts.DestWriter = io.MultiWriter(os.Stderr, rotateLogWriter)
	}
	logger := zap.New(zap.UseFlagOptions(&opts))
	ctrl.SetLogger(logger)
	klog.SetLogger(logger)

	ctx := ctrl.SetupSignalHandler()

	cfg := ctrl.GetConfigOrDie()
	rest.AddUserAgent(cfg, util.UserAgent)
	if kubeAPIQPS <= 0 {
		kubeAPIQPS = defaultQPS
	}
	cfg.QPS = float32(kubeAPIQPS)
	if kubeAPIBurst <= 0 {
		kubeAPIBurst = defaultBurst
	}
	cfg.Burst = kubeAPIBurst

	var (
		cacheObjectSelector = cache.ObjectSelector{}
		err                 error
	)
	if objectLabelSelector != "" {
		cacheObjectSelector.Label, err = labels.Parse(objectLabelSelector)
		if err != nil {
			setupLog.Error(err, "failed to parse object label selector", "selector", objectLabelSelector)
			os.Exit(1)
		}
	}
	if objectFieldSelector != "" {
		cacheObjectSelector.Field, err = fields.ParseSelector(objectFieldSelector)
		if err != nil {
			setupLog.Error(err, "failed to parse object field selector", "selector", objectFieldSelector)
			os.Exit(1)
		}
	}

	mgr, err := ctrl.NewManager(cfg, ctrl.Options{
		Scheme:                 scheme,
		MetricsBindAddress:     metricsAddr,
		Port:                   9443,
		HealthProbeBindAddress: probeAddr,
		LeaderElection:         enableLeaderElection,
		LeaderElectionID:       "kapacity-manager.kapacitystack.io",
		// LeaderElectionReleaseOnCancel defines if the leader should step down voluntarily
		// when the Manager ends. This requires the binary to immediately end when the
		// Manager is stopped, otherwise, this setting is unsafe. Setting this significantly
		// speeds up voluntary leader transitions as the new leader don't have to wait
		// LeaseDuration time first.
		LeaderElectionReleaseOnCancel: true,
		Controller: ctrlconfigv1alpha1.ControllerConfigurationSpec{
			GroupKindConcurrency: map[string]int{
				"ReplicaProfile.autoscaling.kapacitystack.io":                     reconcileConcurrency,
				"IntelligentHorizontalPodAutoscaler.autoscaling.kapacitystack.io": reconcileConcurrency,
				"HorizontalPortrait.autoscaling.kapacitystack.io":                 reconcileConcurrency,
			},
		},
		NewCache: func(config *rest.Config, opts cache.Options) (cache.Cache, error) {
			opts.DefaultSelector = cacheObjectSelector
			return cache.New(config, opts)
		},
	})
	if err != nil {
		setupLog.Error(err, "unable to build manager")
		os.Exit(1)
	}

	dynamicClient, err := dynamic.NewForConfig(cfg)
	if err != nil {
		setupLog.Error(err, "unable to build dynamic client")
		os.Exit(1)
	}

	discoveryClient, err := discovery.NewDiscoveryClientForConfig(cfg)
	if err != nil {
		setupLog.Error(err, "unable to build discovery client")
		os.Exit(1)
	}

	scalesGetter, err := scaleclient.NewForConfig(cfg, mgr.GetRESTMapper(), dynamic.LegacyAPIPathResolverFunc, scaleclient.NewDiscoveryScaleKindResolver(discoveryClient))
	if err != nil {
		setupLog.Error(err, "unable to build scale client")
		os.Exit(1)
	}
	scaler := scale.NewScaler(scalesGetter, mgr.GetRESTMapper())

	ihpaReconcilerEventTrigger := make(chan event.GenericEvent, 1024)
	hpReconcilerEventTrigger := make(chan event.GenericEvent, 1024)

	var metricProvider metricprovider.Interface
	switch metricProviderType {
	case "prometheus":
		metricProvider, err = newPrometheusProviderFromConfig(mgr.GetClient(), dynamicClient, promConfig)
	case "metrics-api":
		metricProvider, err = newMetricsAPIProviderFromConfig(cfg)
	default:
		err = fmt.Errorf("unknown metric provider type %q", metricProviderType)
	}
	if err != nil {
		setupLog.Error(err, "unable to build metric provider", "metricProviderType", metricProviderType)
		os.Exit(1)
	}

	externalHorizontalPortraitAlgorithmJobControllers, err := initExternalHorizontalPortraitAlgorithmJobControllers(mgr.GetClient(), algorithmJobConfig)
	if err != nil {
		setupLog.Error(err, "unable to init external horizontal portrait algorithm job controllers")
		os.Exit(1)
	}
	externalHorizontalPortraitAlgorithmJobResultFetchers, err := initExternalHorizontalPortraitAlgorithmJobResultFetchers(ctx, mgr.GetClient(), mgr.GetCache(), hpReconcilerEventTrigger)
	if err != nil {
		setupLog.Error(err, "unable to init external horizontal portrait algorithm job result fetchers")
		os.Exit(1)
	}

	if enableAdmissionWebhookServer {
		if err := webhook.SetupWithManager(mgr); err != nil {
			setupLog.Error(err, "unable to set up webhook server")
			os.Exit(1)
		}
	}

	if err := (&autoscaling.ReplicaProfileReconciler{
		Client:        mgr.GetClient(),
		Scheme:        mgr.GetScheme(),
		EventRecorder: mgr.GetEventRecorderFor(autoscaling.ReplicaProfileControllerName),
		Scaler:        scaler,
	}).SetupWithManager(ctx, mgr); err != nil {
		setupLog.Error(err, "unable to create controller", "controller", "ReplicaProfile")
		os.Exit(1)
	}
	if err := (&autoscaling.IntelligentHorizontalPodAutoscalerReconciler{
		Client:            mgr.GetClient(),
		Scheme:            mgr.GetScheme(),
		EventRecorder:     mgr.GetEventRecorderFor(autoscaling.IHPAControllerName),
		EventTrigger:      ihpaReconcilerEventTrigger,
		PortraitProviders: initHorizontalPortraitProviders(mgr.GetClient(), ihpaReconcilerEventTrigger),
	}).SetupWithManager(mgr); err != nil {
		setupLog.Error(err, "unable to create controller", "controller", "IntelligentHorizontalPodAutoscaler")
		os.Exit(1)
	}
	if err := (&autoscaling.HorizontalPortraitReconciler{
		Client:                             mgr.GetClient(),
		Scheme:                             mgr.GetScheme(),
		EventRecorder:                      mgr.GetEventRecorderFor(autoscaling.HorizontalPortraitControllerName),
		EventTrigger:                       hpReconcilerEventTrigger,
		PortraitGenerators:                 initPortraitGenerators(mgr.GetClient(), metricProvider, scaler),
		ExternalAlgorithmJobControllers:    externalHorizontalPortraitAlgorithmJobControllers,
		ExternalAlgorithmJobResultFetchers: externalHorizontalPortraitAlgorithmJobResultFetchers,
	}).SetupWithManager(mgr); err != nil {
		setupLog.Error(err, "unable to create controller", "controller", "HorizontalPortrait")
		os.Exit(1)
	}
	//+kubebuilder:scaffold:builder

	if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
		setupLog.Error(err, "unable to set up health check")
		os.Exit(1)
	}
	if err := mgr.AddReadyzCheck("readyz", healthz.Ping); err != nil {
		setupLog.Error(err, "unable to set up ready check")
		os.Exit(1)
	}

	if grpcServerAddr != "" && grpcServerAddr != "0" {
		server := internalgrpc.NewServer(grpcServerAddr)
		metricProviderServer := metricservice.NewProviderServer(metricProvider)
		metricProviderServer.RegisterTo(server.ServiceRegistrar())
		if err := mgr.Add(server); err != nil {
			setupLog.Error(err, "unable to set up gRPC server")
			os.Exit(1)
		}
	}

	if runnable, ok := metricProvider.(manager.Runnable); ok {
		if err := mgr.Add(runnable); err != nil {
			setupLog.Error(err, "failed to add runnable metric provider to manager",
				"metricProviderType", metricProviderType)
			os.Exit(1)
		}
	}

	for sourceType, fetcher := range externalHorizontalPortraitAlgorithmJobResultFetchers {
		if runnable, ok := fetcher.(manager.Runnable); ok {
			if err := mgr.Add(runnable); err != nil {
				setupLog.Error(err, "failed to add runnable external horizontal portrait algorithm job result fetcher to manager",
					"sourceType", sourceType)
				os.Exit(1)
			}
		}
	}

	setupLog.Info("starting manager")
	if err := mgr.Start(ctx); err != nil {
		setupLog.Error(err, "problem running manager")
		os.Exit(1)
	}
}

func newPrometheusProviderFromConfig(kubeClient client.Client, kubeDynamicClient dynamic.Interface, config promConfig) (metricprovider.Interface, error) {
	if config.MetricsMaxAge == 0 {
		config.MetricsMaxAge = config.MetricsRelistInterval
	}

	promClient, err := buildPrometheusClient(config)
	if err != nil {
		return nil, fmt.Errorf("failed to build prometheus client: %v", err)
	}

	metricsConfig, err := prometheus.MetricsConfigFromFile(config.MetricsConfigFile)
	if err != nil {
		return nil, fmt.Errorf("failed to load metrics config from file %q: %v", config.MetricsConfigFile, err)
	}

	return prometheus.NewMetricProvider(kubeClient, kubeDynamicClient, promClient, metricsConfig, config.MetricsRelistInterval, config.MetricsMaxAge)
}

func buildPrometheusClient(config promConfig) (promapi.Client, error) {
	if config.AuthInCluster && config.AuthConfigFile != "" {
		return nil, fmt.Errorf("may not use both in-cluster auth and an explicit kubeconfig at the same time")
	}
	var (
		rt http.RoundTripper = &http.Transport{
			Proxy: http.ProxyFromEnvironment,
		}
		authConfig *rest.Config
		err        error
	)
	if config.AuthInCluster {
		authConfig, err = rest.InClusterConfig()
		if err != nil {
			return nil, fmt.Errorf("failed to build in-cluster auth config: %v", err)
		}
	} else if config.AuthConfigFile != "" {
		authConfig, err = clientcmd.NewNonInteractiveDeferredLoadingClientConfig(
			&clientcmd.ClientConfigLoadingRules{ExplicitPath: config.AuthConfigFile},
			&clientcmd.ConfigOverrides{}).ClientConfig()
		if err != nil {
			return nil, fmt.Errorf("failed to build auth config from kubeconfig %q: %v", config.AuthConfigFile, err)
		}
	}
	if authConfig != nil {
		rt, err = rest.TransportFor(authConfig)
		if err != nil {
			return nil, fmt.Errorf("failed to build transport from auth config: %v", err)
		}
	}
	return promapi.NewClient(promapi.Config{
		Address:      config.Address,
		RoundTripper: rt,
	})
}

func newMetricsAPIProviderFromConfig(kubeConfig *rest.Config) (metricprovider.Interface, error) {
	resourceMetricsClient, err := resourceclient.NewForConfig(kubeConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to build resource metrics client: %v", err)
	}
	return metricsapi.NewMetricProvider(resourceMetricsClient), nil
}

func initHorizontalPortraitProviders(client client.Client, eventTrigger chan event.GenericEvent) map[autoscalingv1alpha1.HorizontalPortraitProviderType]portraitprovider.Horizontal {
	providers := make(map[autoscalingv1alpha1.HorizontalPortraitProviderType]portraitprovider.Horizontal)
	providers[autoscalingv1alpha1.StaticHorizontalPortraitProviderType] = portraitprovider.NewStaticHorizontal()
	providers[autoscalingv1alpha1.CronHorizontalPortraitProviderType] = portraitprovider.NewCronHorizontal(eventTrigger)
	providers[autoscalingv1alpha1.DynamicHorizontalPortraitProviderType] = portraitprovider.NewDynamicHorizontal(client, eventTrigger)
	return providers
}

func initPortraitGenerators(client client.Client, metricProvider metricprovider.Interface, scaler *scale.Scaler) map[autoscalingv1alpha1.PortraitType]portraitgenerator.Interface {
	generators := make(map[autoscalingv1alpha1.PortraitType]portraitgenerator.Interface)
	generators[autoscalingv1alpha1.ReactivePortraitType] = reactive.NewPortraitGenerator(metricProvider, util.NewCtrlPodLister(client), scaler)
	return generators
}

func initExternalHorizontalPortraitAlgorithmJobControllers(client client.Client, config algorithmJobConfig) (map[autoscalingv1alpha1.PortraitAlgorithmJobType]jobcontroller.Horizontal, error) {
	if config.Namespace == "" {
		return nil, fmt.Errorf("shall set namespace for algorithm jobs")
	}
	controllers := make(map[autoscalingv1alpha1.PortraitAlgorithmJobType]jobcontroller.Horizontal)
	controllers[autoscalingv1alpha1.CronJobPortraitAlgorithmJobType] = jobcontroller.NewCronJobHorizontal(client, config.Namespace, config.DefaultServiceAccount, config.DefaultMetricsServerAddr, map[autoscalingv1alpha1.PortraitType]string{
		autoscalingv1alpha1.PredictivePortraitType: config.DefaultImageHorizontalPredictive,
	})
	return controllers, nil
}

func initExternalHorizontalPortraitAlgorithmJobResultFetchers(ctx context.Context, client client.Client, cache cache.Cache, eventTrigger chan event.GenericEvent) (map[autoscalingv1alpha1.PortraitAlgorithmResultSourceType]resultfetcher.Horizontal, error) {
	fetchers := make(map[autoscalingv1alpha1.PortraitAlgorithmResultSourceType]resultfetcher.Horizontal)
	configMapInformer, err := cache.GetInformer(ctx, &corev1.ConfigMap{})
	if err != nil {
		return nil, fmt.Errorf("failed to get ConfigMap informer from cache: %v", err)
	}
	fetchers[autoscalingv1alpha1.ConfigMapPortraitAlgorithmResultSourceType] = resultfetcher.NewConfigMapHorizontal(client, eventTrigger, configMapInformer)
	return fetchers, nil
}