From 2131a13979fe77ab61b5c68fbf6fd13a47386129 Mon Sep 17 00:00:00 2001 From: Petru Anica-Popa Date: Thu, 5 Nov 2026 17:54:27 +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/filters/filters_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/gpu_filters/gpu_filters_test.go | 6 +- .../multi_efa_filters_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 +- .../neuron_filters/neuron_filters_test.go | 2 +- 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 | 594 ++++++++++++ util/otelmetrics/complete_test.go | 869 ++++++++++++++++++ util/otelmetrics/query_cache.go | 33 +- 25 files changed, 1654 insertions(+), 64 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/filters/filters_test.go b/test/otel/filters/filters_test.go index 8ed556578..22ea6d2a1 100644 --- a/test/otel/filters/filters_test.go +++ b/test/otel/filters/filters_test.go @@ -368,7 +368,7 @@ func query(ctx context.Context, metric string, resource map[string]string) ([]ot for key, value := range resource { selector += fmt.Sprintf(`,%q=%q`, key, value) } - return client.Query(ctx, selector+"}") + return client.QueryComplete(ctx, selector+"}") } // logRecord is the part of an OTLP log record, as stored in CloudWatch Logs, 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/gpu_filters/gpu_filters_test.go b/test/otel/gpu_filters/gpu_filters_test.go index 16c73f822..d3e912026 100644 --- a/test/otel/gpu_filters/gpu_filters_test.go +++ b/test/otel/gpu_filters/gpu_filters_test.go @@ -118,7 +118,7 @@ func settled(t *testing.T, metric string) []otelmetrics.MetricResult { if i > 0 { time.Sleep(retryInterval) } - results, err := client.Query(context.Background(), selector) + results, err := client.QueryComplete(context.Background(), selector) if err != nil { problems = []string{err.Error()} continue @@ -155,7 +155,7 @@ func TestMetricNameFilterCase(t *testing.T) { time.Sleep(retryInterval) } var results []otelmetrics.MetricResult - results, err = client.Query(context.Background(), nodeSelector(kept)) + results, err = client.QueryComplete(context.Background(), nodeSelector(kept)) gpus = map[string]bool{} for _, r := range results { if g := uuid(r); g != "" { @@ -166,7 +166,7 @@ func TestMetricNameFilterCase(t *testing.T) { require.NoError(t, err, "querying %s", kept) require.Len(t, gpus, expectedGPUCount, "%s GPUs from node %s after %d attempts", kept, node, attempts) - results, err := client.Query(context.Background(), + results, err := client.QueryComplete(context.Background(), fmt.Sprintf(`{"__name__"=%q,"@resource.k8s.cluster.name"=%q}`, dropped, cfg.ClusterName)) require.NoError(t, err, "querying %s", dropped) assert.Empty(t, results, "%s present, but otelFilters.metricNames.exclude has it", dropped) diff --git a/test/otel/multi_efa_filters/multi_efa_filters_test.go b/test/otel/multi_efa_filters/multi_efa_filters_test.go index e74d13072..b58d8184a 100644 --- a/test/otel/multi_efa_filters/multi_efa_filters_test.go +++ b/test/otel/multi_efa_filters/multi_efa_filters_test.go @@ -180,7 +180,7 @@ func settled(t *testing.T, metric, nodeName string, checkFn func([]otelmetrics.M continue } } - results, err := client.Query(context.Background(), selector) + results, err := client.QueryComplete(context.Background(), selector) if err != nil { problems = []string{err.Error()} continue 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/neuron_filters/neuron_filters_test.go b/test/otel/neuron_filters/neuron_filters_test.go index d8ccc7604..d6901a3ff 100644 --- a/test/otel/neuron_filters/neuron_filters_test.go +++ b/test/otel/neuron_filters/neuron_filters_test.go @@ -173,7 +173,7 @@ func settled(t *testing.T, metric string) []otelmetrics.MetricResult { if i > 0 { time.Sleep(retryInterval) } - results, err := client.Query(context.Background(), selector) + results, err := client.QueryComplete(context.Background(), selector) if err != nil { problems = []string{err.Error()} continue 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..877a90e90 --- /dev/null +++ b/util/otelmetrics/complete.go @@ -0,0 +1,594 @@ +// 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" + "unicode/utf8" +) + +// 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 + // maxPrefixSteps bounds the one-character extensions of a value prefix that every series shares. + maxPrefixSteps = 8 + // prefixLabel is the label that label_replace writes a value prefix to; it must not clash with a real label. + prefixLabel = "__shard_prefix" + 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, 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 + prefix bool // value is a prefix: the group holds every value that starts with it +} + +type shard struct { + matcher string + want int + exact bool // the shard pins the label to one value (or to absent) + prefix string // every value of the label in the shard starts with this +} + +// 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. +// A label with more values than one response can list is split by value prefix instead, after the others. +// prefixes holds, per label, a prefix that every value in selector starts with. +func (c *OtelMetricsClient) split(ctx context.Context, selector, evalTime string, want int, sample []MetricResult, pinned map[string]bool, prefixes map[string]string, 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) + } + candidates := splitCandidates(sample, pinned) + var tooMany, rest []string + for i, label := range candidates { + if i == maxSplitCandidates { + rest = candidates[i:] + break + } + groups, err := c.countBy(ctx, selector, label, evalTime) + if errors.Is(err, ErrTruncated) { + tooMany = append(tooMany, label) + 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, prefixes, depth) + } + // Split by value prefix, which needs no list of every value: first the labels that had too + // many values to list, then the labels past maxSplitCandidates, up to maxSplitCandidates labels. + for i, label := range append(tooMany, rest...) { + if i == maxSplitCandidates { + break + } + groups, prefix, err := c.prefixGroups(ctx, selector, label, prefixes[label], evalTime, want, sample) + if err != nil { + return nil, err + } + if groups == nil { + continue + } + shards := packShards(label, groups, c.shardTarget(depth)) + for j := range shards { + if shards[j].prefix == "" { + shards[j].prefix = prefix + } + } + log.Printf("fetching %d shards by %s prefix %q depth=%d query=%s", len(shards), label, prefix, depth, selector) + return c.fetchShards(ctx, selector, evalTime, label, shards, pinned, prefixes, depth) + } + return nil, fmt.Errorf("%w: no label divides the result further (query %s)", ErrTruncated, selector) +} + +// prefixGroups counts the series of selector by the value prefix of label one character past +// prefix, which every value in selector already starts with. The group without a longer prefix +// holds the series whose value is exactly prefix, or that lack the label when prefix is "". +// While one group holds every series, the prefix grows: first to the longest prefix of the +// sample's common prefix that every series extends, then one character at a time. The groups +// must add up to want. It returns nil groups when the label does not divide selector. +func (c *OtelMetricsClient) prefixGroups(ctx context.Context, selector, label, prefix, evalTime string, want int, sample []MetricResult) ([]shardGroup, string, error) { + groups, longer, err := c.longestSharedPrefix(ctx, selector, label, prefix, evalTime, want, sample) + if err != nil { + return nil, "", err + } + if groups != nil { + prefix = longer + } + for step := 0; ; step++ { + if groups == nil { + var err error + groups, err = c.countByPrefix(ctx, selector, label, prefix, evalTime) + if errors.Is(err, ErrTruncated) { + return nil, "", nil + } + if err != nil { + return nil, "", err + } + } + if groupTotal(groups) != want { + return nil, "", nil // groups were cut short, or do not match count() + } + for _, g := range groups { + if g.prefix && !fitsMatcher(g.value) { + return nil, "", nil + } + } + if len(groups) > 1 { + return groups, prefix, nil + } + if !groups[0].prefix || step == maxPrefixSteps { + return nil, "", nil + } + prefix, groups = groups[0].value, nil + } +} + +// longestSharedPrefix finds the longest prefix of the sample's common prefix, past prefix, that +// every series of selector extends, and returns its groups. The common prefix of all values is a +// prefix of the sample's, and a prefix that holds implies that every shorter one holds, so a +// binary search over its rune boundaries finds it. It returns nil groups when none past prefix holds. +func (c *OtelMetricsClient) longestSharedPrefix(ctx context.Context, selector, label, prefix, evalTime string, want int, sample []MetricResult) ([]shardGroup, string, error) { + common := commonPrefix(sample, label, prefix) + var ends []int // byte lengths of the candidate prefixes, shortest first + for i := range common[len(prefix):] { + if n := len(prefix) + i; n > len(prefix) && fitsMatcher(common[:n]) { + ends = append(ends, n) + } + } + if len(common) > len(prefix) && fitsMatcher(common) { + ends = append(ends, len(common)) + } + holds := func(n int) ([]shardGroup, error) { + g, err := c.countByPrefix(ctx, selector, label, common[:n], evalTime) + if errors.Is(err, ErrTruncated) { + return nil, nil + } + if err != nil || groupTotal(g) != want || hasExact(g) { + return nil, err + } + return g, nil + } + var best []shardGroup + bestLen := len(prefix) + // The sample's whole common prefix usually holds, so try it first. + lo, hi := 0, len(ends)-1 + for first := true; lo <= hi; first = false { + mid := (lo + hi + 1) / 2 + if first { + mid = hi + } + g, err := holds(ends[mid]) + if err != nil { + return nil, "", err + } + if g != nil { + best, bestLen, lo = g, ends[mid], mid+1 + } else { + hi = mid - 1 + } + } + return best, common[:bestLen], nil +} + +// countByPrefix returns the series count per value prefix of label one character longer than +// prefix. Series whose value does not match keep no prefix label and form the non-prefix group. +func (c *OtelMetricsClient) countByPrefix(ctx context.Context, selector, label, prefix, evalTime string) ([]shardGroup, error) { + promql := fmt.Sprintf(`count by (%s)(label_replace(%s, %s, "$1", %s, %s))`, + quoteLabel(prefixLabel), selector, quoteLabel(prefixLabel), quoteLabel(label), + quoteLabel("("+regexp.QuoteMeta(prefix)+".).*")) + 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{value: prefix, count: int(r.Value)} + if v, ok := r.Labels.AllLabels()[prefixLabel]; ok && v != "" { + g = shardGroup{value: v, count: int(r.Value), prefix: true} + } + groups = append(groups, g) + } + return groups, nil +} + +func hasExact(groups []shardGroup) bool { + for _, g := range groups { + if !g.prefix { + return true + } + } + return false +} + +// fitsMatcher reports whether a prefix shard's regex for prefix stays within maxMatcherValueLen. +func fitsMatcher(prefix string) bool { + return len(prefixPattern(prefix)) <= maxMatcherValueLen +} + +func prefixPattern(prefix string) string { + return regexp.QuoteMeta(prefix) + ".*" +} + +// commonPrefix returns the longest prefix of label shared by every sampled value that starts with +// prefix, or prefix when there are none. It is only a guess, which prefixGroups checks. +func commonPrefix(sample []MetricResult, label, prefix string) string { + common, found := "", false + for _, r := range sample { + v, ok := r.Labels.AllLabels()[label] + if !ok || !strings.HasPrefix(v, prefix) { + continue + } + if !found { + common, found = v, true + continue + } + n := 0 + for n < len(common) && n < len(v) && common[n] == v[n] { + n++ + } + common = common[:n] + } + if !found { + return prefix + } + // Do not cut a multi-byte character in half. + for len(common) > len(prefix) && !utf8.ValidString(common) { + common = common[:len(common)-1] + } + return common +} + +func (c *OtelMetricsClient) fetchShards(ctx context.Context, selector, evalTime, label string, shards []shard, pinned map[string]bool, prefixes map[string]string, 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, subPrefixes := pinned, prefixes + if sh.exact { + subPinned = withPinned(pinned, label) + } else if sh.prefix != "" { + subPrefixes = withPrefix(prefixes, label, sh.prefix) + } + results, err = c.split(ctx, sub, evalTime, sh.want, results, subPinned, subPrefixes, 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. +// split counts the first maxSplitCandidates by value and the rest only by value prefix. +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] + }) + return append(out, keys...) +} + +// packShards bins label values (or value prefixes) 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 { + groups []shardGroup + want, chars int + } + var bins []bin + for _, g := range sorted { + if g.value == "" && !g.prefix { + out = append(out, shard{matcher: quoteLabel(label) + `!~".+"`, want: g.count, exact: true}) + continue + } + chars := len(groupPattern(g)) + placed := false + for i := range bins { + if bins[i].want+g.count <= target && bins[i].chars+1+chars <= maxMatcherValueLen { + bins[i].groups = append(bins[i].groups, g) + bins[i].want += g.count + bins[i].chars += 1 + chars + placed = true + break + } + } + if !placed { + bins = append(bins, bin{groups: []shardGroup{g}, want: g.count, chars: chars}) + } + } + for _, b := range bins { + sh := shard{matcher: binMatcher(label, b.groups), want: b.want} + if len(b.groups) == 1 { + if b.groups[0].prefix { + sh.prefix = b.groups[0].value + } else { + sh.exact = true + } + } + out = append(out, sh) + } + return out +} + +// groupPattern is the regex alternative that matches a group's values. +func groupPattern(g shardGroup) string { + if g.prefix { + return prefixPattern(g.value) + } + return regexp.QuoteMeta(g.value) +} + +// binMatcher matches every value of the groups in a bin: = for one exact value, else a regex. +func binMatcher(label string, groups []shardGroup) string { + if len(groups) == 1 && !groups[0].prefix { + return fmt.Sprintf(`%s="%s"`, quoteLabel(label), promqlEscaper.Replace(groups[0].value)) + } + alts := make([]string, len(groups)) + for i, g := range groups { + alts[i] = groupPattern(g) + } + 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 withPrefix(prefixes map[string]string, label, prefix string) map[string]string { + out := make(map[string]string, len(prefixes)+1) + for k, v := range prefixes { + out[k] = v + } + out[label] = prefix + return out +} + +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..bbd8a0ab3 --- /dev/null +++ b/util/otelmetrics/complete_test.go @@ -0,0 +1,869 @@ +// 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 + prefixOffset int // added to the first prefix group of a label_replace count, likewise + 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{} + if lr, isReplace := parseLabelReplace(inner); isReplace { + if lr == nil || lr.dst != label { + f.badRequest(w, "unexpected label_replace query %q", query) + return + } + re := regexp.MustCompile("^(?:" + lr.regex + ")$") + for _, s := range f.match(lr.selector) { + if !s.stale { + key := "" + if m := re.FindStringSubmatch(s.labels[lr.src]); m != nil { + key = m[1] + } + groups[key]++ + } + } + first := "" + for k := range groups { + if k != "" && (first == "" || k < first) { + first = k + } + } + if first != "" { + groups[first] += f.prefixOffset + } + } else { + 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 +} + +type fakeLabelReplace struct{ selector, dst, src, regex string } + +// parseLabelReplace parses label_replace(selector, "dst", "$1", "src", "regex"). It reports +// whether s is a label_replace call, and returns nil when the call has another shape. +func parseLabelReplace(s string) (*fakeLabelReplace, bool) { + if !strings.HasPrefix(s, "label_replace(") { + return nil, false + } + str := `"((?:[^"\\]|\\.)*)"` + m := regexp.MustCompile(`^label_replace\((.*), ` + str + `, "\$1", ` + str + `, ` + str + `\)$`).FindStringSubmatch(s) + if m == nil { + return nil, true + } + return &fakeLabelReplace{selector: m[1], dst: unescape(m[2]), src: unescape(m[3]), regex: unescape(m[4])}, 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`) })) +} + +// unsplittablePrefix is a common value prefix too long for a matcher, so no prefix shard can hold it. +var unsplittablePrefix = strings.Repeat("x", maxMatcherValueLen) + +const karpenterSelector = `karpenter_cloudprovider_instance_type_cpu_cores{"@resource.k8s.cluster.name"="c1"}` + +// karpenterSeries builds one series per instance type, all on the Karpenter controller's one resource. +func karpenterSeries(types []string) []fakeSeries { + out := make([]fakeSeries, 0, len(types)) + for _, it := range types { + out = append(out, fakeSeries{labels: map[string]string{ + "__name__": "karpenter_cloudprovider_instance_type_cpu_cores", + "@resource.k8s.cluster.name": "c1", + "@resource.k8s.node.name": "ip-10-0-0-1.ec2.internal", + "@resource.k8s.pod.name": "karpenter-0", + "instance_type": it, + }}) + } + return out +} + +// instanceTypes returns 4 sizes of each family, in family order. +func instanceTypes(families ...string) []string { + var out []string + for _, f := range families { + for _, size := range []string{"large", "xlarge", "2xlarge", "4xlarge"} { + out = append(out, f+"."+size) + } + } + return out +} + +// assertEachInstanceTypeOnce checks that results hold every one of types exactly once. +func assertEachInstanceTypeOnce(t *testing.T, types []string, results []MetricResult) { + t.Helper() + got := map[string]int{} + for _, r := range results { + got[r.Labels.AllLabels()["instance_type"]]++ + } + want := map[string]int{} + for _, it := range types { + want[it] = 1 + } + assert.Equal(t, want, got) +} + +func isPrefixCount(q string) bool { return strings.Contains(q, `)(label_replace(`) } + +func TestQueryComplete_LabelWithMoreValuesThanLimitSplitsByPrefix(t *testing.T) { + types := instanceTypes("c5", "c6i", "m5", "m6i", "r5", "r6i", "t3") + f, c := newFake(t, 10, karpenterSeries(types)) + + results, err := c.QueryComplete(context.Background(), karpenterSelector) + require.NoError(t, err, "instance_type has 28 values under a cap of 10, so its count by is capped too") + assertEachInstanceTypeOnce(t, types, results) + + reqs := f.queries() + assert.Equal(t, 1, countQueries(reqs, func(q string) bool { + return q == `count by ("__shard_prefix")(label_replace(`+karpenterSelector+`, "__shard_prefix", "$1", "instance_type", "(.).*"))` + })) + for _, family := range []string{"c", "m", "r", "t"} { + assert.Equal(t, 1, countQueries(reqs, func(q string) bool { + return isPlain(q) && strings.Contains(q, `"instance_type"=~"`+family+`.*"`) + }), "one shard per first letter: %s", family) + } + assertOneEvalTime(t, reqs) +} + +func TestQueryComplete_PrefixShardOverLimitUsesLongerPrefix(t *testing.T) { + types := []string{ + "m5.large", "m5.xlarge", "m5.2xlarge", "m5.4xlarge", "m5.8xlarge", "m5.12xlarge", "m5.16xlarge", "m5.24xlarge", "m5.metal", + "m5d.large", "m5d.xlarge", "m5d.2xlarge", "m5d.4xlarge", "m6i.large", "m6i.xlarge", + "c5.large", "c5.xlarge", "r5.large", "t3.micro", "t3.small", + } + f, c := newFake(t, 10, karpenterSeries(types)) + + results, err := c.QueryComplete(context.Background(), karpenterSelector) + require.NoError(t, err) + assertEachInstanceTypeOnce(t, types, results) + + reqs := f.queries() + for _, matcher := range []string{`"instance_type"=~"m.*"`, `"instance_type"=~"m5.*"`, `"instance_type"=~"m5\\..*"`} { + assert.Positive(t, countQueries(reqs, func(q string) bool { return isPlain(q) && strings.HasSuffix(q, ","+matcher+"}") }), + "the 15 m types and then the 13 m5 types are each over the cap of 10: %s", matcher) + } +} + +func TestQueryComplete_LongSharedPrefixPastSampleCommonPrefix(t *testing.T) { + var types []string + for i := 0; i < 40; i++ { + types = append(types, fmt.Sprintf("arn:aws:ec2:us-east-1:111122223333:volume/vol-0%02d", i)) + } + f, c := newFake(t, 10, karpenterSeries(types)) + + results, err := c.QueryComplete(context.Background(), karpenterSelector) + require.NoError(t, err, "the sample holds vol-000 to vol-009, so its common prefix vol-00 overshoots the one that all values share") + assertEachInstanceTypeOnce(t, types, results) + for d := 0; d < 4; d++ { + assert.Equal(t, 1, countQueries(f.queries(), func(q string) bool { + return isPlain(q) && strings.HasSuffix(q, fmt.Sprintf(`,"instance_type"=~"arn:aws:ec2:us-east-1:111122223333:volume/vol-0%d.*"}`, d)) + }), "one shard per digit after vol-0: %d", d) + } +} + +func TestQueryComplete_PrefixShardsEscapeRegexMetacharacters(t *testing.T) { + types := []string{"a.b1", "a.b2", "a.b3", "a+b1", "a+b2", "a|b1", "a(b1", `a"b1`, `a\b1`, "x"} + f, c := newFake(t, 6, karpenterSeries(types)) + + results, err := c.QueryComplete(context.Background(), karpenterSelector) + require.NoError(t, err, "an unescaped prefix such as a. or a| would overlap other shards") + assertEachInstanceTypeOnce(t, types, results) + assert.Positive(t, countQueries(f.queries(), func(q string) bool { return strings.Contains(q, `"instance_type"=~"a\\..*"`) })) +} + +func TestQueryComplete_PrefixTotalsThatDoNotAddUpAreError(t *testing.T) { + for _, offset := range []int{-1, 1} { + t.Run(strconv.Itoa(offset), func(t *testing.T) { + f, c := newFake(t, 10, karpenterSeries(instanceTypes("c5", "c6i", "m5", "m6i", "r5", "r6i", "t3"))) + f.prefixOffset = offset + + results, err := c.QueryComplete(context.Background(), karpenterSelector) + require.Error(t, err) + assert.True(t, errors.Is(err, ErrTruncated), "got %v", err) + assert.Contains(t, err.Error(), "no label divides the result further") + assert.Empty(t, results, "no partial result") + assert.Positive(t, countQueries(f.queries(), isPrefixCount), "the prefix split was tried and rejected") + }) + } +} + +func TestQueryComplete_PrefixSplitTriesLabelPastCandidateLimit(t *testing.T) { + types := instanceTypes("c5", "c6i", "m5", "m6i", "r5") + series := karpenterSeries(types) + for i := range series { + // Two labels that rank before instance_type, have too many values and share a prefix too long to split. + series[i].labels["a_uid"] = unsplittablePrefix + strconv.Itoa(i) + series[i].labels["b_uid"] = unsplittablePrefix + strconv.Itoa(i) + } + f, c := newFake(t, 10, series) + + results, err := c.QueryComplete(context.Background(), karpenterSelector) + require.NoError(t, err, "instance_type is the fourth candidate, past maxSplitCandidates") + assertEachInstanceTypeOnce(t, types, results) + assert.Zero(t, countQueries(f.queries(), func(q string) bool { return strings.HasPrefix(q, `count by ("instance_type")`) }), + "labels past the cut are split by prefix only") +} + +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": unsplittablePrefix + strconv.Itoa(i)}}) + } + _, c := newFake(t, 10, series) + // id has 20 values, so its count by is capped too, and they share a prefix too long for a + // matcher; 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_NoSeriesIsEmptyResult(t *testing.T) { + f, c := newFake(t, 10, podSeries(1, 2)) + + results, err := c.QueryComplete(context.Background(), `{"__name__"="absent_metric"}`) + require.NoError(t, err, "count() of nothing is an empty vector, not an error") + assert.Empty(t, results) + assert.Len(t, f.queries(), 2, "expected the selector plus one count()") +} + +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": unsplittablePrefix + 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{{value: "a.b", count: 30}, {value: "c", count: 25}, {value: "", count: 4}, {value: "d", count: 10}, {value: "e", count: 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{value: fmt.Sprintf("ip-10-0-%d-%d.ec2.internal", i/50, i), count: 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