feat(collector): add http timing metrics and page_size to run_ai input (#2621)

Author: tothandras

collector/benthos/input/run_ai.go

@@ -3,6 +3,7 @@ package input
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"sync"
 	"time"
 
@@ -24,6 +25,7 @@ const (
 	fieldMetrics              = "metrics"
 	fieldSchedule             = "schedule"
 	fieldMetricsOffset        = "metrics_offset"
+	fieldPageSize             = "page_size"
 	fieldHTTPConfig           = "http"
 	fieldHTTPTimeout          = "timeout"
 	fieldHTTPRetryCount       = "retry_count"
@@ -45,6 +47,9 @@ func runAIInputConfig() *service.ConfigSpec {
 			service.NewStringEnumField(fieldResourceType, resourceTypes...).
 				Default("workload").
 				Description("Run AI resource to collect metrics from."),
+			service.NewIntField(fieldPageSize).
+				Description("Run AI page size.").
+				Default(500),
 			service.NewStringListField(fieldMetrics).
 				Description("Run AI metrics to collect.").
 				Default(lo.Map([]runai.MetricType{
@@ -90,9 +95,10 @@ input:
     url: "${RUNAI_URL:}"
     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}"
+    schedule: "*/30 * * * * *"
+    metrics_scrape_offset: "30s"
+    resource_type: "workload"
+    page_size: 500
     metrics:
       - CPU_LIMIT_CORES
       - CPU_MEMORY_LIMIT_BYTES
@@ -116,7 +122,15 @@ input:
 
 func init() {
 	err := service.RegisterBatchInput("run_ai", runAIInputConfig(), func(conf *service.ParsedConfig, mgr *service.Resources) (service.BatchInput, error) {
-		return newRunAIInput(conf, mgr.Logger())
+		httpMetrics := mgr.Metrics().NewTimer("run_ai_http_request", "url")
+		in, err := newRunAIInput(conf, mgr.Logger(), httpMetrics)
+		if err != nil {
+			return nil, err
+		}
+
+		in.timingMetrics = httpMetrics
+
+		return in, nil
 	})
 	if err != nil {
 		panic(err)
@@ -136,9 +150,11 @@ type runAIInput struct {
 	scheduler     gocron.Scheduler
 	store         map[time.Time][]runai.ResourceWithMetrics
 	mu            sync.Mutex
+
+	timingMetrics *service.MetricTimer
 }
 
-func newRunAIInput(conf *service.ParsedConfig, logger *service.Logger) (*runAIInput, error) {
+func newRunAIInput(conf *service.ParsedConfig, logger *service.Logger, httpMetrics *service.MetricTimer) (*runAIInput, error) {
 	url, err := conf.FieldString(fieldURL)
 	if err != nil {
 		return nil, err
@@ -174,6 +190,15 @@ func newRunAIInput(conf *service.ParsedConfig, logger *service.Logger) (*runAIIn
 		return nil, err
 	}
 
+	pageSize, err := conf.FieldInt(fieldPageSize)
+	if err != nil {
+		return nil, err
+	}
+
+	if pageSize < 100 || pageSize > 500 {
+		return nil, errors.New("page size must be between 100 and 500")
+	}
+
 	var interval time.Duration
 	{
 		// Create a cron scheduler
@@ -219,6 +244,8 @@ func newRunAIInput(conf *service.ParsedConfig, logger *service.Logger) (*runAIIn
 		RetryCount:       retryCount,
 		RetryWaitTime:    retryWaitTime,
 		RetryMaxWaitTime: retryMaxWaitTime,
+		TimingMetrics:    httpMetrics,
+		PageSize:         pageSize,
 	})
 	if err != nil {
 		return nil, err

collector/benthos/input/runai/pods.go

@@ -116,7 +116,7 @@ func (s *Service) ListAllPods(ctx context.Context) ([]Pod, error) {
 
 	for {
 		resp, err := s.ListPods(ctx, ListPodsParams{
-			Limit:  100,
+			Limit:  s.pageSize,
 			Offset: len(pods),
 		})
 		if err != nil {

collector/benthos/input/runai/service.go

@@ -2,6 +2,7 @@ package runai
 
 import (
 	"net/http"
+	"regexp"
 	"time"
 
 	"github.com/go-resty/resty/v2"
@@ -14,20 +15,24 @@ type Service struct {
 	appID     string
 	appSecret string
 	token     string
+	pageSize  int
 }
 
 type HTTPRequestConfig struct {
 	Timeout          time.Duration
 	RetryWaitTime    time.Duration
 	RetryMaxWaitTime time.Duration
 	RetryCount       int
+	PageSize         int
+	TimingMetrics    *service.MetricTimer
 }
 
 func NewService(baseURL, appID, appSecret string, logger *service.Logger, requestConfig HTTPRequestConfig) (*Service, error) {
 	service := &Service{
 		logger:    logger,
 		appID:     appID,
 		appSecret: appSecret,
+		pageSize:  requestConfig.PageSize,
 	}
 
 	client := resty.New().
@@ -59,6 +64,20 @@ func NewService(baseURL, appID, appSecret string, logger *service.Logger, reques
 				service.SetToken("")
 			}
 
+			if requestConfig.TimingMetrics != nil {
+				path := response.Request.RawRequest.URL.Path
+				if matched, err := regexp.MatchString("/api/v1/workloads/[0-9a-f-]+/pods/[0-9a-f-]+/metrics", path); err == nil && matched {
+					path = "/api/v1/workloads/:workloadID/pods/:podID/metrics"
+				} else if matched, err := regexp.MatchString("/api/v1/workloads/[0-9a-f-]+/metrics", path); err == nil && matched {
+					path = "/api/v1/workloads/:workloadID/metrics"
+				} else if matched, err := regexp.MatchString("/api/v1/pods", path); err == nil && matched {
+					path = "/api/v1/pods"
+				} else if matched, err := regexp.MatchString("/api/v1/workloads", path); err == nil && matched {
+					path = "/api/v1/workloads"
+				}
+				requestConfig.TimingMetrics.Timing(response.Time().Nanoseconds(), path)
+			}
+
 			return nil
 		})
 

collector/benthos/input/runai/workloads.go

@@ -136,7 +136,7 @@ func (s *Service) ListAllWorkloads(ctx context.Context) ([]Workload, error) {
 
 	for {
 		resp, err := s.ListWorkloads(ctx, ListWorkloadParams{
-			Limit:  100,
+			Limit:  s.pageSize,
 			Offset: len(workloads),
 		})
 		if err != nil {