func initWebRateLimiter(app *App) { rateLimitStore := store.NewMemStore(1000) if app.redis != nil { rateLimitStore = store.NewRedisStore(app.redis, "throttle:", 0) } rateLimiter := throttled.RateLimit( app.config.RateLimit, &throttled.VaryBy{Custom: remoteAddrIP}, rateLimitStore, ) rateLimiter.DeniedHandler = &RateLimitExceededAction{App: app, Action: Action{}} app.web.rateLimiter = rateLimiter }
func throttler(perMinRate int) http.Handler { t := throttled.RateLimit(throttled.PerMin(perMinRate), &throttled.VaryBy{RemoteAddr: true}, store.NewMemStore(1000)) return t.Throttle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // Nothing needed here })) }
func main() { flag.Parse() var h http.Handler var ok, ko int var mu sync.Mutex var st throttled.Store // Keep the start time to print since-time start := time.Now() // Create the rate-limit store switch *storeType { case "mem": st = store.NewMemStore(0) case "redis": st = store.NewRedisStore(setupRedis(), "throttled:", 0) default: log.Fatalf("unsupported store: %s", *storeType) } // Create the rate-limit throttler, varying on path t := throttled.RateLimit(throttled.Q{Requests: *requests, Window: *window}, &throttled.VaryBy{ Path: true, }, st) // Set its denied handler t.DeniedHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if *output == "v" || *output == "ko" { log.Printf("KO: %s", time.Since(start)) } throttled.DefaultDeniedHandler.ServeHTTP(w, r) mu.Lock() defer mu.Unlock() ko++ }) // Throttle the OK handler rand.Seed(time.Now().Unix()) h = t.Throttle(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if *output == "v" || *output == "ok" { log.Printf("ok: %s", time.Since(start)) } if *delayRes > 0 { wait := time.Duration(rand.Intn(int(*delayRes))) time.Sleep(wait) } w.WriteHeader(200) mu.Lock() defer mu.Unlock() ok++ })) // Print stats once in a while go func() { for _ = range time.Tick(10 * time.Second) { mu.Lock() log.Printf("ok: %d, ko: %d", ok, ko) mu.Unlock() } }() fmt.Println("server listening on port 9000") http.ListenAndServe(":9000", h) }