Skip to content

Commit 2606fff

Browse files
coleenquadroscoleenquadros
committed
add handler for ratelimiting
Signed-off-by: Coleen Iona Quadros <[email protected]>
1 parent a92611b commit 2606fff

File tree

5 files changed

+166
-71
lines changed

5 files changed

+166
-71
lines changed

Diff for: ‎api/metrics/legacy/http.go

+17-8
Original file line numberDiff line numberDiff line change
@@ -141,23 +141,32 @@ func NewHandler(url *url.URL, upstreamCA []byte, upstreamCert *stdtls.Certificat
141141
}
142142

143143
r.Group(func(r chi.Router) {
144+
r.Use(func(handler http.Handler) http.Handler {
145+
return server.InjectLabelsCtx(
146+
prometheus.Labels{"group": "metricslegacy", "handler": "query"},
147+
handler,
148+
)
149+
})
144150
r.Use(c.queryMiddlewares...)
145151
r.Handle(QueryRoute,
146152
otelhttp.WithRouteTag(
147153
c.spanRoutePrefix+QueryRoute,
148-
server.InjectLabelsCtx(
149-
prometheus.Labels{"group": "metricslegacy", "handler": "query"},
150-
legacyProxy,
151-
),
154+
legacyProxy,
152155
),
153156
)
157+
})
158+
r.Group(func(r chi.Router) {
159+
r.Use(func(handler http.Handler) http.Handler {
160+
return server.InjectLabelsCtx(
161+
prometheus.Labels{"group": "metricslegacy", "handler": "query_range"},
162+
handler,
163+
)
164+
})
165+
r.Use(c.queryMiddlewares...)
154166
r.Handle(QueryRangeRoute,
155167
otelhttp.WithRouteTag(
156168
c.spanRoutePrefix+QueryRangeRoute,
157-
server.InjectLabelsCtx(
158-
prometheus.Labels{"group": "metricslegacy", "handler": "query_range"},
159-
legacyProxy,
160-
),
169+
legacyProxy,
161170
),
162171
)
163172
})

Diff for: ‎api/metrics/v1/http.go

