Skip to content

Commit

Permalink
Use custom metrics middleware for Palomar to avoid 404DOSing
Browse files Browse the repository at this point in the history
  • Loading branch information
ericvolp12 committed Oct 1, 2023
1 parent b6da8c8 commit 98242a4
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 3 deletions.
93 changes: 93 additions & 0 deletions search/metrics.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,12 @@
package search

import (
"errors"
"net/http"
"strconv"
"time"

"github.com/labstack/echo/v4"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promauto"
)
Expand Down Expand Up @@ -49,3 +55,90 @@ var currentSeq = promauto.NewGauge(prometheus.GaugeOpts{
Name: "search_current_seq",
Help: "Current sequence number",
})

var reqSz = promauto.NewHistogramVec(prometheus.HistogramOpts{
Name: "http_request_size_bytes",
Help: "A histogram of request sizes for requests.",
Buckets: prometheus.ExponentialBuckets(100, 10, 8),
}, []string{"code", "method", "path"})

var reqDur = promauto.NewHistogramVec(prometheus.HistogramOpts{
Name: "http_request_duration_seconds",
Help: "A histogram of latencies for requests.",
Buckets: prometheus.ExponentialBuckets(0.0001, 2, 18),
}, []string{"code", "method", "path"})

var reqCnt = promauto.NewCounterVec(prometheus.CounterOpts{
Name: "http_requests_total",
Help: "A counter for requests to the wrapped handler.",
}, []string{"code", "method", "path"})

var resSz = promauto.NewHistogramVec(prometheus.HistogramOpts{
Name: "http_response_size_bytes",
Help: "A histogram of response sizes for requests.",
Buckets: prometheus.ExponentialBuckets(100, 10, 8),
}, []string{"code", "method", "path"})

// MetricsMiddleware defines handler function for metrics middleware
func MetricsMiddleware(next echo.HandlerFunc) echo.HandlerFunc {
return func(c echo.Context) error {
path := c.Path()
if path == "/metrics" || path == "/_health" {
return next(c)
}

start := time.Now()
requestSize := computeApproximateRequestSize(c.Request())

err := next(c)

status := c.Response().Status
if err != nil {
var httpError *echo.HTTPError
if errors.As(err, &httpError) {
status = httpError.Code
}
if status == 0 || status == http.StatusOK {
status = http.StatusInternalServerError
}
}

elapsed := float64(time.Since(start)) / float64(time.Second)

statusStr := strconv.Itoa(status)
method := c.Request().Method

responseSize := float64(c.Response().Size)

reqDur.WithLabelValues(statusStr, method, path).Observe(elapsed)
reqCnt.WithLabelValues(statusStr, method, path).Inc()
reqSz.WithLabelValues(statusStr, method, path).Observe(float64(requestSize))
resSz.WithLabelValues(statusStr, method, path).Observe(responseSize)

return err
}
}

func computeApproximateRequestSize(r *http.Request) int {
s := 0
if r.URL != nil {
s = len(r.URL.Path)
}

s += len(r.Method)
s += len(r.Proto)
for name, values := range r.Header {
s += len(name)
for _, value := range values {
s += len(value)
}
}
s += len(r.Host)

// N.B. r.Form and r.MultipartForm are assumed to be included in r.URL.

if r.ContentLength != -1 {
s += int(r.ContentLength)
}
return s
}
6 changes: 3 additions & 3 deletions search/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@ import (
"github.com/bluesky-social/indigo/backfill"
"github.com/bluesky-social/indigo/util/version"
"github.com/bluesky-social/indigo/xrpc"
"github.com/prometheus/client_golang/prometheus/promhttp"

"github.com/labstack/echo-contrib/echoprometheus"
"github.com/labstack/echo/v4"
"github.com/labstack/echo/v4/middleware"
es "github.com/opensearch-project/opensearch-go/v2"
Expand Down Expand Up @@ -183,7 +183,7 @@ func (s *Server) RunAPI(listen string) error {
e.HideBanner = true
e.Use(slogecho.New(s.logger))
e.Use(middleware.Recover())
e.Use(echoprometheus.NewMiddleware("palomar"))
e.Use(MetricsMiddleware)
e.Use(middleware.BodyLimit("64M"))

e.HTTPErrorHandler = func(err error, ctx echo.Context) {
Expand All @@ -197,7 +197,7 @@ func (s *Server) RunAPI(listen string) error {

e.Use(middleware.CORS())
e.GET("/_health", s.handleHealthCheck)
e.GET("/metrics", echoprometheus.NewHandler())
e.GET("/metrics", echo.WrapHandler(promhttp.Handler()))
e.GET("/xrpc/app.bsky.unspecced.searchPostsSkeleton", s.handleSearchPostsSkeleton)
e.GET("/xrpc/app.bsky.unspecced.searchActorsSkeleton", s.handleSearchActorsSkeleton)
e.GET("/xrpc/app.bsky.unspecced.indexRepos", s.handleIndexRepos)
Expand Down

0 comments on commit 98242a4

Please sign in to comment.