From 63a5e4fd59f304235a90f075bdbc964c6add5d50 Mon Sep 17 00:00:00 2001 From: Martin Montes Date: Sun, 17 Mar 2024 12:24:47 +0100 Subject: [PATCH] Support for prometheus scrape timeout in probe endpoint Signed-off-by: Martin Montes --- mysqld_exporter.go | 57 ++++++++++++++++++++---------- mysqld_exporter_test.go | 77 +++++++++++++++++++++++++++++++++++++++++ probe.go | 16 +++++++++ 3 files changed, 131 insertions(+), 19 deletions(-) diff --git a/mysqld_exporter.go b/mysqld_exporter.go index 76fb107bf..e3fc0139f 100644 --- a/mysqld_exporter.go +++ b/mysqld_exporter.go @@ -15,6 +15,7 @@ package main import ( "context" + "fmt" "net/http" "os" "strconv" @@ -127,6 +128,32 @@ func filterScrapers(scrapers []collector.Scraper, collectParams []string) []coll return filteredScrapers } +func getScrapeTimeoutSeconds(r *http.Request, offset float64) (float64, error) { + var timeoutSeconds float64 + if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" { + var err error + timeoutSeconds, err = strconv.ParseFloat(v, 64) + if err != nil { + return 0, fmt.Errorf("failed to parse timeout from Prometheus header: %v", err) + } + } + if timeoutSeconds == 0 { + return 0, nil + } + if timeoutSeconds < 0 { + return 0, fmt.Errorf("timeout value from Prometheus header is invalid: %f", timeoutSeconds) + } + + if offset >= timeoutSeconds { + // Ignore timeout offset if it doesn't leave time to scrape. + return 0, fmt.Errorf("timeout offset (%f) should be lower than prometheus scrape timeout (%f)", offset, timeoutSeconds) + } else { + // Subtract timeout offset from timeout. + timeoutSeconds -= offset + } + return timeoutSeconds, nil +} + func init() { prometheus.MustRegister(version.NewCollector("mysqld_exporter")) } @@ -155,25 +182,17 @@ func newHandler(scrapers []collector.Scraper, logger log.Logger) http.HandlerFun // Use request context for cancellation when connection gets closed. ctx := r.Context() // If a timeout is configured via the Prometheus header, add it to the context. - if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" { - timeoutSeconds, err := strconv.ParseFloat(v, 64) - if err != nil { - level.Error(logger).Log("msg", "Failed to parse timeout from Prometheus header", "err", err) - } else { - if *timeoutOffset >= timeoutSeconds { - // Ignore timeout offset if it doesn't leave time to scrape. - level.Error(logger).Log("msg", "Timeout offset should be lower than prometheus scrape timeout", "offset", *timeoutOffset, "prometheus_scrape_timeout", timeoutSeconds) - } else { - // Subtract timeout offset from timeout. - timeoutSeconds -= *timeoutOffset - } - // Create new timeout context with request context as parent. - var cancel context.CancelFunc - ctx, cancel = context.WithTimeout(ctx, time.Duration(timeoutSeconds*float64(time.Second))) - defer cancel() - // Overwrite request with timeout context. - r = r.WithContext(ctx) - } + timeoutSeconds, err := getScrapeTimeoutSeconds(r, *timeoutOffset) + if err != nil { + level.Error(logger).Log("msg", "Error getting timeout from Prometheus header", "err", err) + } + if timeoutSeconds > 0 { + // Create new timeout context with request context as parent. + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, time.Duration(timeoutSeconds*float64(time.Second))) + defer cancel() + // Overwrite request with timeout context. + r = r.WithContext(ctx) } filteredScrapers := filterScrapers(scrapers, collect) diff --git a/mysqld_exporter_test.go b/mysqld_exporter_test.go index 9417c0d56..a631f7da4 100644 --- a/mysqld_exporter_test.go +++ b/mysqld_exporter_test.go @@ -304,3 +304,80 @@ func Test_filterScrapers(t *testing.T) { }) } } + +func Test_getScrapeTimeoutSeconds(t *testing.T) { + type args struct { + timeoutHeader string + offset float64 + } + tests := []struct { + name string + args args + wantTimeout float64 + wantErr bool + }{ + {"no_timeout_header", + args{}, + 0, false, + }, + {"zero_timeout_header", + args{ + timeoutHeader: "0", + }, + 0, false, + }, + {"negative_timeout_header", + args{ + timeoutHeader: "-5", + }, + 0, true, + }, + {"offset_greater_than_timeout", + args{ + timeoutHeader: "5", + offset: 6, + }, + 0, true, + }, + {"offset_equal_timeout", + args{ + timeoutHeader: "5", + offset: 5, + }, + 0, true, + }, + {"offset_less_than_timeout", + args{ + timeoutHeader: "5", + offset: 1, + }, + 4, false, + }, + {"no_offset", + args{ + timeoutHeader: "5", + }, + 5, false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + request, err := http.NewRequest(http.MethodGet, "", nil) + if err != nil { + t.Fatalf("unexpected error creating http request: %v", err) + } + request.Header.Set("X-Prometheus-Scrape-Timeout-Seconds", tt.args.timeoutHeader) + + timeout, err := getScrapeTimeoutSeconds(request, tt.args.offset) + if err != nil && !tt.wantErr { + t.Fatalf("unexpected error: %v", err) + } + if err == nil && tt.wantErr { + t.Fatal("expecting an error, got nil") + } + if timeout != tt.wantTimeout { + t.Fatalf("unexpected timeout, got '%f' but expected '%f'", timeout, tt.wantTimeout) + } + }) + } +} diff --git a/probe.go b/probe.go index a7c6606a5..c101aee6f 100644 --- a/probe.go +++ b/probe.go @@ -14,8 +14,10 @@ package main import ( + "context" "fmt" "net/http" + "time" "github.com/go-kit/log" "github.com/go-kit/log/level" @@ -54,6 +56,20 @@ func handleProbe(scrapers []collector.Scraper, logger log.Logger) http.HandlerFu return } + // If a timeout is configured via the Prometheus header, add it to the context. + timeoutSeconds, err := getScrapeTimeoutSeconds(r, *timeoutOffset) + if err != nil { + level.Error(logger).Log("msg", "Error getting timeout from Prometheus header", "err", err) + } + if timeoutSeconds > 0 { + // Create new timeout context with request context as parent. + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, time.Duration(timeoutSeconds*float64(time.Second))) + defer cancel() + // Overwrite request with timeout context. + r = r.WithContext(ctx) + } + filteredScrapers := filterScrapers(scrapers, collectParams) registry := prometheus.NewRegistry()