feat: allow end-users to adjust K8S client QPS and burst (#5910)

aquasecurity · Jan 15, 2024 · 2c9d7c6 · 2c9d7c6
1 parent ffe2ca7
commit 2c9d7c6
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/docs/docs/references/configuration/cli/trivy_kubernetes.md b/docs/docs/references/configuration/cli/trivy_kubernetes.md
@@ -28,6 +28,7 @@ trivy kubernetes [flags] { cluster | all | specific resources like kubectl. eg:
 
 ```
   -A, --all-namespaces                    fetch resources from all cluster namespaces
+      --burst int                         specify the maximum burst for throttle (default 10)
       --cache-backend string              cache backend (e.g. redis://localhost:6379) (default "fs")
       --cache-ttl duration                cache TTL when using redis as cache backend
       --clear-cache                       clear image caches without scanning
@@ -72,6 +73,7 @@ trivy kubernetes [flags] { cluster | all | specific resources like kubectl. eg:
       --password strings                  password. Comma-separated passwords allowed. TRIVY_PASSWORD should be used for security reasons.
       --policy-bundle-repository string   OCI registry URL to retrieve policy bundle from (default "ghcr.io/aquasecurity/trivy-policies:0")
       --policy-namespaces strings         Rego namespaces
+      --qps float                         specify the maximum QPS to the master from this client (default 5)
       --redis-ca string                   redis ca file location, if using redis as cache backend
       --redis-cert string                 redis certificate file location, if using redis as cache backend
       --redis-key string                  redis key file location, if using redis as cache backend

diff --git a/pkg/flag/kubernetes_flags.go b/pkg/flag/kubernetes_flags.go
@@ -88,6 +88,18 @@ var (
 		Default:    "ghcr.io/aquasecurity/node-collector:0.0.9",
 		Usage:      "indicate the image reference for the node-collector scan job",
 	}
+	QPS = Flag{
+		Name:       "qps",
+		ConfigName: "kubernetes.qps",
+		Default:    5.0,
+		Usage:      "specify the maximum QPS to the master from this client",
+	}
+	Burst = Flag{
+		Name:       "burst",
+		ConfigName: "kubernetes.burst",
+		Default:    10,
+		Usage:      "specify the maximum burst for throttle",
+	}
 )
 
 type K8sFlagGroup struct {
@@ -102,6 +114,8 @@ type K8sFlagGroup struct {
 	NodeCollectorNamespace *Flag
 	ExcludeOwned           *Flag
 	ExcludeNodes           *Flag
+	QPS                    *Flag
+	Burst                  *Flag
 }
 
 type K8sOptions struct {
@@ -116,6 +130,8 @@ type K8sOptions struct {
 	NodeCollectorNamespace string
 	ExcludeOwned           bool
 	ExcludeNodes           map[string]string
+	QPS                    float32
+	Burst                  int
 }
 
 func NewK8sFlagGroup() *K8sFlagGroup {
@@ -131,6 +147,8 @@ func NewK8sFlagGroup() *K8sFlagGroup {
 		ExcludeOwned:           &ExcludeOwned,
 		ExcludeNodes:           &ExcludeNodes,
 		NodeCollectorImageRef:  &NodeCollectorImageRef,
+		QPS:                    &QPS,
+		Burst:                  &Burst,
 	}
 }
 
@@ -151,6 +169,8 @@ func (f *K8sFlagGroup) Flags() []*Flag {
 		f.ExcludeOwned,
 		f.ExcludeNodes,
 		f.NodeCollectorImageRef,
+		f.QPS,
+		f.Burst,
 	}
 }
 

diff --git a/pkg/k8s/commands/run.go b/pkg/k8s/commands/run.go
@@ -30,6 +30,8 @@ func Run(ctx context.Context, args []string, opts flag.Options) error {
 	cluster, err := k8s.GetCluster(
 		k8s.WithContext(opts.K8sOptions.ClusterContext),
 		k8s.WithKubeConfig(opts.K8sOptions.KubeConfig),
+		k8s.WithBurst(opts.K8sOptions.Burst),
+		k8s.WithQPS(opts.K8sOptions.QPS),
 	)
 	if err != nil {
 		return xerrors.Errorf("failed getting k8s cluster: %w", err)