From 1ba6c8c8187363560251876347c693291a405173 Mon Sep 17 00:00:00 2001 From: Petru Anica-Popa Date: Thu, 29 Oct 2026 16:28:32 +0000 Subject: [PATCH] fix(otelmetrics): return complete query results and honor Retry-After --- test/otel/efa/dedup_test.go | 2 +- test/otel/efa/efa_test.go | 2 +- test/otel/efa/node_color_test.go | 2 +- test/otel/gpu/dcgm_test.go | 2 +- test/otel/gpu/dedup_test.go | 2 +- test/otel/gpu/node_color_test.go | 2 +- test/otel/neuron/multi_device_test.go | 4 +- test/otel/neuron/neuron_test.go | 2 +- test/otel/neuron/node_color_test.go | 4 +- test/otel/performance/performance_test.go | 2 +- test/otel/standard/cadvisor_test.go | 4 +- test/otel/standard/dedup_test.go | 10 +- test/otel/standard/ksm_test.go | 12 +- test/otel/standard/kubeletstats_test.go | 2 +- test/otel/standard/labels_common_test.go | 4 +- test/otel/standard/node_exporter_test.go | 2 +- util/otelmetrics/client.go | 94 ++- util/otelmetrics/client_test.go | 58 ++ util/otelmetrics/complete.go | 370 ++++++++++++ util/otelmetrics/complete_test.go | 670 ++++++++++++++++++++++ util/otelmetrics/query_cache.go | 33 +- 21 files changed, 1225 insertions(+), 58 deletions(-) create mode 100644 util/otelmetrics/client_test.go create mode 100644 util/otelmetrics/complete.go create mode 100644 util/otelmetrics/complete_test.go diff --git a/test/otel/efa/dedup_test.go b/test/otel/efa/dedup_test.go index 4d8b7e73c..51b968bd4 100644 --- a/test/otel/efa/dedup_test.go +++ b/test/otel/efa/dedup_test.go @@ -48,7 +48,7 @@ func TestEFANoDuplicateSeries(t *testing.T) { `efa_rx_bytes{"@resource.k8s.cluster.name"="%s","@resource.k8s.node.name"="%s","aws.efa.device"="%s"}`, escaped, pinNode, pinDevice, ) - pinned, err := client.Query(ctx, promql) + pinned, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying pinned efa_rx_bytes") require.NotEmpty(t, pinned, "pinned efa_rx_bytes returned 0 results") require.Equal(t, 1, len(pinned), diff --git a/test/otel/efa/efa_test.go b/test/otel/efa/efa_test.go index 2d80ff1ef..0f2993192 100644 --- a/test/otel/efa/efa_test.go +++ b/test/otel/efa/efa_test.go @@ -406,7 +406,7 @@ func TestEFANodeGroupCoverage(t *testing.T) { promql := fmt.Sprintf( `efa_rx_bytes{"@resource.k8s.cluster.name"="%s","@resource.host.type"="%s"}`, escaped, "c5n.9xlarge") - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying efa_rx_bytes on c5n.9xlarge") require.True(t, len(results) > 0, "EFA metrics missing from c5n.9xlarge nodes — awsefareceiver not collecting?") diff --git a/test/otel/efa/node_color_test.go b/test/otel/efa/node_color_test.go index ab80aca2d..29dd6bf4c 100644 --- a/test/otel/efa/node_color_test.go +++ b/test/otel/efa/node_color_test.go @@ -31,7 +31,7 @@ func TestEFAHasYellowNodeColor(t *testing.T) { promql := fmt.Sprintf( `%s{"@resource.k8s.cluster.name"="%s","@resource.%s"="yellow"}`, metricName, escaped, nodeColorLabel) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s with node-color=yellow filter", metricName) require.True(t, len(results) > 0, "%s: no results from c5n.9xlarge nodes (expected node-color=yellow)", metricName) diff --git a/test/otel/gpu/dcgm_test.go b/test/otel/gpu/dcgm_test.go index 4a096ab5b..6b02f6651 100644 --- a/test/otel/gpu/dcgm_test.go +++ b/test/otel/gpu/dcgm_test.go @@ -422,7 +422,7 @@ func TestDCGMNodeGroupCoverage(t *testing.T) { promql := fmt.Sprintf( `DCGM_FI_DEV_GPU_UTIL{"@resource.k8s.cluster.name"="%s","@resource.host.type"="%s"}`, escaped, ng.InstanceType) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying DCGM_FI_DEV_GPU_UTIL on %s", ng.Description) require.True(t, len(results) > 0, "DCGM missing from %s nodes (%s) — dcgm-exporter not running?", diff --git a/test/otel/gpu/dedup_test.go b/test/otel/gpu/dedup_test.go index abc5d860e..7a796bf73 100644 --- a/test/otel/gpu/dedup_test.go +++ b/test/otel/gpu/dedup_test.go @@ -50,7 +50,7 @@ func TestDCGMNoDuplicateSeries(t *testing.T) { `DCGM_FI_DEV_GPU_UTIL{"@resource.k8s.cluster.name"="%s","@resource.k8s.node.name"="%s",UUID="%s"}`, escaped, pinNode, pinUUID, ) - pinned, err := client.Query(ctx, promql) + pinned, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying pinned DCGM_FI_DEV_GPU_UTIL") require.NotEmpty(t, pinned, "pinned DCGM_FI_DEV_GPU_UTIL returned 0 results") require.Equal(t, 1, len(pinned), diff --git a/test/otel/gpu/node_color_test.go b/test/otel/gpu/node_color_test.go index c9d1be06e..3c5f44057 100644 --- a/test/otel/gpu/node_color_test.go +++ b/test/otel/gpu/node_color_test.go @@ -35,7 +35,7 @@ func TestDCGMHasGreenNodeColor(t *testing.T) { promql := fmt.Sprintf( `%s{"@resource.k8s.cluster.name"="%s","@resource.%s"="green"}`, metricName, escaped, nodeColorLabel) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s with node-color=green filter", metricName) require.True(t, len(results) > 0, "%s: no results from GPU nodes (expected node-color=green)", metricName) diff --git a/test/otel/neuron/multi_device_test.go b/test/otel/neuron/multi_device_test.go index 9b2e96dc2..c82698b32 100644 --- a/test/otel/neuron/multi_device_test.go +++ b/test/otel/neuron/multi_device_test.go @@ -273,7 +273,7 @@ func TestMultiNeuronCorrelatedCoreCount(t *testing.T) { promql := fmt.Sprintf( `neuroncore_utilization_ratio{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"neuron-burn-multi.*"}`, escaped) - burnResults, err := client.Query(ctx, promql) + burnResults, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying neuroncore_utilization_ratio for neuron-burn-multi") require.Equal(t, 2, len(burnResults), fmt.Sprintf("Expected 2 cores correlated to neuron-burn-multi (1 whole device × 2 cores), got %d", len(burnResults))) @@ -289,7 +289,7 @@ func TestMultiNeuronBurnDevicePodLabels(t *testing.T) { promql := fmt.Sprintf( `neuroncore_utilization_ratio{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"neuron-burn-multi.*"}`, escaped) - burnResults, err := client.Query(ctx, promql) + burnResults, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying neuroncore_utilization_ratio for neuron-burn-multi") require.Equal(t, 2, len(burnResults), fmt.Sprintf("Expected 2 cores correlated to neuron-burn-multi, got %d", len(burnResults))) diff --git a/test/otel/neuron/neuron_test.go b/test/otel/neuron/neuron_test.go index abe2e578e..9882d12b8 100644 --- a/test/otel/neuron/neuron_test.go +++ b/test/otel/neuron/neuron_test.go @@ -528,7 +528,7 @@ func TestNeuronNodeGroupCoverage(t *testing.T) { promql := fmt.Sprintf( `neuroncore_utilization_ratio{"@resource.k8s.cluster.name"="%s","@resource.host.type"="%s"}`, escaped, ng.InstanceType) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying neuroncore_utilization_ratio on %s", ng.Description) require.True(t, len(results) > 0, "Neuron metrics missing from %s nodes (%s) — neuron-monitor not running?", diff --git a/test/otel/neuron/node_color_test.go b/test/otel/neuron/node_color_test.go index e033946e4..7d24db650 100644 --- a/test/otel/neuron/node_color_test.go +++ b/test/otel/neuron/node_color_test.go @@ -33,7 +33,7 @@ func TestNeuronHasRedNodeColor(t *testing.T) { promql := fmt.Sprintf( `%s{"@resource.k8s.cluster.name"="%s","@resource.%s"="red"}`, metricName, escaped, nodeColorLabel) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s with node-color=red filter", metricName) require.True(t, len(results) > 0, "%s: no results from inf2.xlarge nodes (expected node-color=red)", metricName) @@ -55,7 +55,7 @@ func TestNeuronHasPurpleNodeColor(t *testing.T) { promql := fmt.Sprintf( `%s{"@resource.k8s.cluster.name"="%s","@resource.%s"="purple"}`, md.Name, escaped, nodeColorLabel) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s with node-color=purple filter", md.Name) require.True(t, len(results) > 0, "%s: no results from inf2.24xlarge node (expected node-color=purple)", md.Name) diff --git a/test/otel/performance/performance_test.go b/test/otel/performance/performance_test.go index 3c1fa7be8..bd8798a3c 100644 --- a/test/otel/performance/performance_test.go +++ b/test/otel/performance/performance_test.go @@ -62,7 +62,7 @@ func getNodeAllocatable(t *testing.T, metricName string, resource string) (float // and skew the percent-of-node denominator the thresholds are calibrated to. query := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s", resource="%s"}`, metricName, otelmetrics.EscapePromQLValue(cfg.ClusterName), resource) - results, err := client.Query(context.Background(), query) + results, err := client.QueryComplete(context.Background(), query) require.NoError(t, err, "failed to query node allocatable for resource=%s", resource) require.NotEmpty(t, results, "no kube_node_status_allocatable data for resource=%s — are KSM metrics being collected?", resource) var sum float64 diff --git a/test/otel/standard/cadvisor_test.go b/test/otel/standard/cadvisor_test.go index 8e1a8ccdb..7d1d67ea1 100644 --- a/test/otel/standard/cadvisor_test.go +++ b/test/otel/standard/cadvisor_test.go @@ -138,7 +138,7 @@ func TestCadvisorNginxWorkloadLabels(t *testing.T) { t.Run(name, func(t *testing.T) { promql := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"nginx-test.*"}`, name, cfg.ClusterName) - nginx, err := client.Query(context.Background(), promql) + nginx, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s for nginx-test", name) require.True(t, len(nginx) > 0, "No %s from nginx-test pods", name) for _, r := range nginx { @@ -237,7 +237,7 @@ func TestCadvisorNodeGroupCoverage(t *testing.T) { promql := fmt.Sprintf( `container_memory_working_set_bytes{"@resource.k8s.cluster.name"="%s","@resource.host.type"="%s"}`, otelmetrics.EscapePromQLValue(cfg.ClusterName), ng.InstanceType) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying cadvisor on %s", ng.Description) require.True(t, len(results) > 0, "cadvisor missing from %s (%s)", ng.Description, ng.InstanceType) }) diff --git a/test/otel/standard/dedup_test.go b/test/otel/standard/dedup_test.go index e5c95006d..62a7b69a5 100644 --- a/test/otel/standard/dedup_test.go +++ b/test/otel/standard/dedup_test.go @@ -29,7 +29,7 @@ func TestNodeExporterNoDuplicateSeries(t *testing.T) { `node_cpu_seconds_total{"@resource.k8s.cluster.name"="%s","@resource.host.type"="%s"}`, escaped, clusterHostTypes[0], ) - results, err := client.Query(ctx, promql) + results, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying node_cpu_seconds_total") require.NotEmpty(t, results, "node_cpu_seconds_total not available") @@ -44,7 +44,7 @@ func TestNodeExporterNoDuplicateSeries(t *testing.T) { `node_cpu_seconds_total{"@resource.k8s.cluster.name"="%s","@resource.k8s.node.name"="%s",cpu="%s",mode="%s"}`, escaped, pinNode, pinCPU, pinMode, ) - pinned, err := client.Query(ctx, promql) + pinned, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying pinned node_cpu_seconds_total") require.NotEmpty(t, pinned, "pinned node_cpu_seconds_total returned 0 results") require.Equal(t, 1, len(pinned), @@ -62,7 +62,7 @@ func TestCadvisorNoDuplicateSeries(t *testing.T) { `container_memory_working_set_bytes{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"nginx-test.*"}`, escaped, ) - results, err := client.Query(ctx, promql) + results, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying container_memory_working_set_bytes for nginx-test") require.NotEmpty(t, results, "container_memory_working_set_bytes not available for nginx-test pods") @@ -84,7 +84,7 @@ func TestCadvisorNoDuplicateSeries(t *testing.T) { `container_memory_working_set_bytes{"@resource.k8s.cluster.name"="%s","@resource.k8s.node.name"="%s","@resource.k8s.pod.name"="%s","@resource.k8s.container.name"="%s"}`, escaped, pinNode, pinPod, pinContainer, ) - pinned, err := client.Query(ctx, promql) + pinned, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying pinned container_memory_working_set_bytes") require.NotEmpty(t, pinned, "pinned container_memory_working_set_bytes returned 0 results") require.Equal(t, 1, len(pinned), @@ -128,7 +128,7 @@ func TestKubeletstatsNoDuplicateSeries(t *testing.T) { `{"__name__"="k8s.node.cpu.usage","@resource.k8s.cluster.name"="%s","@resource.k8s.node.name"="%s"}`, escaped, pinNode, ) - pinned, err := client.Query(ctx, promql) + pinned, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying pinned k8s.node.cpu.usage") require.NotEmpty(t, pinned, "pinned k8s.node.cpu.usage returned 0 results") require.Equal(t, 1, len(pinned), diff --git a/test/otel/standard/ksm_test.go b/test/otel/standard/ksm_test.go index d2ad15d88..45270585a 100644 --- a/test/otel/standard/ksm_test.go +++ b/test/otel/standard/ksm_test.go @@ -544,7 +544,7 @@ func TestKSM_PodBucket_NginxDeploymentName(t *testing.T) { ctx := context.Background() promql := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"nginx-test.*"}`, metricName, cfg.ClusterName) - results, err := client.Query(ctx, promql) + results, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying %s for nginx-test", metricName) require.True(t, len(results) > 0, "No %s results from nginx-test pods", metricName) for _, r := range results { @@ -643,7 +643,7 @@ func TestKSM_PodBucket_StatefulSetPodName(t *testing.T) { t.Parallel() promql := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"ksm-test-statefulset.*"}`, metricName, cfg.ClusterName) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s for ksm-test-statefulset", metricName) require.True(t, len(results) > 0, "No %s results from ksm-test-statefulset pods", metricName) for _, r := range results { @@ -665,7 +665,7 @@ func TestKSM_PodBucket_JobPodName(t *testing.T) { t.Parallel() promql := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"ksm-test-job.*"}`, metricName, cfg.ClusterName) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s for ksm-test-job", metricName) require.True(t, len(results) > 0, "No %s results from ksm-test-job pods", metricName) for _, r := range results { @@ -687,7 +687,7 @@ func TestKSM_PodBucket_CronJobPodName(t *testing.T) { t.Parallel() promql := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"ksm-test-cronjob.*"}`, metricName, cfg.ClusterName) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s for ksm-test-cronjob", metricName) if len(results) == 0 { t.Skipf("No %s results from ksm-test-cronjob pods (may have been cleaned up)", metricName) @@ -714,7 +714,7 @@ func TestKSM_PodBucket_ReplicaSetPodName(t *testing.T) { t.Parallel() promql := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"ksm-test-replicaset.*"}`, metricName, cfg.ClusterName) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s for ksm-test-replicaset", metricName) require.True(t, len(results) > 0, "No %s results from ksm-test-replicaset pods", metricName) for _, r := range results { @@ -1008,7 +1008,7 @@ func TestKSM_ContainerBucket_NginxDeploymentName(t *testing.T) { ctx := context.Background() promql := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"nginx-test.*"}`, metricName, cfg.ClusterName) - results, err := client.Query(ctx, promql) + results, err := client.QueryComplete(ctx, promql) require.NoError(t, err, "querying %s for nginx-test", metricName) require.True(t, len(results) > 0, "No %s results from nginx-test pods", metricName) for _, r := range results { diff --git a/test/otel/standard/kubeletstats_test.go b/test/otel/standard/kubeletstats_test.go index 27a40860b..993a13543 100644 --- a/test/otel/standard/kubeletstats_test.go +++ b/test/otel/standard/kubeletstats_test.go @@ -247,7 +247,7 @@ func TestKubeletstatsNodeGroupCoverage(t *testing.T) { promql := fmt.Sprintf( `{"__name__"="k8s.node.cpu.usage","@resource.k8s.cluster.name"="%s","@resource.host.type"="%s"}`, otelmetrics.EscapePromQLValue(cfg.ClusterName), ng.InstanceType) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying k8s.node.cpu.usage on %s", ng.Description) require.True(t, len(results) > 0, "kubeletstats missing from %s nodes (%s) — kubeletstatsreceiver not scraping?", diff --git a/test/otel/standard/labels_common_test.go b/test/otel/standard/labels_common_test.go index 8834acf3c..79ced4255 100644 --- a/test/otel/standard/labels_common_test.go +++ b/test/otel/standard/labels_common_test.go @@ -503,7 +503,7 @@ func TestWorkloadLabels(t *testing.T) { t.Run(metricName+"/cadvisor_nginx_workload", func(t *testing.T) { promql := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s","@resource.k8s.pod.name"=~"nginx-test.*"}`, metricName, cfg.ClusterName) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s for nginx-test", metricName) require.True(t, len(results) > 0, "No %s results from nginx-test pods", metricName) for _, r := range results { @@ -607,7 +607,7 @@ func TestCustomNodeLabels(t *testing.T) { t.Run(metricName+"/has_blue", func(t *testing.T) { promql := fmt.Sprintf(`%s{"@resource.k8s.cluster.name"="%s","@resource.%s"="blue"}`, metricName, cfg.ClusterName, nodeColorLabel) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying %s with blue filter", metricName) require.True(t, len(results) > 0, "%s no results from standard nodes (color=blue)", metricName) diff --git a/test/otel/standard/node_exporter_test.go b/test/otel/standard/node_exporter_test.go index 69db04781..105e519fc 100644 --- a/test/otel/standard/node_exporter_test.go +++ b/test/otel/standard/node_exporter_test.go @@ -137,7 +137,7 @@ func TestNodeExporterNodeGroupCoverage(t *testing.T) { promql := fmt.Sprintf( `node_load1{"@resource.k8s.cluster.name"="%s","@resource.host.type"="%s"}`, otelmetrics.EscapePromQLValue(cfg.ClusterName), ng.InstanceType) - results, err := client.Query(context.Background(), promql) + results, err := client.QueryComplete(context.Background(), promql) require.NoError(t, err, "querying node_load1 on %s", ng.Description) require.True(t, len(results) > 0, "node_exporter missing from %s nodes (%s) — DaemonSet not scheduling?", diff --git a/util/otelmetrics/client.go b/util/otelmetrics/client.go index c9315830b..6d308bdbe 100644 --- a/util/otelmetrics/client.go +++ b/util/otelmetrics/client.go @@ -12,6 +12,7 @@ import ( "net/http" "net/url" "strconv" + "strings" "time" "github.com/aws/aws-sdk-go-v2/aws" @@ -39,11 +40,13 @@ type OtelMetricsClient struct { region string signingService string maxRetries int + seriesLimit int // endpoint's per-query series cap; 0 means defaultSeriesLimit } type promqlResponse struct { - Status string `json:"status"` - Data promqlData `json:"data"` + Status string `json:"status"` + Data promqlData `json:"data"` + Warnings []string `json:"warnings"` } type promqlData struct { @@ -74,18 +77,43 @@ func NewClient(ctx context.Context, config TestConfig) (*OtelMetricsClient, erro }, nil } -// Query executes a PromQL instant query and returns parsed results. +// Query executes a single PromQL instant query and returns parsed results. The endpoint caps +// results at 500 series, often without saying so; use QueryComplete when every series matters. func (c *OtelMetricsClient) Query(ctx context.Context, promql string) ([]MetricResult, error) { + results, truncated, err := c.queryAt(ctx, promql, "") + if truncated { + log.Printf("WARNING: result truncated by the endpoint series limit query=%s", promql) + } + return results, err +} + +// queryAt runs one instant query, evaluated at evalTime when non-empty. It also returns +// whether the endpoint warned of truncation. +func (c *OtelMetricsClient) queryAt(ctx context.Context, promql, evalTime string) ([]MetricResult, bool, error) { params := url.Values{"query": {promql}} + if evalTime != "" { + params.Set("time", evalTime) + } log.Printf("querying promql=%s", promql) raw, err := c.requestWithRetry(ctx, c.queryURL, params) if err != nil { - return nil, err + return nil, false, err } results := c.parseResponse(raw) log.Printf("query returned count=%d", len(results)) - return results, nil + return results, hasTruncationWarning(raw.Warnings), nil +} + +// hasTruncationWarning reports the endpoint's "results truncated due to series limit" warning. +// Aggregations carry it reliably; plain selectors only sometimes, so it is never proof of completeness. +func hasTruncationWarning(warnings []string) bool { + for _, w := range warnings { + if strings.Contains(strings.ToLower(w), "truncated") { + return true + } + } + return false } // MetricNames returns the distinct metric names present for the given PromQL matcher over [start, end]. @@ -143,8 +171,13 @@ func (c *OtelMetricsClient) requestRawWithRetry(ctx context.Context, baseURL str return nil, fmt.Errorf("signing request: %w", err) } + release, err := acquireRequestSlot(ctx) + if err != nil { + return nil, err + } resp, err := c.httpClient.Do(req) if err != nil { + release() lastErr = err log.Printf("request failed attempt=%d error=%v", attempt+1, err) if attempt < c.maxRetries-1 { @@ -155,6 +188,7 @@ func (c *OtelMetricsClient) requestRawWithRetry(ctx context.Context, baseURL str body, readErr := io.ReadAll(resp.Body) resp.Body.Close() + release() if readErr != nil { return nil, fmt.Errorf("reading response: %w", readErr) } @@ -164,7 +198,7 @@ func (c *OtelMetricsClient) requestRawWithRetry(ctx context.Context, baseURL str lastErr = fmt.Errorf("server error %d: %s", resp.StatusCode, truncate(string(body), 200)) log.Printf("server error, retrying status=%d attempt=%d", resp.StatusCode, attempt+1) if attempt < c.maxRetries-1 { - time.Sleep(time.Duration(1<= 0 { + d = time.Duration(secs) * time.Second + } else if at, err := http.ParseTime(retryAfter); err == nil { + d = at.Sub(now) + } + if d < 0 { + d = 0 + } + if d > maxRetryAfter { + d = maxRetryAfter + } + return d +} + +type requestSlotsKey struct{} + +// withRequestSlots makes every request issued under ctx hold one slot of sem while it is in flight. +func withRequestSlots(ctx context.Context, sem chan struct{}) context.Context { + return context.WithValue(ctx, requestSlotsKey{}, sem) +} + +func acquireRequestSlot(ctx context.Context) (release func(), err error) { + sem, _ := ctx.Value(requestSlotsKey{}).(chan struct{}) + if sem == nil { + return func() {}, nil + } + select { + case sem <- struct{}{}: + return func() { <-sem }, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + func (c *OtelMetricsClient) parseResponse(response *promqlResponse) []MetricResult { var results []MetricResult for _, series := range response.Data.Result { @@ -242,8 +316,9 @@ type RangeResult struct { } type promqlRangeResponse struct { - Status string `json:"status"` - Data promqlRangeData `json:"data"` + Status string `json:"status"` + Data promqlRangeData `json:"data"` + Warnings []string `json:"warnings"` } type promqlRangeData struct { @@ -275,6 +350,9 @@ func (c *OtelMetricsClient) QueryRange(ctx context.Context, promql string, start if err := json.Unmarshal(rawBytes, &rangeResp); err != nil { return nil, fmt.Errorf("parsing range response: %w", err) } + if hasTruncationWarning(rangeResp.Warnings) { + log.Printf("WARNING: range result truncated by the endpoint series limit query=%s", promql) + } var results []RangeResult for _, series := range rangeResp.Data.Result { diff --git a/util/otelmetrics/client_test.go b/util/otelmetrics/client_test.go new file mode 100644 index 000000000..811141381 --- /dev/null +++ b/util/otelmetrics/client_test.go @@ -0,0 +1,58 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: MIT + +package otelmetrics + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRetryDelay(t *testing.T) { + now := time.Date(2026, 10, 15, 12, 0, 0, 0, time.UTC) + for _, tc := range []struct { + header string + attempt int + want time.Duration + }{ + {"", 0, time.Second}, + {"", 2, 4 * time.Second}, + {"not-a-delay", 1, 2 * time.Second}, + {"0", 2, 0}, + {"5", 0, 5 * time.Second}, + {"120", 0, maxRetryAfter}, + {now.Add(3 * time.Second).Format(http.TimeFormat), 0, 3 * time.Second}, + {now.Add(-time.Minute).Format(http.TimeFormat), 0, 0}, + } { + assert.Equal(t, tc.want, retryDelay(tc.header, tc.attempt, now), "Retry-After %q attempt %d", tc.header, tc.attempt) + } +} + +func TestQuery_ThrottledRequestWaitsForRetryAfter(t *testing.T) { + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if atomic.AddInt32(&calls, 1) == 1 { + w.Header().Set("Retry-After", "0") + http.Error(w, "throttled", http.StatusTooManyRequests) + return + } + fmt.Fprint(w, `{"status":"success","data":{"resultType":"vector","result":[]}}`) + })) + defer srv.Close() + c := newTestClient(srv.URL) + c.maxRetries = 2 + + start := time.Now() + _, err := c.Query(context.Background(), "up") + require.NoError(t, err) + assert.EqualValues(t, 2, atomic.LoadInt32(&calls)) + assert.Less(t, time.Since(start), 500*time.Millisecond, "Retry-After: 0 replaces the 1 s backoff") +} diff --git a/util/otelmetrics/complete.go b/util/otelmetrics/complete.go new file mode 100644 index 000000000..1fae4ccec --- /dev/null +++ b/util/otelmetrics/complete.go @@ -0,0 +1,370 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: MIT + +package otelmetrics + +import ( + "context" + "errors" + "fmt" + "log" + "regexp" + "sort" + "strconv" + "strings" + "sync" + "time" +) + +// ErrTruncated reports that the endpoint dropped series to stay under its per-query +// series limit and the complete result could not be recovered. +var ErrTruncated = errors.New("query result truncated by the endpoint series limit") + +const ( + // defaultSeriesLimit is the endpoint's per-query cap ("series limit (500)" in its warnings). + defaultSeriesLimit = 500 + // maxMatcherValueLen is the endpoint's limit on one matcher value, which bounds a shard's regex. + maxMatcherValueLen = 1000 + sampleSize = 50 + maxShardDepth = 4 + maxSplitCandidates = 3 + shardConcurrency = 4 + nodeNameLabel = "@resource.k8s.node.name" + // evalTimeLag keeps the evaluation time that the client picks behind the endpoint's clock + // and its newest samples. + evalTimeLag = 5 * time.Second +) + +var metricNameRE = regexp.MustCompile(`^[a-zA-Z_:][a-zA-Z0-9_:]*$`) + +// QueryComplete returns every series matched by an instant vector selector. The endpoint caps +// selector results and does not reliably say so, so the result is compared with count() at the +// same evaluation time and, only when short, re-fetched as disjoint shards (by node first). +// The client picks that time, so the selector and its count run at once. +// Anything that is not a plain vector selector is passed to Query unchanged. +func (c *OtelMetricsClient) QueryComplete(ctx context.Context, selector string) ([]MetricResult, error) { + if !isVectorSelector(selector) { + return c.Query(ctx, selector) + } + evalTime := evalTimeAt(time.Now()) + var before int + var countErr error + counted := make(chan struct{}) + go func() { + defer close(counted) + before, countErr = c.countAt(ctx, selector, evalTime) + }() + results, _, err := c.queryAt(ctx, selector, evalTime) + <-counted + if err != nil { + return nil, err + } + if countErr != nil { + return nil, countErr + } + if len(results) >= before { + return results, nil + } + log.Printf("result truncated: got %d of %d series, re-fetching in shards query=%s", len(results), before, selector) + + all, err := c.split(ctx, selector, evalTime, before, results, nil, 1) + if err != nil { + return nil, err + } + after, err := c.countAt(ctx, selector, evalTime) + if err != nil { + return nil, err + } + lo, hi := before, after + if lo > hi { + lo, hi = hi, lo + } + if len(all) < lo || len(all) > hi { + return nil, fmt.Errorf("%w: shards returned %d series but count() was %d before and %d after (query %s)", + ErrTruncated, len(all), before, after, selector) + } + log.Printf("sharded fetch complete: %d series query=%s", len(all), selector) + return all, nil +} + +type shardGroup struct { + value string // "" when the label is absent + count int +} + +type shard struct { + matcher string + want int + exact bool // the shard pins the label to one value (or to absent) +} + +// evalTimeAt is the evaluation time, in Unix seconds, that QueryComplete pins its queries to. +func evalTimeAt(now time.Time) string { + return strconv.FormatInt(now.Add(-evalTimeLag).Unix(), 10) +} + +// countAt returns count(selector) at evalTime. +func (c *OtelMetricsClient) countAt(ctx context.Context, selector, evalTime string) (int, error) { + groups, err := c.countBy(ctx, selector, "", evalTime) + if err != nil { + return 0, err + } + return groupTotal(groups), nil +} + +func groupTotal(groups []shardGroup) int { + total := 0 + for _, g := range groups { + total += g.count + } + return total +} + +// countBy returns the series count per value of label (or the total when label is ""). +// It fails with ErrTruncated when the label has too many values to list in one response. +func (c *OtelMetricsClient) countBy(ctx context.Context, selector, label, evalTime string) ([]shardGroup, error) { + promql := "count(" + selector + ")" + if label != "" { + promql = fmt.Sprintf("count by (%s)(%s)", quoteLabel(label), selector) + } + results, truncated, err := c.queryAt(ctx, promql, evalTime) + if err != nil { + return nil, err + } + if truncated { + return nil, fmt.Errorf("%w: too many groups (query %s)", ErrTruncated, promql) + } + groups := make([]shardGroup, 0, len(results)) + for _, r := range results { + g := shardGroup{count: int(r.Value)} + if label != "" { + g.value = r.Labels.AllLabels()[label] + } + groups = append(groups, g) + } + return groups, nil +} + +// split re-fetches a truncated selector of want series as disjoint shards over the first label that divides it. +func (c *OtelMetricsClient) split(ctx context.Context, selector, evalTime string, want int, sample []MetricResult, pinned map[string]bool, depth int) ([]MetricResult, error) { + if depth > maxShardDepth { + return nil, fmt.Errorf("%w: still truncated after %d shard levels (query %s)", ErrTruncated, maxShardDepth, selector) + } + if len(sample) == 0 { + // Stale series filled the cap; topk evaluates fully, so it yields live series to inspect. + sample, _, _ = c.queryAt(ctx, fmt.Sprintf("topk(%d,%s)", sampleSize, selector), evalTime) + } + for _, label := range splitCandidates(sample, pinned) { + groups, err := c.countBy(ctx, selector, label, evalTime) + if errors.Is(err, ErrTruncated) { + continue + } + if err != nil { + return nil, err + } + if len(groups) < 2 || groupTotal(groups) < want { + continue // label does not divide the set, or its groups were cut short + } + shards := packShards(label, groups, c.shardTarget(depth)) + log.Printf("fetching %d shards by %s depth=%d query=%s", len(shards), label, depth, selector) + return c.fetchShards(ctx, selector, evalTime, label, shards, pinned, depth) + } + return nil, fmt.Errorf("%w: no label divides the result further (query %s)", ErrTruncated, selector) +} + +func (c *OtelMetricsClient) fetchShards(ctx context.Context, selector, evalTime, label string, shards []shard, pinned map[string]bool, depth int) ([]MetricResult, error) { + type outcome struct { + results []MetricResult + err error + } + outcomes := make([]outcome, len(shards)) + sem := make(chan struct{}, shardConcurrency) + var wg sync.WaitGroup + for i := range shards { + wg.Add(1) + go func(i int) { + defer wg.Done() + sem <- struct{}{} + defer func() { <-sem }() + sh := shards[i] + sub := withMatcher(selector, sh.matcher) + results, _, err := c.queryAt(ctx, sub, evalTime) + if err == nil && len(results) < sh.want { + subPinned := pinned + if sh.exact { + subPinned = withPinned(pinned, label) + } + results, err = c.split(ctx, sub, evalTime, sh.want, results, subPinned, depth+1) + } + outcomes[i] = outcome{results: results, err: err} + }(i) + } + wg.Wait() + + var all []MetricResult + for _, o := range outcomes { + if o.err != nil { + return nil, o.err + } + all = append(all, o.results...) + } + return all, nil +} + +// shardTarget is the series budget per shard: 80% of the cap, since stale series also count +// against it, halved at each level so a re-split always yields smaller shards. +func (c *OtelMetricsClient) shardTarget(depth int) int { + limit := c.seriesLimit + if limit <= 0 { + limit = defaultSeriesLimit + } + return (limit * 4 / 5) >> (depth - 1) +} + +// splitCandidates orders the labels worth splitting on: node name first, then the sample's +// labels by number of distinct values. Labels already pinned to one value are skipped. +func splitCandidates(sample []MetricResult, pinned map[string]bool) []string { + distinct := make(map[string]map[string]struct{}) + for _, r := range sample { + for k, v := range r.Labels.AllLabels() { + if distinct[k] == nil { + distinct[k] = make(map[string]struct{}) + } + distinct[k][v] = struct{}{} + } + } + var out []string + if !pinned[nodeNameLabel] && (len(sample) == 0 || distinct[nodeNameLabel] != nil) { + out = append(out, nodeNameLabel) + } + var keys []string + for k, values := range distinct { + if k != nodeNameLabel && !pinned[k] && len(values) > 1 { + keys = append(keys, k) + } + } + sort.Slice(keys, func(i, j int) bool { + if ni, nj := len(distinct[keys[i]]), len(distinct[keys[j]]); ni != nj { + return ni > nj + } + return keys[i] < keys[j] + }) + out = append(out, keys...) + if len(out) > maxSplitCandidates { + out = out[:maxSplitCandidates] + } + return out +} + +// packShards bins label values into shards of at most target series. The absent value gets +// its own !~".+" shard because ="" does not match a missing label on this endpoint. +func packShards(label string, groups []shardGroup, target int) []shard { + if target < 1 { + target = 1 + } + sorted := append([]shardGroup(nil), groups...) + sort.Slice(sorted, func(i, j int) bool { + if sorted[i].count != sorted[j].count { + return sorted[i].count > sorted[j].count + } + return sorted[i].value < sorted[j].value + }) + + var out []shard + type bin struct { + values []string + want, chars int + } + var bins []bin + for _, g := range sorted { + if g.value == "" { + out = append(out, shard{matcher: quoteLabel(label) + `!~".+"`, want: g.count, exact: true}) + continue + } + chars := len(regexp.QuoteMeta(g.value)) + placed := false + for i := range bins { + if bins[i].want+g.count <= target && bins[i].chars+1+chars <= maxMatcherValueLen { + bins[i].values = append(bins[i].values, g.value) + bins[i].want += g.count + bins[i].chars += 1 + chars + placed = true + break + } + } + if !placed { + bins = append(bins, bin{values: []string{g.value}, want: g.count, chars: chars}) + } + } + for _, b := range bins { + out = append(out, shard{matcher: valueMatcher(label, b.values), want: b.want, exact: len(b.values) == 1}) + } + return out +} + +func valueMatcher(label string, values []string) string { + if len(values) == 1 { + return fmt.Sprintf(`%s="%s"`, quoteLabel(label), promqlEscaper.Replace(values[0])) + } + alts := make([]string, len(values)) + for i, v := range values { + alts[i] = regexp.QuoteMeta(v) + } + sort.Strings(alts) + return fmt.Sprintf(`%s=~"%s"`, quoteLabel(label), promqlEscaper.Replace(strings.Join(alts, "|"))) +} + +func quoteLabel(label string) string { + return `"` + promqlEscaper.Replace(label) + `"` +} + +func withPinned(pinned map[string]bool, label string) map[string]bool { + out := make(map[string]bool, len(pinned)+1) + for k := range pinned { + out[k] = true + } + out[label] = true + return out +} + +// withMatcher adds one label matcher to a vector selector. +func withMatcher(selector, matcher string) string { + s := strings.TrimSpace(selector) + if !strings.HasSuffix(s, "}") { + return s + "{" + matcher + "}" + } + body := strings.TrimSpace(s[:len(s)-1]) + if strings.HasSuffix(body, "{") || strings.HasSuffix(body, ",") { + return body + matcher + "}" + } + return body + "," + matcher + "}" +} + +// isVectorSelector reports whether s is a single instant vector selector (name, name{...} +// or {...}) with no range, offset, function or operator around it. +func isVectorSelector(s string) bool { + s = strings.TrimSpace(s) + open := strings.IndexByte(s, '{') + if open < 0 { + return metricNameRE.MatchString(s) + } + if name := strings.TrimSpace(s[:open]); name != "" && !metricNameRE.MatchString(name) { + return false + } + var quote byte + for i := open + 1; i < len(s); i++ { + ch := s[i] + switch { + case quote != 0 && ch == '\\' && quote != '`': + i++ + case quote != 0 && ch == quote: + quote = 0 + case quote != 0: + case ch == '"' || ch == '\'' || ch == '`': + quote = ch + case ch == '}': + return i == len(s)-1 + } + } + return false +} diff --git a/util/otelmetrics/complete_test.go b/util/otelmetrics/complete_test.go new file mode 100644 index 000000000..816710928 --- /dev/null +++ b/util/otelmetrics/complete_test.go @@ -0,0 +1,670 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: MIT + +package otelmetrics + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "regexp" + "sort" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const fakeEvalTime = "1700000000.123" + +type fakeSeries struct { + labels map[string]string + stale bool // occupies the cap like a recently ended series, but has no sample +} + +type fakeRequest struct { + query, time string +} + +// fakePromQL mimics the endpoint: plain selectors are capped at limit series (stale ones +// included), aggregations warn when capped, and `=""` never matches an absent label. +type fakePromQL struct { + t *testing.T + series []fakeSeries + limit int + countOffset int // added to count() without by, to simulate an inconsistent endpoint + brokenNotRegex bool // make !~ match everything, so shards overlap + mu sync.Mutex + requests []fakeRequest +} + +func (f *fakePromQL) handler(w http.ResponseWriter, r *http.Request) { + q := r.URL.Query() + query, ts := q.Get("query"), q.Get("time") + f.mu.Lock() + f.requests = append(f.requests, fakeRequest{query: query, time: ts}) + f.mu.Unlock() + if ts == "" { + ts = fakeEvalTime + } + + type result struct { + Metric map[string]string `json:"metric"` + Value []json.RawMessage `json:"value"` + } + resp := struct { + Status string `json:"status"` + Warnings []string `json:"warnings,omitempty"` + Data struct { + ResultType string `json:"resultType"` + Result []result `json:"result"` + } `json:"data"` + }{Status: "success"} + resp.Data.ResultType = "vector" + sample := func(metric map[string]string, v int) result { + return result{Metric: metric, Value: []json.RawMessage{json.RawMessage(ts), json.RawMessage(strconv.Quote(strconv.Itoa(v)))}} + } + + switch { + case strings.HasPrefix(query, "topk("): + m := regexp.MustCompile(`^topk\((\d+),(.*)\)$`).FindStringSubmatch(query) + if m == nil { + f.badRequest(w, "unexpected topk query %q", query) + return + } + k, _ := strconv.Atoi(m[1]) + for _, s := range f.match(m[2]) { + if !s.stale && len(resp.Data.Result) < k { + resp.Data.Result = append(resp.Data.Result, sample(s.labels, 1)) + } + } + case strings.HasPrefix(query, "count by ("): + label, inner, ok := parseCountBy(query) + if !ok { + f.badRequest(w, "unexpected count by query %q", query) + return + } + groups := map[string]int{} + for _, s := range f.match(inner) { + if !s.stale { + groups[s.labels[label]]++ + } + } + keys := make([]string, 0, len(groups)) + for k := range groups { + keys = append(keys, k) + } + sort.Strings(keys) + if len(keys) > f.limit { + keys = keys[:f.limit] + resp.Warnings = []string{fmt.Sprintf("results truncated due to series limit (%d).", f.limit)} + } + for _, k := range keys { + m := map[string]string{} + if k != "" { + m[label] = k + } + resp.Data.Result = append(resp.Data.Result, sample(m, groups[k])) + } + case strings.HasPrefix(query, "count("): + n := 0 + for _, s := range f.match(query[len("count(") : len(query)-1]) { + if !s.stale { + n++ + } + } + if n > 0 { + resp.Data.Result = append(resp.Data.Result, sample(map[string]string{}, n+f.countOffset)) + } + default: + matched := f.match(query) + capped := len(matched) > f.limit + if capped { + matched = matched[:f.limit] + } + for _, s := range matched { + if !s.stale { + resp.Data.Result = append(resp.Data.Result, sample(s.labels, 1)) + } + } + // The endpoint warns only when the capped output is full, so stale series hide it. + if capped && len(resp.Data.Result) == f.limit { + resp.Warnings = []string{fmt.Sprintf("results truncated due to series limit (%d).", f.limit)} + } + } + if err := json.NewEncoder(w).Encode(resp); err != nil { + f.t.Errorf("encoding response: %v", err) + } +} + +// badRequest reports a query the fake cannot parse; the handler runs off the test goroutine. +func (f *fakePromQL) badRequest(w http.ResponseWriter, format string, args ...interface{}) { + f.t.Errorf(format, args...) + http.Error(w, fmt.Sprintf(format, args...), http.StatusBadRequest) +} + +func (f *fakePromQL) match(selector string) []fakeSeries { + matchers, ok := parseFakeSelector(selector) + if !ok { + f.t.Errorf("unparseable selector %q", selector) + return nil + } + var out []fakeSeries + for _, s := range f.series { + ok := true + for _, m := range matchers { + v, present := s.labels[m.key] + switch m.op { + case "=": + ok = present && v == m.val + case "=~": + ok = regexp.MustCompile("^(?:" + m.val + ")$").MatchString(v) + case "!~": + ok = f.brokenNotRegex || !regexp.MustCompile("^(?:"+m.val+")$").MatchString(v) + default: + f.t.Errorf("unsupported matcher op %q", m.op) + ok = false + } + if !ok { + break + } + } + if ok { + out = append(out, s) + } + } + return out +} + +func (f *fakePromQL) queries() []fakeRequest { + f.mu.Lock() + defer f.mu.Unlock() + return append([]fakeRequest(nil), f.requests...) +} + +type fakeMatcher struct{ key, op, val string } + +func parseCountBy(q string) (string, string, bool) { + m := regexp.MustCompile(`^count by \("((?:[^"\\]|\\.)*)"\)\((.*)\)$`).FindStringSubmatch(q) + if m == nil { + return "", "", false + } + return unescape(m[1]), m[2], true +} + +func parseFakeSelector(s string) ([]fakeMatcher, bool) { + open := strings.IndexByte(s, '{') + if open < 0 { + return []fakeMatcher{{"__name__", "=", s}}, true + } + var out []fakeMatcher + if open > 0 { + out = append(out, fakeMatcher{"__name__", "=", s[:open]}) + } + re := regexp.MustCompile(`^\s*(?:"((?:[^"\\]|\\.)*)"|([a-zA-Z_][a-zA-Z0-9_]*))\s*(=~|!~|=)\s*"((?:[^"\\]|\\.)*)"\s*,?`) + body := s[open+1 : len(s)-1] + for strings.TrimSpace(body) != "" { + m := re.FindStringSubmatch(body) + if m == nil { + return nil, false + } + key := m[2] + if m[1] != "" { + key = unescape(m[1]) + } + out = append(out, fakeMatcher{key, m[3], unescape(m[4])}) + body = body[len(m[0]):] + } + return out, true +} + +func unescape(s string) string { + return strings.NewReplacer(`\\`, `\`, `\"`, `"`).Replace(s) +} + +func newFake(t *testing.T, limit int, series []fakeSeries) (*fakePromQL, *OtelMetricsClient) { + f := &fakePromQL{t: t, series: series, limit: limit} + srv := httptest.NewServer(http.HandlerFunc(f.handler)) + t.Cleanup(srv.Close) + c := newTestClient(srv.URL) + c.seriesLimit = limit + return f, c +} + +// podSeries builds kube_pod_info-like series: perNode pods on each node, plus extra. +func podSeries(nodes, perNode int) []fakeSeries { + var out []fakeSeries + for n := 0; n < nodes; n++ { + for p := 0; p < perNode; p++ { + out = append(out, fakeSeries{labels: map[string]string{ + "__name__": "kube_pod_info", + "@resource.k8s.cluster.name": "c1", + "@resource.k8s.node.name": fmt.Sprintf("ip-10-0-%d-1.ec2.internal", n), + "@resource.k8s.pod.name": fmt.Sprintf("pod-%d-%d", n, p), + "@resource.k8s.namespace.name": "default", + "@instrumentation.@name": "ksm", + "@aws.account": "123456789012", + "__type__": "Gauge", + }}) + } + } + return out +} + +const podSelector = `kube_pod_info{"@resource.k8s.cluster.name"="c1"}` + +func podNames(results []MetricResult) map[string]bool { + out := map[string]bool{} + for _, r := range results { + out[r.Labels.Resource["k8s.pod.name"]] = true + } + return out +} + +func countQueries(reqs []fakeRequest, pred func(string) bool) int { + n := 0 + for _, r := range reqs { + if pred(r.query) { + n++ + } + } + return n +} + +func isPlain(q string) bool { return !strings.HasPrefix(q, "count") } + +// assertOneEvalTime checks that the client pinned every request to the same evaluation time. +func assertOneEvalTime(t *testing.T, reqs []fakeRequest) { + t.Helper() + require.NotEmpty(t, reqs) + require.NotEmpty(t, reqs[0].time, "the client must pick the evaluation time: %s", reqs[0].query) + for _, r := range reqs { + assert.Equal(t, reqs[0].time, r.time, "every query is pinned to one evaluation time: %s", r.query) + } +} + +func TestQueryComplete_NotTruncatedIsOneDataQuery(t *testing.T) { + f, c := newFake(t, 10, podSeries(2, 3)) + + results, err := c.QueryComplete(context.Background(), podSelector) + require.NoError(t, err) + assert.Len(t, results, 6) + + reqs := f.queries() + require.Len(t, reqs, 2, "expected the selector plus one count(): %+v", reqs) + assert.ElementsMatch(t, []string{podSelector, "count(" + podSelector + ")"}, []string{reqs[0].query, reqs[1].query}) + assertOneEvalTime(t, reqs) +} + +func TestQueryComplete_SelectorAndCountRunTogether(t *testing.T) { + f := &fakePromQL{t: t, series: podSeries(2, 3), limit: 10} + var mu sync.Mutex + arrived := 0 + both := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + if arrived++; arrived == 2 { + close(both) + } + mu.Unlock() + // Answer only once both requests are in flight, which a client that waits for one cannot do. + select { + case <-both: + case <-time.After(2 * time.Second): + t.Errorf("request %q was not sent together with the other one", r.URL.Query().Get("query")) + } + f.handler(w, r) + })) + t.Cleanup(srv.Close) + c := newTestClient(srv.URL) + c.seriesLimit = 10 + + results, err := c.QueryComplete(context.Background(), podSelector) + require.NoError(t, err) + assert.Len(t, results, 6) + assertOneEvalTime(t, f.queries()) +} + +func TestEvalTimeAt(t *testing.T) { + assert.Equal(t, "1700000005", evalTimeAt(time.Unix(1700000010, 900e6)), "Unix seconds, evalTimeLag in the past") +} + +func TestQueryComplete_TruncatedShardsByNode(t *testing.T) { + series := podSeries(6, 5) + series[0].stale, series[7].stale = true, true // stale series still use up the cap + f, c := newFake(t, 10, series) + + results, err := c.QueryComplete(context.Background(), podSelector) + require.NoError(t, err) + assert.Len(t, results, 28) + assert.Len(t, podNames(results), 28, "every live pod exactly once") + + reqs := f.queries() + assert.Equal(t, 1, countQueries(reqs, func(q string) bool { + return q == `count by ("@resource.k8s.node.name")(`+podSelector+")" + })) + shards := countQueries(reqs, func(q string) bool { return isPlain(q) && strings.Contains(q, `"@resource.k8s.node.name"`) }) + assert.GreaterOrEqual(t, shards, 4, "28 series under a cap of 10 need at least 4 shards") + assert.Equal(t, 2, countQueries(reqs, func(q string) bool { return q == "count("+podSelector+")" }), "count() before and after") + assertOneEvalTime(t, reqs) +} + +func TestQueryComplete_AbsentNodeShardUsesNotRegex(t *testing.T) { + series := podSeries(3, 4) + for i := 0; i < 3; i++ { + series = append(series, fakeSeries{labels: map[string]string{ + "__name__": "kube_pod_info", + "@resource.k8s.cluster.name": "c1", + "@resource.k8s.pod.name": fmt.Sprintf("pending-%d", i), + }}) + } + f, c := newFake(t, 5, series) + + results, err := c.QueryComplete(context.Background(), podSelector) + require.NoError(t, err) + assert.Len(t, results, 15) + names := podNames(results) + for i := 0; i < 3; i++ { + assert.True(t, names[fmt.Sprintf("pending-%d", i)], "series without a node must be returned") + } + + reqs := f.queries() + assert.Equal(t, 1, countQueries(reqs, func(q string) bool { + return strings.Contains(q, `"@resource.k8s.node.name"!~".+"`) + })) + assert.Zero(t, countQueries(reqs, func(q string) bool { + return strings.Contains(q, `"@resource.k8s.node.name"=""`) + }), `="" does not match an absent label on this endpoint`) +} + +func TestQueryComplete_SingleNodeOverLimitSplitsFurther(t *testing.T) { + var series []fakeSeries + for cpu := 0; cpu < 8; cpu++ { + for _, mode := range []string{"idle", "user", "system"} { + series = append(series, fakeSeries{labels: map[string]string{ + "__name__": "node_cpu_seconds_total", + "@resource.k8s.cluster.name": "c1", + "@resource.k8s.node.name": "big-node", + "cpu": strconv.Itoa(cpu), + "mode": mode, + }}) + } + } + f, c := newFake(t, 10, series) + + results, err := c.QueryComplete(context.Background(), `node_cpu_seconds_total{"@resource.k8s.cluster.name"="c1"}`) + require.NoError(t, err) + assert.Len(t, results, 24) + assert.Positive(t, countQueries(f.queries(), func(q string) bool { return strings.HasPrefix(q, `count by ("cpu")`) }), + "one node holds every series, so the split must fall through to the cpu label") +} + +func TestQueryComplete_SkipsLabelWithTooManyValues(t *testing.T) { + var series []fakeSeries + for i := 0; i < 20; i++ { + series = append(series, fakeSeries{labels: map[string]string{ + "__name__": "m", + "id": strconv.Itoa(i), + "zone": fmt.Sprintf("z%d", i%3), + }}) + } + f, c := newFake(t, 10, series) + + results, err := c.QueryComplete(context.Background(), `m{}`) + require.NoError(t, err) + assert.Len(t, results, 20) + reqs := f.queries() + assert.Equal(t, 1, countQueries(reqs, func(q string) bool { return strings.HasPrefix(q, `count by ("id")`) })) + assert.Equal(t, 3, countQueries(reqs, func(q string) bool { return isPlain(q) && strings.Contains(q, `"zone"="z`) })) +} + +func TestQueryComplete_InconsistentCountIsError(t *testing.T) { + f, c := newFake(t, 10, podSeries(6, 5)) + f.countOffset = 3 // count() claims 33, but every split only accounts for 30 + + _, err := c.QueryComplete(context.Background(), podSelector) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrTruncated), "got %v", err) +} + +func TestQueryComplete_ShardTotalMismatchIsError(t *testing.T) { + series := podSeries(3, 4) + for i := 0; i < 3; i++ { + series = append(series, fakeSeries{labels: map[string]string{ + "__name__": "kube_pod_info", + "@resource.k8s.cluster.name": "c1", + "@resource.k8s.pod.name": fmt.Sprintf("pending-%d", i), + }}) + } + f, c := newFake(t, 5, series) + f.brokenNotRegex = true // the absent-node shard overlaps the node shards + + _, err := c.QueryComplete(context.Background(), podSelector) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrTruncated), "got %v", err) + assert.Contains(t, err.Error(), "shards returned 17 series but count() was 15 before and 15 after") +} + +func TestQueryComplete_UnsplittableIsError(t *testing.T) { + var series []fakeSeries + for i := 0; i < 20; i++ { + series = append(series, fakeSeries{labels: map[string]string{"__name__": "m", "id": strconv.Itoa(i)}}) + } + _, c := newFake(t, 10, series) + // id has 20 values, so its count by is capped too; nothing else varies. + _, err := c.QueryComplete(context.Background(), `m{}`) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrTruncated), "got %v", err) +} + +func TestQueryComplete_EmptyResultWithLiveSeries(t *testing.T) { + series := podSeries(2, 3) + stale := make([]fakeSeries, 0, 5) + for i := 0; i < 5; i++ { + s := podSeries(1, 1)[0] + s.labels["@resource.k8s.pod.name"] = fmt.Sprintf("done-%d", i) + s.stale = true + stale = append(stale, s) + } + f, c := newFake(t, 5, append(stale, series...)) + + results, err := c.QueryComplete(context.Background(), podSelector) + require.NoError(t, err) + assert.Len(t, results, 6, "the first query sees only stale series, yet six are live") + assertOneEvalTime(t, f.queries()) +} + +func TestQueryComplete_NonSelectorPassesThrough(t *testing.T) { + f, c := newFake(t, 10, podSeries(1, 2)) + _, err := c.QueryComplete(context.Background(), "count("+podSelector+")") + require.NoError(t, err) + assert.Len(t, f.queries(), 1) +} + +func TestQueryComplete_WarnedSelectorStillShards(t *testing.T) { + f, c := newFake(t, 10, podSeries(4, 5)) + + results, err := c.QueryComplete(context.Background(), podSelector) + require.NoError(t, err) + assert.Len(t, results, 20) + assert.Positive(t, countQueries(f.queries(), func(q string) bool { return strings.HasPrefix(q, "count by (") })) +} + +func TestQuery_TruncationWarningKeepsResults(t *testing.T) { + _, c := newFake(t, 3, podSeries(2, 3)) + results, err := c.Query(context.Background(), podSelector) + require.NoError(t, err, "Query stays a single best-effort request") + assert.Len(t, results, 3) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{"status":"success","warnings":["results truncated due to series limit (500)."],"data":{"resultType":"matrix","result":[{"metric":{"a":"b"},"values":[[1,"1"]]}]}}`) + })) + defer srv.Close() + rr, err := newTestClient(srv.URL).QueryRange(context.Background(), "up", time.Unix(0, 0), time.Unix(60, 0), time.Minute) + require.NoError(t, err) + assert.Len(t, rr, 1) +} + +func TestQueryCache_GetReturnsEverySeriesWhenTruncated(t *testing.T) { + f, c := newFake(t, 10, podSeries(5, 7)) + qc := NewQueryCache(c, "c1") + + results, err := qc.Get(context.Background(), "kube_pod_info") + require.NoError(t, err) + assert.Len(t, results, 35) + + n := len(f.queries()) + again, err := qc.Get(context.Background(), "kube_pod_info") + require.NoError(t, err) + assert.Len(t, again, 35) + assert.Len(t, f.queries(), n, "a cached metric must not be re-queried") +} + +func TestQueryCache_TruncatedHostTypeFailsFetch(t *testing.T) { + series := podSeries(1, 3) + for i := range series { + series[i].labels["@resource.host.type"] = "a" + } + for i := 0; i < 20; i++ { // host type b cannot be split, so its shard fetch fails + series = append(series, fakeSeries{labels: map[string]string{ + "__name__": "kube_pod_info", "@resource.k8s.cluster.name": "c1", + "@resource.host.type": "b", "id": strconv.Itoa(i), + }}) + } + _, c := newFake(t, 10, series) + qc := NewQueryCache(c, "c1", WithHostTypes([]string{"a", "b"})) + + results, err := qc.Get(context.Background(), "kube_pod_info") + require.Error(t, err, "the 3 series of host type a are not the whole result") + assert.True(t, errors.Is(err, ErrTruncated), "got %v", err) + assert.Empty(t, results) +} + +func TestQueryCache_ShardRequestsStayWithinConcurrencyLimit(t *testing.T) { + var series []fakeSeries + for _, ht := range []string{"a", "b"} { + for _, s := range podSeries(6, 5) { + s.labels["@resource.host.type"] = ht + s.labels["@resource.k8s.node.name"] += "-" + ht + series = append(series, s) + } + } + f := &fakePromQL{t: t, series: series, limit: 10} + var inflight, peak int + var mu sync.Mutex + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + inflight++ + if inflight > peak { + peak = inflight + } + mu.Unlock() + time.Sleep(5 * time.Millisecond) + f.handler(w, r) + mu.Lock() + inflight-- + mu.Unlock() + })) + t.Cleanup(srv.Close) + c := newTestClient(srv.URL) + c.seriesLimit = 10 + qc := NewQueryCache(c, "c1", WithHostTypes([]string{"a", "b"}), WithMaxConcurrency(2)) + + results, err := qc.Get(context.Background(), "kube_pod_info") + require.NoError(t, err) + assert.Len(t, results, 60) + assert.LessOrEqual(t, peak, 2, "shard re-fetches must share the cache's concurrency limit") +} + +func TestIsVectorSelector(t *testing.T) { + for s, want := range map[string]bool{ + `up`: true, + `kube_pod_info{"a"="b"}`: true, + `{"__name__"="k8s.pod.cpu.usage","a"="}"}`: true, + ` m{a=~"x|y", b!~".+"} `: true, + `k8s.pod.cpu.usage`: false, + `count(m{a="b"})`: false, + `m{a="b"}[5m]`: false, + `m{a="b"} or n{c="d"}`: false, + `rate(m[5m])`: false, + `m{a="b"} offset 5m`: false, + } { + assert.Equal(t, want, isVectorSelector(s), s) + } +} + +func TestWithMatcher(t *testing.T) { + assert.Equal(t, `m{x="1"}`, withMatcher("m", `x="1"`)) + assert.Equal(t, `m{x="1"}`, withMatcher("m{}", `x="1"`)) + assert.Equal(t, `m{a="b",x="1"}`, withMatcher(`m{a="b"}`, `x="1"`)) + assert.Equal(t, `{"__name__"="a.b",x="1"}`, withMatcher(`{"__name__"="a.b",}`, `x="1"`)) +} + +func TestPackShards(t *testing.T) { + groups := []shardGroup{{"a.b", 30}, {"c", 25}, {"", 4}, {"d", 10}, {"e", 50}} + shards := packShards("@resource.k8s.node.name", groups, 40) + + total := 0 + for _, s := range shards { + total += s.want + assert.LessOrEqual(t, s.want, 50) + } + assert.Equal(t, 119, total) + assert.Contains(t, shards, shard{matcher: `"@resource.k8s.node.name"!~".+"`, want: 4, exact: true}) + assert.Contains(t, shards, shard{matcher: `"@resource.k8s.node.name"="e"`, want: 50, exact: true}) + assert.Contains(t, shards, shard{matcher: `"@resource.k8s.node.name"=~"a\\.b|d"`, want: 40, exact: false}) + + var many []shardGroup + for i := 0; i < 200; i++ { + many = append(many, shardGroup{fmt.Sprintf("ip-10-0-%d-%d.ec2.internal", i/50, i), 1}) + } + for _, s := range packShards("@resource.k8s.node.name", many, 400) { + value := s.matcher[strings.Index(s.matcher, `"=~"`)+4 : len(s.matcher)-1] + assert.LessOrEqual(t, len(unescape(value)), maxMatcherValueLen, "the endpoint rejects longer matcher values") + } +} + +func TestQueryCache_ParallelGetWithFilterStaysWithinConcurrencyLimit(t *testing.T) { + f := &fakePromQL{t: t, series: podSeries(6, 5), limit: 10} + var inflight, peak int + var mu sync.Mutex + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + inflight++ + if inflight > peak { + peak = inflight + } + mu.Unlock() + time.Sleep(5 * time.Millisecond) + f.handler(w, r) + mu.Lock() + inflight-- + mu.Unlock() + })) + t.Cleanup(srv.Close) + c := newTestClient(srv.URL) + c.seriesLimit = 10 + qc := NewQueryCache(c, "c1", WithMaxConcurrency(2)) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + var wg sync.WaitGroup + for i := 0; i < 4; i++ { // more callers than slots, each needing several shard requests + wg.Add(1) + go func() { + defer wg.Done() + results, err := qc.GetWithFilter(ctx, "kube_pod_info", map[string]string{"@resource.k8s.namespace.name": "default"}) + assert.NoError(t, err, "a call that held a slot while its shards waited for one would deadlock here") + assert.Len(t, results, 30) + }() + } + wg.Wait() + assert.LessOrEqual(t, peak, 2, "shard re-fetches must share the cache's concurrency limit") +} diff --git a/util/otelmetrics/query_cache.go b/util/otelmetrics/query_cache.go index 51aef0f4f..3c39c75a3 100644 --- a/util/otelmetrics/query_cache.go +++ b/util/otelmetrics/query_cache.go @@ -32,6 +32,8 @@ type cacheEntry struct { // QueryCache provides session-scoped caching of PromQL queries. // Each unique metric name is queried exactly once; subsequent calls return cached data. // Concurrent requests for the same metric are deduplicated via singleflight. +// Every fetch goes through QueryComplete, so results are never silently truncated, and the +// concurrency limit bounds HTTP requests, so shard re-fetches stay inside it. type QueryCache struct { mu sync.RWMutex filtered map[string]cacheEntry @@ -142,6 +144,7 @@ func (qc *QueryCache) Get(ctx context.Context, metricName string) ([]MetricResul // fetchFiltered performs the actual HTTP query for a filtered metric. func (qc *QueryCache) fetchFiltered(ctx context.Context, metricName string) cacheEntry { + ctx = withRequestSlots(ctx, qc.sem) escaped := promqlEscaper.Replace(qc.cluster) sel := promqlMetricSelector(metricName) @@ -163,7 +166,7 @@ func (qc *QueryCache) fetchFiltered(ctx context.Context, metricName string) cach if clusterScopedOnly { promql := fmt.Sprintf(`%s"@resource.k8s.cluster.name"="%s","@resource.host.type"=""}`, sel, escaped) - r, err := qc.client.Query(ctx, promql) + r, err := qc.client.QueryComplete(ctx, promql) if err != nil { firstErr = err } else { @@ -181,15 +184,8 @@ func (qc *QueryCache) fetchFiltered(ctx context.Context, metricName string) cach wg.Add(1) go func(hostType string) { defer wg.Done() - select { - case qc.sem <- struct{}{}: - defer func() { <-qc.sem }() - case <-ctx.Done(): - ch <- queryResult{err: ctx.Err()} - return - } promql := fmt.Sprintf(`%s"@resource.k8s.cluster.name"="%s","@resource.host.type"="%s"}`, sel, escaped, hostType) - r, err := qc.client.Query(ctx, promql) + r, err := qc.client.QueryComplete(ctx, promql) ch <- queryResult{results: r, err: err} }(ht) } @@ -207,12 +203,13 @@ func (qc *QueryCache) fetchFiltered(ctx context.Context, metricName string) cach } } else { promql := fmt.Sprintf(`%s"@resource.k8s.cluster.name"="%s"}`, sel, escaped) - r, err := qc.client.Query(ctx, promql) + r, err := qc.client.QueryComplete(ctx, promql) results = r firstErr = err } - if len(results) == 0 && firstErr != nil { + // One failed host type makes the union incomplete, so it fails the whole fetch. + if firstErr != nil { slog.Debug("query failed", "metric", metricName, "error", firstErr) return cacheEntry{err: firstErr} } @@ -220,16 +217,10 @@ func (qc *QueryCache) fetchFiltered(ctx context.Context, metricName string) cach } // GetWithFilter returns results with additional PromQL label filters. Not cached. -// Shares the cache's concurrency limit, so a suite of parallel filtered queries -// doesn't burst past the endpoint's rate limit. +// Each HTTP request holds one slot of the cache's concurrency limit, so a suite of +// parallel filtered queries, and the shard re-fetches of a truncated one, stay inside it. func (qc *QueryCache) GetWithFilter(ctx context.Context, metricName string, extraFilters map[string]string) ([]MetricResult, error) { - select { - case qc.sem <- struct{}{}: - defer func() { <-qc.sem }() - case <-ctx.Done(): - return nil, ctx.Err() - } - return qc.client.Query(ctx, promqlMetricSelector(metricName)+qc.filterMatchers(extraFilters)+"}") + return qc.client.QueryComplete(withRequestSlots(ctx, qc.sem), promqlMetricSelector(metricName)+qc.filterMatchers(extraFilters)+"}") } // filterMatchers renders the cluster matcher plus extraFilters as PromQL label @@ -267,7 +258,7 @@ func (qc *QueryCache) GetUnfiltered(ctx context.Context, metricName string) ([]M } qc.mu.Unlock() - results, queryErr := qc.client.Query(ctx, metricName) + results, queryErr := qc.client.QueryComplete(withRequestSlots(ctx, qc.sem), metricName) entry := cacheEntry{results: results, err: queryErr} // Do not cache empty successful results; a transient miss would poison