snapshot of https://github.com/go-chi/chi.git (master, 2026-09-15; blobs > 131072B omitted)

middleware/throttle_test.go

@@ -0,0 +1,329 @@
+package middleware
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/go-chi/chi/v5"
+)
+
+var testContent = []byte("Hello world!")
+
+func TestThrottleBacklog(t *testing.T) {
+	r := chi.NewRouter()
+
+	r.Use(ThrottleBacklog(10, 50, time.Second*10))
+
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		time.Sleep(time.Second * 1) // Expensive operation.
+		w.Write(testContent)
+	})
+
+	server := httptest.NewServer(r)
+	defer server.Close()
+
+	client := http.Client{
+		Timeout: time.Second * 5, // Maximum waiting time.
+	}
+
+	var wg sync.WaitGroup
+
+	// The throttler processes 10 consecutive requests, each one of those
+	// requests lasts 1s. The maximum number of requests this can possible serve
+	// before the clients time out (5s) is 40.
+	for i := range 40 {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+
+			res, err := client.Get(server.URL)
+			assertNoError(t, err)
+
+			assertEqual(t, http.StatusOK, res.StatusCode)
+			buf, err := io.ReadAll(res.Body)
+			assertNoError(t, err)
+			assertEqual(t, testContent, buf)
+		}(i)
+	}
+
+	wg.Wait()
+}
+
+func TestThrottleClientTimeout(t *testing.T) {
+	r := chi.NewRouter()
+
+	r.Use(ThrottleBacklog(10, 50, time.Second*10))
+
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		time.Sleep(time.Second * 5) // Expensive operation.
+		w.Write(testContent)
+	})
+
+	server := httptest.NewServer(r)
+	defer server.Close()
+
+	client := http.Client{
+		Timeout: time.Second * 3, // Maximum waiting time.
+	}
+
+	var wg sync.WaitGroup
+
+	for i := range 10 {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			_, err := client.Get(server.URL)
+			assertError(t, err)
+		}(i)
+	}
+
+	wg.Wait()
+}
+
+func TestThrottleTriggerGatewayTimeout(t *testing.T) {
+	r := chi.NewRouter()
+
+	r.Use(ThrottleBacklog(50, 100, time.Second*5))
+
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		time.Sleep(time.Second * 10) // Expensive operation.
+		w.Write(testContent)
+	})
+
+	server := httptest.NewServer(r)
+	defer server.Close()
+
+	client := http.Client{
+		Timeout: time.Second * 60, // Maximum waiting time.
+	}
+
+	var wg sync.WaitGroup
+
+	// These requests will be processed normally until they finish.
+	for i := range 50 {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+
+			res, err := client.Get(server.URL)
+			assertNoError(t, err)
+			assertEqual(t, http.StatusOK, res.StatusCode)
+		}(i)
+	}
+
+	time.Sleep(time.Second * 1)
+
+	// These requests will wait for the first batch to complete but it will take
+	// too much time, so they will eventually receive a timeout error.
+	for i := range 50 {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+
+			res, err := client.Get(server.URL)
+			assertNoError(t, err)
+
+			buf, err := io.ReadAll(res.Body)
+			assertNoError(t, err)
+			assertEqual(t, http.StatusTooManyRequests, res.StatusCode)
+			assertEqual(t, errTimedOut, strings.TrimSpace(string(buf)))
+		}(i)
+	}
+
+	wg.Wait()
+}
+
+func TestThrottleMaximum(t *testing.T) {
+	r := chi.NewRouter()
+
+	r.Use(ThrottleBacklog(10, 10, time.Second*5))
+
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		time.Sleep(time.Second * 3) // Expensive operation.
+		w.Write(testContent)
+	})
+
+	server := httptest.NewServer(r)
+	defer server.Close()
+
+	client := http.Client{
+		Timeout: time.Second * 60, // Maximum waiting time.
+	}
+
+	var wg sync.WaitGroup
+
+	for i := range 20 {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+
+			res, err := client.Get(server.URL)
+			assertNoError(t, err)
+			assertEqual(t, http.StatusOK, res.StatusCode)
+
+			buf, err := io.ReadAll(res.Body)
+			assertNoError(t, err)
+			assertEqual(t, testContent, buf)
+		}(i)
+	}
+
+	// Wait less time than what the server takes to reply.
+	time.Sleep(time.Second * 2)
+
+	// At this point the server is still processing, all the following request
+	// will be beyond the server capacity.
+	for i := range 20 {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+
+			res, err := client.Get(server.URL)
+			assertNoError(t, err)
+
+			buf, err := io.ReadAll(res.Body)
+			assertNoError(t, err)
+			assertEqual(t, http.StatusTooManyRequests, res.StatusCode)
+			assertEqual(t, errCapacityExceeded, strings.TrimSpace(string(buf)))
+		}(i)
+	}
+
+	wg.Wait()
+}
+
+func TestThrottleRetryAfter(t *testing.T) {
+	r := chi.NewRouter()
+	retryAfterFn := func(ctxDone bool) time.Duration { return time.Hour }
+
+	r.Use(ThrottleWithOpts(ThrottleOpts{
+		Limit:        5,
+		BacklogLimit: 0,
+		RetryAfterFn: retryAfterFn,
+	}))
+
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(time.Second * 1) // Expensive operation.
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("ok"))
+	})
+
+	server := httptest.NewServer(r)
+	defer server.Close()
+	client := http.Client{}
+
+	type result struct {
+		status int
+		header http.Header
+	}
+
+	var wg sync.WaitGroup
+	totalRequests := 10
+	resultsCh := make(chan result, totalRequests)
+
+	for i := 0; i < totalRequests; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			res, _ := client.Get(server.URL)
+			resultsCh <- result{status: res.StatusCode, header: res.Header}
+		}()
+	}
+
+	wg.Wait()
+	close(resultsCh)
+
+	count200 := 0
+	count429 := 0
+	for res := range resultsCh {
+		switch res.status {
+		case http.StatusOK:
+			count200++
+			continue
+		case http.StatusTooManyRequests:
+			count429++
+			assertEqual(t, "3600", res.header.Get("Retry-After"))
+			continue
+		default:
+			t.Fatalf("Unexpected status code: %d", res.status)
+			continue
+		}
+	}
+
+	assertEqual(t, 5, count200)
+	assertEqual(t, 5, count429)
+}
+
+func TestThrottleCustomStatusCode(t *testing.T) {
+	const timeout = time.Second * 3
+
+	wait := make(chan struct{})
+
+	r := chi.NewRouter()
+	r.Use(ThrottleWithOpts(ThrottleOpts{Limit: 1, StatusCode: http.StatusServiceUnavailable}))
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
+		select {
+		case <-wait:
+		case <-time.After(timeout):
+		}
+		w.WriteHeader(http.StatusOK)
+	})
+	server := httptest.NewServer(r)
+	defer server.Close()
+
+	const totalRequestCount = 5
+
+	codes := make(chan int, totalRequestCount)
+	errs := make(chan error, totalRequestCount)
+	client := &http.Client{Timeout: timeout}
+	for range totalRequestCount {
+		go func() {
+			resp, err := client.Get(server.URL)
+			if err != nil {
+				errs <- err
+				return
+			}
+			codes <- resp.StatusCode
+		}()
+	}
+
+	waitResponse := func(wantCode int) {
+		select {
+		case err := <-errs:
+			t.Fatal(err)
+		case code := <-codes:
+			assertEqual(t, wantCode, code)
+		case <-time.After(timeout):
+			t.Fatalf("waiting %d code, timeout exceeded", wantCode)
+		}
+	}
+
+	for range totalRequestCount - 1 {
+		waitResponse(http.StatusServiceUnavailable)
+	}
+	close(wait) // Allow the last request to proceed.
+	waitResponse(http.StatusOK)
+}
+
+func BenchmarkThrottle(b *testing.B) {
+	throttleMiddleware := ThrottleBacklog(1000, 50, time.Second)
+
+	handler := throttleMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	req := httptest.NewRequest("GET", "/", nil)
+
+	b.ReportAllocs()
+	for b.Loop() {
+		w := httptest.NewRecorder()
+		handler.ServeHTTP(w, req)
+	}
+}