+108-53
Original file line numberDiff line numberDiff line change
@@ -213,24 +213,34 @@ func NewHandler(endpoints Endpoints, upstreamCA []byte, upstreamCert *stdtls.Cer
213213
}
214214
}
215215
r.Group(func(r chi.Router) {
216+
r.Use(func(handler http.Handler) http.Handler {
217+
return server.InjectLabelsCtx(
218+
prometheus.Labels{"group": "metricsv1", "handler": "query"},
219+
handler,
220+
)
221+
})
216222
r.Use(c.queryMiddlewares...)
217223
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
218224
r.Handle(QueryRoute,
219225
otelhttp.WithRouteTag(
220226
c.spanRoutePrefix+QueryRoute,
221-
server.InjectLabelsCtx(
222-
prometheus.Labels{"group": "metricsv1", "handler": "query"},
223-
proxyQuery,
224-
),
227+
proxyQuery,
225228
),
226229
)
230+
})
231+
r.Group(func(r chi.Router) {
232+
r.Use(func(handler http.Handler) http.Handler {
233+
return server.InjectLabelsCtx(
234+
prometheus.Labels{"group": "metricsv1", "handler": "query_range"},
235+
handler,
236+
)
237+
})
238+
r.Use(c.queryMiddlewares...)
239+
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
227240
r.Handle(QueryRangeRoute,
228241
otelhttp.WithRouteTag(
229242
c.spanRoutePrefix+QueryRangeRoute,
230-
server.InjectLabelsCtx(
231-
prometheus.Labels{"group": "metricsv1", "handler": "query_range"},
232-
proxyQuery,
233-
),
243+
proxyQuery,
234244
),
235245
)
236246
})
@@ -258,45 +268,69 @@ func NewHandler(endpoints Endpoints, upstreamCA []byte, upstreamCert *stdtls.Cer
258268
}
259269
}
260270
r.Group(func(r chi.Router) {
271+
r.Use(func(handler http.Handler) http.Handler {
272+
return server.InjectLabelsCtx(
273+
prometheus.Labels{"group": "metricsv1", "handler": "series"},
274+
handler,
275+
)
276+
})
261277
r.Use(c.readMiddlewares...)
262278
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
263-
r.Handle(RulesRoute, otelhttp.WithRouteTag(c.spanRoutePrefix+RulesRoute, proxyRead))
264279
r.Handle(SeriesRoute,
265280
otelhttp.WithRouteTag(
266281
c.spanRoutePrefix+SeriesRoute,
267-
server.InjectLabelsCtx(
268-
prometheus.Labels{"group": "metricsv1", "handler": "series"},
269-
proxyRead,
270-
),
282+
proxyRead,
271283
),
272284
)
285+
})
286+
r.Group(func(r chi.Router) {
287+
r.Use(func(handler http.Handler) http.Handler {
288+
return server.InjectLabelsCtx(
289+
prometheus.Labels{"group": "metricsv1", "handler": "label_names"},
290+
handler,
291+
)
292+
})
293+
r.Use(c.readMiddlewares...)
294+
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
273295
r.Handle(LabelNamesRoute,
274296
otelhttp.WithRouteTag(
275297
c.spanRoutePrefix+LabelNamesRoute,
276-
server.InjectLabelsCtx(
277-
prometheus.Labels{"group": "metricsv1", "handler": "label_names"},
278-
proxyRead,
279-
),
298+
proxyRead,
280299
),
281300
)
301+
})
302+
r.Group(func(r chi.Router) {
303+
r.Use(func(handler http.Handler) http.Handler {
304+
return server.InjectLabelsCtx(
305+
prometheus.Labels{"group": "metricsv1", "handler": "label_values"},
306+
handler,
307+
)
308+
})
309+
r.Use(c.readMiddlewares...)
310+
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
282311
r.Handle(LabelValuesRoute,
283312
otelhttp.WithRouteTag(
284313
c.spanRoutePrefix+LabelValuesRoute,
285-
server.InjectLabelsCtx(
286-
prometheus.Labels{"group": "metricsv1", "handler": "label_values"},
287-
proxyRead,
288-
),
314+
proxyRead,
289315
),
290316
)
317+
})
318+
319+
r.Group(func(r chi.Router) {
320+
r.Use(func(handler http.Handler) http.Handler {
321+
return server.InjectLabelsCtx(
322+
prometheus.Labels{"group": "metricsv1", "handler": "rules"},
323+
handler,
324+
)
325+
})
326+
r.Use(c.readMiddlewares...)
327+
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
291328
// Thanos Query Rules API supports matchers from v0.25 so the WithEnforceTenancyOnMatchers
292329
// middleware will not work here if prior versions are used.
293330
r.Handle(RulesRoute,
294331
otelhttp.WithRouteTag(
295332
c.spanRoutePrefix+RulesRoute,
296-
server.InjectLabelsCtx(
297-
prometheus.Labels{"group": "metricsv1", "handler": "rules"},
298-
proxyRead,
299-
),
333+
proxyRead,
300334
),
301335
)
302336
})
@@ -319,15 +353,18 @@ func NewHandler(endpoints Endpoints, upstreamCA []byte, upstreamCert *stdtls.Cer
319353
}
320354
}
321355
r.Group(func(r chi.Router) {
356+
r.Use(func(handler http.Handler) http.Handler {
357+
return server.InjectLabelsCtx(
358+
prometheus.Labels{"group": "metricsv1", "handler": "ui"},
359+
handler,
360+
)
361+
})
322362
r.Use(c.uiMiddlewares...)
323363
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
324364
r.Mount(UIRoute,
325365
otelhttp.WithRouteTag(
326366
c.spanRoutePrefix+UIRoute,
327-
server.InjectLabelsCtx(
328-
prometheus.Labels{"group": "metricsv1", "handler": "ui"},
329-
uiProxy,
330-
),
367+
uiProxy,
331368
),
332369
)
333370
})
@@ -357,15 +394,18 @@ func NewHandler(endpoints Endpoints, upstreamCA []byte, upstreamCert *stdtls.Cer
357394
}
358395
}
359396
r.Group(func(r chi.Router) {
397+
r.Use(func(handler http.Handler) http.Handler {
398+
return server.InjectLabelsCtx(
399+
prometheus.Labels{"group": "metricsv1", "handler": "receive"},
400+
handler,
401+
)
402+
})
360403
r.Use(c.writeMiddlewares...)
361404
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
362405
r.Handle(ReceiveRoute,
363406
otelhttp.WithRouteTag(
364407
c.spanRoutePrefix+ReceiveRoute,
365-
server.InjectLabelsCtx(
366-
prometheus.Labels{"group": "metricsv1", "handler": "receive"},
367-
proxyWrite,
368-
),
408+
proxyWrite,
369409
),
370410
)
371411
})
@@ -381,29 +421,35 @@ func NewHandler(endpoints Endpoints, upstreamCA []byte, upstreamCert *stdtls.Cer
381421
rh := rulesHandler{client: client, logger: c.logger, tenantLabel: c.tenantLabel}
382422

383423
r.Group(func(r chi.Router) {
424+
r.Use(func(handler http.Handler) http.Handler {
425+
return server.InjectLabelsCtx(
426+
prometheus.Labels{"group": "metricsv1", "handler": "rules"},
427+
handler,
428+
)
429+
})
384430
r.Use(c.uiMiddlewares...)
385431
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
386432
r.Method(http.MethodGet, RulesRawRoute,
387433
otelhttp.WithRouteTag(
388434
c.spanRoutePrefix+RulesRawRoute,
389-
server.InjectLabelsCtx(
390-
prometheus.Labels{"group": "metricsv1", "handler": "rules"},
391-
http.HandlerFunc(rh.get),
392-
),
435+
http.HandlerFunc(rh.get),
393436
),
394437
)
395438
})
396439

397440
r.Group(func(r chi.Router) {
441+
r.Use(func(handler http.Handler) http.Handler {
442+
return server.InjectLabelsCtx(
443+
prometheus.Labels{"group": "metricsv1", "handler": "rules-raw"},
444+
handler,
445+
)
446+
})
398447
r.Use(c.writeMiddlewares...)
399448
r.Use(server.StripTenantPrefix("/api/metrics/v1"))
400449
r.Method(http.MethodPut, RulesRawRoute,
401450
otelhttp.WithRouteTag(
402451
c.spanRoutePrefix+RulesRawRoute,
403-
server.InjectLabelsCtx(
404-
prometheus.Labels{"group": "metricsv1", "handler": "rules-raw"},
405-
http.HandlerFunc(rh.put),
406-
),
452+
http.HandlerFunc(rh.put),
407453
),
408454
)
409455
})
@@ -434,41 +480,50 @@ func NewHandler(endpoints Endpoints, upstreamCA []byte, upstreamCert *stdtls.Cer
434480
}
435481

436482
r.Group(func(r chi.Router) {
483+
r.Use(func(handler http.Handler) http.Handler {
484+
return server.InjectLabelsCtx(
485+
prometheus.Labels{"group": "metricsv1", "handler": "alerts"},
486+
handler,
487+
)
488+
})
437489
r.Use(c.alertmanagerMiddleware.alertsReadMiddlewares...)
438490
r.Use(server.StripTenantPrefixWithSubRoute("/api/metrics/v1", "/am"))
439491

440492
r.Method(http.MethodGet, AlertmanagerAlertsRoute, otelhttp.WithRouteTag(
441493
c.spanRoutePrefix+AlertmanagerAlertsRoute,
442-
server.InjectLabelsCtx(
443-
prometheus.Labels{"group": "metricsv1", "handler": "alerts"},
444-
proxyAlertmanager,
445-
),
494+
proxyAlertmanager,
446495
))
447496
})
448497

449498
r.Group(func(r chi.Router) {
499+
r.Use(func(handler http.Handler) http.Handler {
500+
return server.InjectLabelsCtx(
501+
prometheus.Labels{"group": "metricsv1", "handler": "silences"},
502+
handler,
503+
)
504+
})
450505
r.Use(c.alertmanagerMiddleware.silenceReadMiddlewares...)
451506
r.Use(server.StripTenantPrefixWithSubRoute("/api/metrics/v1", "/am"))
452507

453508
r.Method(http.MethodGet, AlertmanagerSilencesRoute, otelhttp.WithRouteTag(
454509
c.spanRoutePrefix+AlertmanagerSilencesRoute,
455-
server.InjectLabelsCtx(
456-
prometheus.Labels{"group": "metricsv1", "handler": "silences"},
457-
proxyAlertmanager,
458-
),
510+
proxyAlertmanager,
459511
))
460512
})
461513

462514
r.Group(func(r chi.Router) {
515+
r.Use(func(handler http.Handler) http.Handler {
516+
return server.InjectLabelsCtx(
517+
prometheus.Labels{"group": "metricsv1", "handler": "silences"},
518+
handler,
519+
)
520+
})
463521
r.Use(c.alertmanagerMiddleware.silenceWriteMiddlewares...)
464522
r.Use(server.StripTenantPrefixWithSubRoute("/api/metrics/v1", "/am"))
465523

466524
r.Method(http.MethodPost, AlertmanagerSilencesRoute, otelhttp.WithRouteTag(
467525
c.spanRoutePrefix+AlertmanagerSilencesRoute,
468-
server.InjectLabelsCtx(
469-
prometheus.Labels{"group": "metricsv1", "handler": "silences"},
470-
proxyAlertmanager,
471-
),
526+
proxyAlertmanager,
472527
))
473528
})
474529
}

