Add max-pod-sample-size annotation for the pod collector

pkg/annotations/parser.go

@@ -2,6 +2,7 @@ package annotations
 
 import (
 	"fmt"
+	"strconv"
 	"strings"
 	"time"
 
@@ -13,14 +14,16 @@ const (
 	perReplicaMetricsConfKey = "per-replica"
 	intervalMetricsConfKey   = "interval"
 	minPodReadyAgeConfKey    = "min-pod-ready-age"
+	maxPodSampleSizeConfKey  = "max-pod-sample-size"
 )
 
 type AnnotationConfigs struct {
-	CollectorType  string
-	Configs        map[string]string
-	PerReplica     bool
-	Interval       time.Duration
-	MinPodReadyAge time.Duration
+	CollectorType    string
+	Configs          map[string]string
+	PerReplica       bool
+	Interval         time.Duration
+	MinPodReadyAge   time.Duration
+	MaxPodSampleSize int
 }
 
 type MetricConfigKey struct {
@@ -100,6 +103,15 @@ func (m AnnotationConfigMap) Parse(annotations map[string]string) error {
 			continue
 		}
 
+		if parts[1] == maxPodSampleSizeConfKey {
+			maxPodSampleSize, err := strconv.Atoi(val)
+			if err != nil {
+				return fmt.Errorf("failed to parse max-pod-sample-size value %s for %s: %v", val, key, err)
+			}
+			config.MaxPodSampleSize = maxPodSampleSize
+			continue
+		}
+
 		config.Configs[parts[1]] = val
 	}
 	return nil

pkg/collector/collector.go

@@ -202,6 +202,7 @@ type MetricConfig struct {
 	PerReplica      bool
 	Interval        time.Duration
 	MinPodReadyAge  time.Duration
+	MaxPodSampleSize int
 	MetricSpec      autoscalingv2.MetricSpec
 }
 
@@ -267,6 +268,7 @@ func ParseHPAMetrics(hpa *autoscalingv2.HorizontalPodAutoscaler) ([]*MetricConfi
 			config.Interval = annotationConfigs.Interval
 			config.PerReplica = annotationConfigs.PerReplica
 			config.MinPodReadyAge = annotationConfigs.MinPodReadyAge
+			config.MaxPodSampleSize = annotationConfigs.MaxPodSampleSize
 			// configs specified in annotations takes precedence
 			// over labels
 			for k, v := range annotationConfigs.Configs {

pkg/collector/pod_collector.go

@@ -42,6 +42,7 @@ type PodCollector struct {
 	metric           autoscalingv2.MetricIdentifier
 	metricType       autoscalingv2.MetricSourceType
 	minPodReadyAge   time.Duration
+	maxPodSampleSize int
 	interval         time.Duration
 	logger           *log.Entry
 }
@@ -59,6 +60,7 @@ func NewPodCollector(ctx context.Context, client kubernetes.Interface, argoRollo
 		metric:           config.Metric,
 		metricType:       config.Type,
 		minPodReadyAge:   config.MinPodReadyAge,
+		maxPodSampleSize: config.MaxPodSampleSize,
 		interval:         interval,
 		podLabelSelector: selector,
 		logger:           log.WithFields(log.Fields{"Collector": "Pod"}),
@@ -94,6 +96,7 @@ func (c *PodCollector) GetMetrics(ctx context.Context) ([]CollectedMetric, error
 	ch := make(chan CollectedMetric)
 	errCh := make(chan error)
 	skippedPodsCount := 0
+	sampledPodsCount := 0
 
 	for _, pod := range pods.Items {
 
@@ -108,6 +111,9 @@ func (c *PodCollector) GetMetrics(ctx context.Context) ([]CollectedMetric, error
 				c.logger.Warnf("Skipping metrics collection for pod %s/%s because it's ready age is %s and min-pod-ready-age is set to %s", pod.Namespace, pod.Name, podReadyAge, c.minPodReadyAge)
 			} else {
 				go c.getPodMetric(pod, ch, errCh)
+				if sampledPodsCount++; sampledPodsCount > c.maxPodSampleSize {
+					break
+				}
 			}
 		} else {
 			skippedPodsCount++