feat(collector): extend kubernetes collector (#2387)

tothandras · web-flow · commit d742455e62ef · 2025-03-07T17:00:56.000+01:00
diff --git a/collector/benthos/input/kubernetes.go b/collector/benthos/input/kubernetes.go
@@ -25,10 +25,13 @@ func kubernetesResourcesInputConfig() *service.ConfigSpec {
 	return service.NewConfigSpec().
 		Beta().
 		Categories("Services").
-		Summary("List pods in Kubernetes.").
+		Summary("List resources in Kubernetes.").
 		Fields(
 			service.NewStringListField("namespaces").
-				Description("List of namespaces to list pods from."),
+				Description("List of namespaces to list resources from."),
+			service.NewStringEnumField("resource_type", "pod", "node", "persistentvolume", "persistentvolumeclaim").
+				Description("Type of resource to list.").
+				Default("pod"),
 			service.NewStringField("label_selector").
 				Description("Label selector applied to each list operation.").
 				Optional(),
@@ -46,6 +49,7 @@ func init() {
 
 type kubernetesResourcesInput struct {
 	namespaces    []string
+	resourceType  string
 	labelSelector labels.Selector
 	logger        *service.Logger
 
@@ -63,6 +67,11 @@ func newKubernetesResourcesInput(conf *service.ParsedConfig, logger *service.Log
 		return nil, err
 	}
 
+	resourceType, err := conf.FieldString("resource_type")
+	if err != nil {
+		return nil, err
+	}
+
 	// Normalize the namespaces to lowercase and deduplicate.
 	namespaces = lo.Uniq(lo.Map(namespaces, func(s string, _ int) string { return strings.ToLower(s) }))
 
@@ -116,6 +125,7 @@ func newKubernetesResourcesInput(conf *service.ParsedConfig, logger *service.Log
 	return &kubernetesResourcesInput{
 		namespaces:    namespaces,
 		labelSelector: selector,
+		resourceType:  resourceType,
 		manager:       mgr,
 		client:        client,
 		logger:        logger,
@@ -153,36 +163,84 @@ func (in *kubernetesResourcesInput) ReadBatch(ctx context.Context) (service.Mess
 
 	// Iterate over each namespace and list pods.
 	for _, ns := range in.namespaces {
-		podList := &corev1.PodList{}
 		opts := []client.ListOption{client.InNamespace(ns)}
 		if in.labelSelector != nil {
 			opts = append(opts, client.MatchingLabelsSelector{
 				Selector: in.labelSelector,
 			})
 		}
 
-		if err := in.client.List(ctx, podList, opts...); err != nil {
-			return nil, nil, err
-		}
+		switch in.resourceType {
+		case "pod":
+			podList := &corev1.PodList{}
+			if err := in.client.List(ctx, podList, opts...); err != nil {
+				return nil, nil, err
+			}
+
+			for _, pod := range podList.Items {
+				if !lo.EveryBy(pod.Status.ContainerStatuses, func(cs corev1.ContainerStatus) bool {
+					return cs.State.Running != nil
+				}) {
+					continue
+				}
+
+				encoded, err := json.Marshal(pod)
+				if err != nil {
+					return nil, nil, err
+				}
+
+				in.logger.Debugf("adding pod %s to batch", pod.Name)
+				batch = append(batch, service.NewMessage(encoded))
+			}
+		case "node":
+			nodeList := &corev1.NodeList{}
+			if err := in.client.List(ctx, nodeList, opts...); err != nil {
+				return nil, nil, err
+			}
 
-		for _, pod := range podList.Items {
-			if !lo.EveryBy(pod.Status.ContainerStatuses, func(cs corev1.ContainerStatus) bool {
-				return cs.State.Running != nil
-			}) {
-				continue
+			for _, node := range nodeList.Items {
+				encoded, err := json.Marshal(node)
+				if err != nil {
+					return nil, nil, err
+				}
+
+				in.logger.Debugf("adding node %s to batch", node.Name)
+				batch = append(batch, service.NewMessage(encoded))
+			}
+		case "persistentvolume":
+			persistentVolumeList := &corev1.PersistentVolumeList{}
+			if err := in.client.List(ctx, persistentVolumeList, opts...); err != nil {
+				return nil, nil, err
 			}
 
-			encoded, err := json.Marshal(pod)
-			if err != nil {
+			for _, persistentVolume := range persistentVolumeList.Items {
+				encoded, err := json.Marshal(persistentVolume)
+				if err != nil {
+					return nil, nil, err
+				}
+
+				in.logger.Debugf("adding persistent volume %s to batch", persistentVolume.Name)
+				batch = append(batch, service.NewMessage(encoded))
+			}
+		case "persistentvolumeclaim":
+			persistentVolumeClaimList := &corev1.PersistentVolumeClaimList{}
+			if err := in.client.List(ctx, persistentVolumeClaimList, opts...); err != nil {
 				return nil, nil, err
 			}
 
-			in.logger.Debugf("adding pod %s to batch", pod.Name)
-			batch = append(batch, service.NewMessage(encoded))
+			for _, persistentVolumeClaim := range persistentVolumeClaimList.Items {
+				encoded, err := json.Marshal(persistentVolumeClaim)
+				if err != nil {
+					return nil, nil, err
+				}
+
+				in.logger.Debugf("adding persistent volume claim %s to batch", persistentVolumeClaim.Name)
+				batch = append(batch, service.NewMessage(encoded))
+			}
 		}
 	}
 
-	in.logger.Debugf("batch size: %d", len(batch))
+	in.logger.Debugf("batch size of %s: %d", in.resourceType, len(batch))
 
 	return batch, func(context.Context, error) error {
 		// A nack (when err is non-nil) is handled automatically when we
diff --git a/collector/benthos/input/run_ai.go b/collector/benthos/input/run_ai.go
@@ -23,6 +23,7 @@ const (
 	fieldResourceType         = "resource_type"
 	fieldMetrics              = "metrics"
 	fieldSchedule             = "schedule"
+	fieldMetricsScrapeOffset  = "metrics_scrape_offset"
 	fieldHTTPConfig           = "http"
 	fieldHTTPTimeout          = "timeout"
 	fieldHTTPRetryCount       = "retry_count"
@@ -66,6 +67,9 @@ func runAIInputConfig() *service.ConfigSpec {
 				Description("The cron expression to use for the scrape job.").
 				Examples("*/30 * * * * *", "@every 30s").
 				Default("*/30 * * * * *"),
+			service.NewDurationField(fieldMetricsScrapeOffset).
+				Description("Indicates how far back in time the scraping window should start to account for delays in metric availability.").
+				Default("0s"),
 			service.NewObjectField(fieldHTTPConfig,
 				service.NewDurationField(fieldHTTPTimeout).
 					Description("Request timeout.").
@@ -87,6 +91,7 @@ input:
     app_id: "${RUNAI_APP_ID:}"
     app_secret: "${RUNAI_APP_SECRET:}"
     schedule: "${RUNAI_SCRAPE_SCHEDULE:*/30 * * * * *}"
+    metrics_scrape_offset: "${RUNAI_METRICS_SCRAPE_OFFSET:30s}"
     resource_type: "${RUNAI_RESOURCE_TYPE:workload}"
     metrics:
       - CPU_LIMIT_CORES
@@ -121,15 +126,16 @@ func init() {
 var _ service.BatchInput = (*runAIInput)(nil)
 
 type runAIInput struct {
-	logger       *service.Logger
-	service      *runai.Service
-	resourceType string
-	metrics      []runai.MetricType
-	interval     time.Duration
-	schedule     string
-	scheduler    gocron.Scheduler
-	store        map[time.Time][]runai.ResourceWithMetrics
-	mu           sync.Mutex
+	logger              *service.Logger
+	service             *runai.Service
+	resourceType        string
+	metrics             []runai.MetricType
+	interval            time.Duration
+	schedule            string
+	metricsScrapeOffset time.Duration
+	scheduler           gocron.Scheduler
+	store               map[time.Time][]runai.ResourceWithMetrics
+	mu                  sync.Mutex
 }
 
 func newRunAIInput(conf *service.ParsedConfig, logger *service.Logger) (*runAIInput, error) {
@@ -163,6 +169,11 @@ func newRunAIInput(conf *service.ParsedConfig, logger *service.Logger) (*runAIIn
 		return nil, err
 	}
 
+	metricsScrapeOffset, err := conf.FieldDuration(fieldMetricsScrapeOffset)
+	if err != nil {
+		return nil, err
+	}
+
 	var interval time.Duration
 	{
 		// Create a cron scheduler
@@ -219,12 +230,13 @@ func newRunAIInput(conf *service.ParsedConfig, logger *service.Logger) (*runAIIn
 	}
 
 	return &runAIInput{
-		logger:       logger,
-		service:      service,
-		resourceType: resourceType,
-		interval:     interval,
-		schedule:     schedule,
-		scheduler:    scheduler,
+		logger:              logger,
+		service:             service,
+		resourceType:        resourceType,
+		interval:            interval,
+		schedule:            schedule,
+		metricsScrapeOffset: metricsScrapeOffset,
+		scheduler:           scheduler,
 		metrics: lo.Map(metrics, func(metric string, _ int) runai.MetricType {
 			return runai.MetricType(metric)
 		}),
@@ -241,8 +253,8 @@ func (in *runAIInput) scrape(ctx context.Context, t time.Time) error {
 	case "workload":
 		workloadsWithMetrics, err := in.service.GetAllWorkloadWithMetrics(ctx, runai.MeasurementParams{
 			MetricType: in.metrics,
-			StartTime:  t.Add(-in.interval),
-			EndTime:    t,
+			StartTime:  t.Add(-in.interval).Add(-in.metricsScrapeOffset),
+			EndTime:    t.Add(-in.metricsScrapeOffset),
 		})
 		if err != nil {
 			return err
@@ -256,8 +268,8 @@ func (in *runAIInput) scrape(ctx context.Context, t time.Time) error {
 	case "pod":
 		podsWithMetrics, err := in.service.GetAllPodWithMetrics(ctx, runai.MeasurementParams{
 			MetricType: in.metrics,
-			StartTime:  t.Add(-in.interval),
-			EndTime:    t,
+			StartTime:  t.Add(-in.interval).Add(-in.metricsScrapeOffset),
+			EndTime:    t.Add(-in.metricsScrapeOffset),
 		})
 		if err != nil {
 			return err