Diff for: ‎ratelimit/http.go

-1
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,6 @@ func (l rateLimiter) Handler(next http.Handler) http.Handler {
127127
httperr.PrometheusAPIError(w, err.Error(), http.StatusInternalServerError)
128128
return
129129
}
130-
131130
next.ServeHTTP(w, r)
132131
})
133132
}

Diff for: ‎ratelimit/http_test.go

+31-4
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,12 @@ import (
1010
"testing"
1111
"time"
1212

13+
"github.com/prometheus/client_golang/prometheus"
14+
1315
"github.com/go-chi/chi"
1416
"github.com/observatorium/api/authentication"
1517
"github.com/observatorium/api/logger"
18+
"github.com/observatorium/api/server"
1619
)
1720

1821
const (
@@ -107,14 +110,20 @@ func TestWithLocalRateLimiter(t *testing.T) {
107110

108111
for _, c := range cases {
109112
t.Run(c.name, func(t *testing.T) {
110-
rlmw := WithLocalRateLimiter(c.configs...)
111-
112113
r := chi.NewMux()
113-
114+
reg := prometheus.NewRegistry()
115+
hardcodedLabels := []string{"group", "handler"}
116+
f := server.NewInstrumentedHandlerFactory(reg, hardcodedLabels)
117+
rlmw := WithLocalRateLimiter(c.configs...)
118+
r.Use(func(handler http.Handler) http.Handler {
119+
return f.NewHandler(nil, handler)
120+
})
114121
r.Group(func(r chi.Router) {
122+
r.Use(func(next http.Handler) http.Handler {
123+
return server.InjectLabelsCtx(prometheus.Labels{"group": "test-group", "handler": "test-handler"}, next)
124+
})
115125
r.Use(authentication.WithTenant)
116126
r.Use(rlmw)
117-
118127
r.HandleFunc(testPathOne+"/{tenant}", func(res http.ResponseWriter, req *http.Request) {
119128
res.WriteHeader(http.StatusOK)
120129
})
@@ -144,6 +153,24 @@ func TestWithLocalRateLimiter(t *testing.T) {
144153
gotTooManyRequests,
145154
)
146155
}
156+
// Check for labels only if rate limit is hit
157+
if pathTest.expectedTooManyRequests > 0 {
158+
metrics, err := reg.Gather()
159+
if err != nil {
160+
t.Fatal(err)
161+
}
162+
for _, metric := range metrics {
163+
for _, m := range metric.GetMetric() {
164+
if m.GetLabel()[1].GetValue() != "test-group" {
165+
t.Fatalf("expected label value to be 'test-group', got %s", m.GetLabel()[1].GetValue())
166+
}
167+
if m.GetLabel()[2].GetValue() != "test-handler" {
168+
t.Fatalf("expected label value to be 'test-handler', got %s", m.GetLabel()[2].GetValue())
169+
}
170+
171+
}
172+
}
173+
}
147174
}
148175
})
149176
}

0 commit comments

Comments
 (0)