func StartServer() { l4g.Info("Starting Server...") l4g.Info("Server is listening on " + utils.Cfg.ServiceSettings.ListenAddress) var handler http.Handler = Srv.Router if utils.Cfg.RateLimitSettings.EnableRateLimiter { l4g.Info("RateLimiter is enabled") vary := throttled.VaryBy{} if utils.Cfg.RateLimitSettings.VaryByRemoteAddr { vary.RemoteAddr = true } if len(utils.Cfg.RateLimitSettings.VaryByHeader) > 0 { vary.Headers = strings.Fields(utils.Cfg.RateLimitSettings.VaryByHeader) if utils.Cfg.RateLimitSettings.VaryByRemoteAddr { l4g.Warn("RateLimitSettings not configured properly using VaryByHeader and disabling VaryByRemoteAddr") vary.RemoteAddr = false } } th := throttled.RateLimit(throttled.PerSec(utils.Cfg.RateLimitSettings.PerSec), &vary, throttledStore.NewMemStore(utils.Cfg.RateLimitSettings.MemoryStoreSize)) th.DeniedHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { l4g.Error("%v: code=429 ip=%v", r.URL.Path, GetIpAddress(r)) throttled.DefaultDeniedHandler.ServeHTTP(w, r) }) handler = th.Throttle(Srv.Router) } go func() { err := Srv.Server.ListenAndServe(utils.Cfg.ServiceSettings.ListenAddress, handler) if err != nil { l4g.Critical("Error starting server, err:%v", err) time.Sleep(time.Second) panic("Error starting server " + err.Error()) } }() }
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) }