From 8116fa9f2b1a6d310d5efa3531649048add6fb9f Mon Sep 17 00:00:00 2001
From: CYJiang <86391540+googs1025@users.noreply.github.com>
Date: Wed, 15 May 2024 11:58:57 +0800
Subject: [PATCH] add restconfig flag: qps burst (#32)

Use flags param to set QPS and Burst values for kubeconfig, enabling
more flexibility in config

Signed-off-by: googs1025 <googs1025@gmail.com>
---
 cmd/knavigator/main.go       | 13 +++++++++++--
 pkg/config/config.go         |  7 +++++++
 pkg/utils/k8s_config.go      | 18 +++++++++++-------
 pkg/utils/k8s_config_test.go | 10 +++++++++-
 4 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/cmd/knavigator/main.go b/cmd/knavigator/main.go
index 75da8e3..3617455 100644
--- a/cmd/knavigator/main.go
+++ b/cmd/knavigator/main.go
@@ -32,9 +32,13 @@ import (
 
 func mainInternal() error {
 	var kubeConfigPath, kubeCtx, taskConfigs string
+	var qps float64
+	var burst int
 	flag.StringVar(&kubeConfigPath, "kubeconfig", "", "kubeconfig file path")
 	flag.StringVar(&kubeCtx, "kubectx", "", "kube context")
 	flag.StringVar(&taskConfigs, "tasks", "", "comma-separated list of task config files and dirs")
+	flag.Float64Var(&qps, "kube-api-qps", 500, "Maximum QPS to use while talking with Kubernetes API")
+	flag.IntVar(&burst, "kube-api-burst", 500, "Maximum burst for throttle while talking with Kubernetes API")
 
 	klog.InitFlags(nil)
 	flag.Parse()
@@ -53,8 +57,13 @@ func mainInternal() error {
 	}
 
 	log := textlogger.NewLogger(textlogger.NewConfig(textlogger.Verbosity(utils.Flag2Verbosity(flag.Lookup("v")))))
-
-	restConfig, err := utils.GetK8sConfig(log, kubeConfigPath, kubeCtx)
+	cfg := &config.KubeConfig{
+		KubeConfigPath: kubeConfigPath,
+		KubeCtx:        kubeCtx,
+		QPS:            float32(qps),
+		Burst:          burst,
+	}
+	restConfig, err := utils.GetK8sConfig(log, cfg)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 956f87c..d5f9480 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -25,6 +25,13 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
+type KubeConfig struct {
+	KubeConfigPath string
+	KubeCtx        string
+	QPS            float32
+	Burst          int
+}
+
 type TaskConfig struct {
 	Name        string  `yaml:"name"`
 	Description string  `yaml:"description,omitempty"`
diff --git a/pkg/utils/k8s_config.go b/pkg/utils/k8s_config.go
index 2f8170e..26dffb0 100644
--- a/pkg/utils/k8s_config.go
+++ b/pkg/utils/k8s_config.go
@@ -24,12 +24,16 @@ import (
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
 	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
+
+	cfg "github.com/NVIDIA/knavigator/pkg/config"
 )
 
-func GetK8sConfig(log logr.Logger, kubeConfigPath, kubectx string) (*rest.Config, error) {
+func GetK8sConfig(log logr.Logger, cfg *cfg.KubeConfig) (*rest.Config, error) {
 	// checking in-cluster kubeconfig
 	restConfig, err := rest.InClusterConfig()
 	if err == nil {
+		restConfig.Burst = cfg.Burst
+		restConfig.QPS = cfg.QPS
 		log.Info("Using in-cluster kubeconfig")
 		return restConfig, err
 	}
@@ -37,8 +41,8 @@ func GetK8sConfig(log logr.Logger, kubeConfigPath, kubectx string) (*rest.Config
 	// checking external kubeconfig
 	log.Info("Using external kubeconfig")
 	configAccess := clientcmd.NewDefaultPathOptions()
-	if len(kubeConfigPath) != 0 {
-		configAccess.GlobalFile = kubeConfigPath
+	if len(cfg.KubeConfigPath) != 0 {
+		configAccess.GlobalFile = cfg.KubeConfigPath
 	}
 
 	config, err := configAccess.GetStartingConfig()
@@ -54,14 +58,14 @@ func GetK8sConfig(log logr.Logger, kubeConfigPath, kubectx string) (*rest.Config
 		}
 	}
 
-	if len(kubectx) != 0 {
-		log.Info("Setting kubecontext", "name", kubectx)
+	if len(cfg.KubeCtx) != 0 {
+		log.Info("Setting kubecontext", "name", cfg.KubeCtx)
 
-		err = validateKubeContext(config, kubectx)
+		err = validateKubeContext(config, cfg.KubeCtx)
 		if err != nil {
 			return nil, err
 		}
-		config.CurrentContext = kubectx
+		config.CurrentContext = cfg.KubeCtx
 	}
 
 	return clientcmd.NewDefaultClientConfig(*config, &clientcmd.ConfigOverrides{}).ClientConfig()
diff --git a/pkg/utils/k8s_config_test.go b/pkg/utils/k8s_config_test.go
index 6eb0cff..f800774 100644
--- a/pkg/utils/k8s_config_test.go
+++ b/pkg/utils/k8s_config_test.go
@@ -22,6 +22,8 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"k8s.io/klog/v2/textlogger"
+
+	"github.com/NVIDIA/knavigator/pkg/config"
 )
 
 const testKubeCfg1 = `
@@ -134,7 +136,13 @@ func TestGetK8sConfig(t *testing.T) {
 
 				cfgPath = f.Name()
 			}
-			cfg, err := GetK8sConfig(testLogger, cfgPath, tc.kubeCtx)
+			c := &config.KubeConfig{
+				KubeConfigPath: cfgPath,
+				KubeCtx:        tc.kubeCtx,
+				QPS:            10,
+				Burst:          10,
+			}
+			cfg, err := GetK8sConfig(testLogger, c)
 			if len(tc.expectedErr) != 0 {
 				require.EqualError(t, err, tc.expectedErr)
 			} else {