mirror of
https://github.com/superseriousbusiness/gotosocial
synced 2025-06-05 21:59:39 +02:00
[performance] simpler throttling logic (#2407)
* reduce complexity of throttling logic to use 1 queue and an atomic int * use atomic add instead of CAS, add throttling test
This commit is contained in:
149
internal/middleware/throttling_test.go
Normal file
149
internal/middleware/throttling_test.go
Normal file
@ -0,0 +1,149 @@
|
||||
// GoToSocial
|
||||
// Copyright (C) GoToSocial Authors admin@gotosocial.org
|
||||
// SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
/*
|
||||
The code in this file is adapted from MIT-licensed code in github.com/go-chi/chi. Thanks chi (thi)!
|
||||
|
||||
See: https://github.com/go-chi/chi/blob/e6baba61759b26ddf7b14d1e02d1da81a4d76c08/middleware/throttle.go
|
||||
|
||||
And: https://github.com/sponsors/pkieltyka
|
||||
*/
|
||||
|
||||
package middleware_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/superseriousbusiness/gotosocial/internal/middleware"
|
||||
)
|
||||
|
||||
func TestThrottlingMiddleware(t *testing.T) {
|
||||
testThrottlingMiddleware(t, 2, time.Second*10)
|
||||
testThrottlingMiddleware(t, 4, time.Second*15)
|
||||
testThrottlingMiddleware(t, 8, time.Second*30)
|
||||
}
|
||||
|
||||
func testThrottlingMiddleware(t *testing.T, cpuMulti int, retryAfter time.Duration) {
|
||||
// Calculate expected request limit + queue.
|
||||
limit := runtime.GOMAXPROCS(0) * cpuMulti
|
||||
queueLimit := limit * cpuMulti
|
||||
|
||||
// Calculate expected retry-after header string.
|
||||
retryAfterStr := strconv.FormatUint(uint64(retryAfter/time.Second), 10)
|
||||
|
||||
// Gin test http engine
|
||||
// (used for ctx init).
|
||||
e := gin.New()
|
||||
|
||||
// Add middleware to the gin engine handler stack.
|
||||
middleware := middleware.Throttle(cpuMulti, retryAfter)
|
||||
e.Use(middleware)
|
||||
|
||||
// Set the blocking gin handler.
|
||||
handler := blockingHandler()
|
||||
e.Handle("GET", "/", handler)
|
||||
|
||||
var cncls []func()
|
||||
|
||||
for i := 0; i < queueLimit+limit; i++ {
|
||||
// Prepare a gin test context.
|
||||
r := httptest.NewRequest("GET", "/", nil)
|
||||
rw := httptest.NewRecorder()
|
||||
|
||||
// Wrap request with new cancel context.
|
||||
ctx, cncl := context.WithCancel(r.Context())
|
||||
r = r.WithContext(ctx)
|
||||
|
||||
// Pass req through
|
||||
// engine handler.
|
||||
go e.ServeHTTP(rw, r)
|
||||
time.Sleep(time.Millisecond)
|
||||
|
||||
// Get http result.
|
||||
res := rw.Result()
|
||||
|
||||
if i < queueLimit {
|
||||
|
||||
// Check status == 200 (default, i.e not set).
|
||||
if res.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status code was set (%d) with queueLimit=%d and request=%d", res.StatusCode, queueLimit, i)
|
||||
}
|
||||
|
||||
// Add cancel to func slice.
|
||||
cncls = append(cncls, cncl)
|
||||
|
||||
} else {
|
||||
|
||||
// Check the returned status code is expected.
|
||||
if res.StatusCode != http.StatusTooManyRequests {
|
||||
t.Fatalf("did not return status 429 (%d) with queueLimit=%d and request=%d", res.StatusCode, queueLimit, i)
|
||||
}
|
||||
|
||||
// Check the returned retry-after header is set.
|
||||
if res.Header.Get("Retry-After") != retryAfterStr {
|
||||
t.Fatalf("did not return retry-after %s with queueLimit=%d and request=%d", retryAfterStr, queueLimit, i)
|
||||
}
|
||||
|
||||
// Cancel on return.
|
||||
defer cncl()
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
// Cancel all blocked reqs.
|
||||
for _, cncl := range cncls {
|
||||
cncl()
|
||||
}
|
||||
time.Sleep(time.Second)
|
||||
|
||||
// Check a bunchh more requests
|
||||
// can now make it through after
|
||||
// previous requests were released!
|
||||
for i := 0; i < limit; i++ {
|
||||
|
||||
// Prepare a gin test context.
|
||||
r := httptest.NewRequest("GET", "/", nil)
|
||||
rw := httptest.NewRecorder()
|
||||
|
||||
// Pass req through
|
||||
// engine handler.
|
||||
go e.ServeHTTP(rw, r)
|
||||
time.Sleep(time.Millisecond)
|
||||
|
||||
// Get http result.
|
||||
res := rw.Result()
|
||||
|
||||
// Check status == 200 (default, i.e not set).
|
||||
if res.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status code was set (%d) with queueLimit=%d and request=%d", res.StatusCode, queueLimit, i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func blockingHandler() gin.HandlerFunc {
|
||||
return func(ctx *gin.Context) {
|
||||
<-ctx.Done()
|
||||
ctx.Status(201) // specifically not 200
|
||||
}
|
||||
}
|
Reference in New Issue
Block a user