mirror of https://github.com/grafana/loki
Add rate limiting and metrics to hedging (#4860)
* Add rate limiting and metrics to hedging Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com> * Add changelog entry Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com> * Allow to pass a custom registerer Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>pull/4872/head
parent
77dcb16b4e
commit
c33a651645
@ -0,0 +1,90 @@ |
||||
package hedging |
||||
|
||||
import ( |
||||
"net/http" |
||||
"strings" |
||||
"testing" |
||||
"time" |
||||
|
||||
"github.com/prometheus/client_golang/prometheus" |
||||
"github.com/prometheus/client_golang/prometheus/testutil" |
||||
"github.com/stretchr/testify/require" |
||||
"go.uber.org/atomic" |
||||
) |
||||
|
||||
type RoundTripperFunc func(*http.Request) (*http.Response, error) |
||||
|
||||
func (fn RoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { |
||||
return fn(req) |
||||
} |
||||
|
||||
func resetMetrics() { |
||||
reg := prometheus.NewRegistry() |
||||
prometheus.DefaultRegisterer = reg |
||||
prometheus.DefaultGatherer = reg |
||||
initMetrics() |
||||
} |
||||
|
||||
func TestHedging(t *testing.T) { |
||||
resetMetrics() |
||||
cfg := &Config{ |
||||
At: time.Duration(1), |
||||
UpTo: 3, |
||||
MaxPerSecond: 1000, |
||||
} |
||||
count := atomic.NewInt32(0) |
||||
client := cfg.Client(&http.Client{ |
||||
Transport: RoundTripperFunc(func(r *http.Request) (*http.Response, error) { |
||||
count.Inc() |
||||
time.Sleep(200 * time.Millisecond) |
||||
return &http.Response{ |
||||
StatusCode: http.StatusOK, |
||||
}, nil |
||||
}), |
||||
}) |
||||
_, _ = client.Get("http://example.com") |
||||
|
||||
require.Equal(t, int32(3), count.Load()) |
||||
require.NoError(t, testutil.GatherAndCompare(prometheus.DefaultGatherer, |
||||
strings.NewReader(` |
||||
# HELP hedged_requests_rate_limited_total The total number of hedged requests rejected via rate limiting. |
||||
# TYPE hedged_requests_rate_limited_total counter |
||||
hedged_requests_rate_limited_total 0 |
||||
# HELP hedged_requests_total The total number of hedged requests. |
||||
# TYPE hedged_requests_total counter |
||||
hedged_requests_total 2 |
||||
`, |
||||
), "hedged_requests_total", "hedged_requests_rate_limited_total")) |
||||
} |
||||
|
||||
func TestHedgingRateLimit(t *testing.T) { |
||||
resetMetrics() |
||||
cfg := &Config{ |
||||
At: time.Duration(1), |
||||
UpTo: 20, |
||||
MaxPerSecond: 1, |
||||
} |
||||
count := atomic.NewInt32(0) |
||||
client := cfg.Client(&http.Client{ |
||||
Transport: RoundTripperFunc(func(r *http.Request) (*http.Response, error) { |
||||
count.Inc() |
||||
time.Sleep(200 * time.Millisecond) |
||||
return &http.Response{ |
||||
StatusCode: http.StatusOK, |
||||
}, nil |
||||
}), |
||||
}) |
||||
_, _ = client.Get("http://example.com") |
||||
|
||||
require.Equal(t, int32(2), count.Load()) |
||||
require.NoError(t, testutil.GatherAndCompare(prometheus.DefaultGatherer, |
||||
strings.NewReader(` |
||||
# HELP hedged_requests_rate_limited_total The total number of hedged requests rejected via rate limiting. |
||||
# TYPE hedged_requests_rate_limited_total counter |
||||
hedged_requests_rate_limited_total 18 |
||||
# HELP hedged_requests_total The total number of hedged requests. |
||||
# TYPE hedged_requests_total counter |
||||
hedged_requests_total 1 |
||||
`, |
||||
), "hedged_requests_total", "hedged_requests_rate_limited_total")) |
||||
} |
||||
Loading…
Reference in new issue