Sys level rps

view/json.go

@@ -11,6 +11,8 @@ import (
 	"log"
 	"net/http"
 	"strings"
+
+	"golang.org/x/time/rate"
 )
 
 func JSON(g storage.Graph) error {
@@ -67,7 +69,7 @@ func jsonHandler(g storage.Graph) http.Handler {
 		})
 	}
 
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	return rateLimited(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if gziphttp.Can(r) {
 			gz := gziphttp.New(w)
 			defer gz.Close()
@@ -81,6 +83,18 @@ func jsonHandler(g storage.Graph) http.Handler {
 			Closer: r.Body,
 		}
 		mux.ServeHTTP(w, r)
+	}))
+}
+
+func rateLimited(foo http.HandlerFunc) http.HandlerFunc {
+	sysRPS := config.New().SysRPS
+	limiter := rate.NewLimiter(rate.Limit(sysRPS), sysRPS)
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := limiter.Wait(r.Context()); err != nil {
+			http.Error(w, err.Error(), http.StatusTooManyRequests)
+		}
+		foo(w, r)
 	})
 }
 

view/json_test.go

@@ -1,11 +1,14 @@
 package view
 
 import (
+	"context"
 	"local/dndex/config"
 	"net/http"
+	"net/http/httptest"
 	"net/url"
 	"os"
 	"testing"
+	"time"
 )
 
 func TestGetNamespace(t *testing.T) {
@@ -59,3 +62,31 @@ func TestGetNamespace(t *testing.T) {
 		})
 	}
 }
+
+func TestRateLimited(t *testing.T) {
+	os.Args = os.Args[:1]
+	os.Setenv("SYS_RPS", "10")
+	foo := rateLimited(func(w http.ResponseWriter, r *http.Request) {})
+
+	ok := 0
+	tooMany := 0
+	for i := 0; i < 20; i++ {
+		w := httptest.NewRecorder()
+		r := httptest.NewRequest("GET", "/", nil)
+		ctx, can := context.WithTimeout(r.Context(), time.Millisecond*50)
+		defer can()
+		r = r.WithContext(ctx)
+		foo(w, r)
+		switch w.Code {
+		case http.StatusOK:
+			ok += 1
+		case http.StatusTooManyRequests:
+			tooMany += 1
+		default:
+			t.Fatal("unexpected status", w.Code)
+		}
+	}
+	if ok < 9 || ok > 11 {
+		t.Fatal(ok, tooMany)
+	}
+}