Golang memstore.New函数代码示例

func TestMemStoreLRU(t *testing.T) {
	st, err := memstore.New(10)
	if err != nil {
		t.Fatal(err)
	}
	storetest.TestGCRAStore(t, st)
}

func BenchmarkMemStoreUnlimited(b *testing.B) {
	st, err := memstore.New(0)
	if err != nil {
		b.Fatal(err)
	}
	storetest.BenchmarkGCRAStore(b, st)
}

// Demonstrates direct use of GCRARateLimiter's RateLimit function (and the
// more general RateLimiter interface). This should be used anywhere where
// granular control over rate limiting is required.
func ExampleGCRARateLimiter() {
	store, err := memstore.New(65536)
	if err != nil {
		log.Fatal(err)
	}

	// Maximum burst of 5 which refills at 1 token per hour.
	quota := throttled.RateQuota{throttled.PerHour(1), 5}

	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
	if err != nil {
		log.Fatal(err)
	}

	// Bucket according to the number i / 10 (so 1 falls into the bucket 0
	// while 11 falls into the bucket 1). This has the effect of allowing a
	// burst of 5 plus 1 (a single emission interval) on every ten iterations
	// of the loop. See the output for better clarity here.
	//
	// We also refill the bucket at 1 token per hour, but that has no effect
	// for the purposes of this example.
	for i := 0; i < 20; i++ {
		bucket := fmt.Sprintf("by-order:%v", i/10)

		limited, result, err := rateLimiter.RateLimit(bucket, 1)
		if err != nil {
			log.Fatal(err)
		}

		if limited {
			fmt.Printf("Iteration %2v; bucket %v: FAILED. Rate limit exceeded.\n",
				i, bucket)
		} else {
			fmt.Printf("Iteration %2v; bucket %v: Operation successful (remaining=%v).\n",
				i, bucket, result.Remaining)
		}
	}

	// Output:
	// Iteration  0; bucket by-order:0: Operation successful (remaining=5).
	// Iteration  1; bucket by-order:0: Operation successful (remaining=4).
	// Iteration  2; bucket by-order:0: Operation successful (remaining=3).
	// Iteration  3; bucket by-order:0: Operation successful (remaining=2).
	// Iteration  4; bucket by-order:0: Operation successful (remaining=1).
	// Iteration  5; bucket by-order:0: Operation successful (remaining=0).
	// Iteration  6; bucket by-order:0: FAILED. Rate limit exceeded.
	// Iteration  7; bucket by-order:0: FAILED. Rate limit exceeded.
	// Iteration  8; bucket by-order:0: FAILED. Rate limit exceeded.
	// Iteration  9; bucket by-order:0: FAILED. Rate limit exceeded.
	// Iteration 10; bucket by-order:1: Operation successful (remaining=5).
	// Iteration 11; bucket by-order:1: Operation successful (remaining=4).
	// Iteration 12; bucket by-order:1: Operation successful (remaining=3).
	// Iteration 13; bucket by-order:1: Operation successful (remaining=2).
	// Iteration 14; bucket by-order:1: Operation successful (remaining=1).
	// Iteration 15; bucket by-order:1: Operation successful (remaining=0).
	// Iteration 16; bucket by-order:1: FAILED. Rate limit exceeded.
	// Iteration 17; bucket by-order:1: FAILED. Rate limit exceeded.
	// Iteration 18; bucket by-order:1: FAILED. Rate limit exceeded.
	// Iteration 19; bucket by-order:1: FAILED. Rate limit exceeded.
}

func throttledHandler(h http.HandlerFunc) http.Handler {
	// for now just use an inmemory storage, so per server. In the future we
	// can change to store the state on a remote DB if we want to distribute
	// the counts
	store, err := memstore.New(65536)
	if err != nil {
		// panics only if memstore.New() receives an integer number, so this is
		// OK, this means it's a human error and needs to be fixed
		log.Fatal(err)
	}

	// Based on datadog metrics, kloud.info is called on average 200
	// req/minute.
	quota := throttled.RateQuota{
		MaxRate:  throttled.PerMin(200),
		MaxBurst: 300,
	}

	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
	if err != nil {
		// we exit because this is code error and must be handled
		log.Fatalln(err)
	}

	httpRateLimiter := throttled.HTTPRateLimiter{
		RateLimiter: rateLimiter,
	}

	return httpRateLimiter.RateLimit(http.HandlerFunc(h))
}

// DEPRECATED. NewMemStore is a compatible alias for mem.New
func NewMemStore(maxKeys int) *memstore.MemStore {
	st, err := memstore.New(maxKeys)
	if err != nil {
		// As of this writing, `lru.New` can only return an error if you pass
		// maxKeys <= 0 so this should never occur.
		panic(err)
	}
	return st
}

func TestRateLimitUpdateFailures(t *testing.T) {
	rq := throttled.RateQuota{throttled.PerSec(1), 1}
	mst, err := memstore.New(0)
	if err != nil {
		t.Fatal(err)
	}
	st := testStore{store: mst, failUpdates: true}
	rl, err := throttled.NewGCRARateLimiter(&st, rq)
	if err != nil {
		t.Fatal(err)
	}

	if _, _, err := rl.RateLimit("foo", 1); err == nil {
		t.Error("Expected limiting to fail when store updates fail")
	}
}

func buildLimiter(quota throttled.RateQuota) (*throttled.HTTPRateLimiter, error) {
	store, err := memstore.New(65536)
	if err != nil {
		return nil, err
	}

	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
	if err != nil {
		return nil, err
	}

	return &throttled.HTTPRateLimiter{
		RateLimiter: rateLimiter,
		VaryBy:      &throttled.VaryBy{RemoteAddr: true},
	}, nil
}

func StartServer() {
	l4g.Info(utils.T("api.server.start_server.starting.info"))
	l4g.Info(utils.T("api.server.start_server.listening.info"), utils.Cfg.ServiceSettings.ListenAddress)

	var handler http.Handler = &CorsWrapper{Srv.Router}

	if *utils.Cfg.RateLimitSettings.Enable {
		l4g.Info(utils.T("api.server.start_server.rate.info"))

		store, err := memstore.New(utils.Cfg.RateLimitSettings.MemoryStoreSize)
		if err != nil {
			l4g.Critical(utils.T("api.server.start_server.rate_limiting_memory_store"))
			return
		}

		quota := throttled.RateQuota{
			MaxRate:  throttled.PerSec(utils.Cfg.RateLimitSettings.PerSec),
			MaxBurst: *utils.Cfg.RateLimitSettings.MaxBurst,
		}

		rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
		if err != nil {
			l4g.Critical(utils.T("api.server.start_server.rate_limiting_rate_limiter"))
			return
		}

		httpRateLimiter := throttled.HTTPRateLimiter{
			RateLimiter: rateLimiter,
			VaryBy:      &VaryBy{},
			DeniedHandler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				l4g.Error("%v: Denied due to throttling settings code=429 ip=%v", r.URL.Path, GetIpAddress(r))
				throttled.DefaultDeniedHandler.ServeHTTP(w, r)
			}),
		}

		handler = httpRateLimiter.RateLimit(handler)
	}

	go func() {
		err := manners.ListenAndServe(utils.Cfg.ServiceSettings.ListenAddress, handlers.RecoveryHandler(handlers.PrintRecoveryStack(true))(handler))
		if err != nil {
			l4g.Critical(utils.T("api.server.start_server.starting.critical"), err)
			time.Sleep(time.Second)
		}
	}()
}

// WithRateLimit wraps an ctxhttp.Handler to limit incoming requests.
// Requests that are not limited will be passed to the handler
// unchanged.  Limited requests will be passed to the DeniedHandler.
// X-RateLimit-Limit, X-RateLimit-Remaining, X-RateLimit-Reset and
// Retry-After headers will be written to the response based on the
// values in the RateLimitResult.
func (t *HTTPRateLimit) WithRateLimit(rlStore throttled.GCRAStore, h ctxhttp.Handler) ctxhttp.Handler {
	if t.Config == nil {
		t.Config = config.DefaultManager
	}
	if t.DeniedHandler == nil {
		t.DeniedHandler = DefaultDeniedHandler
	}

	if t.RateLimiter == nil {
		if rlStore == nil {
			var err error
			rlStore, err = memstore.New(65536)
			if err != nil {
				panic(err)
			}
		}

		var err error
		t.RateLimiter, err = throttled.NewGCRARateLimiter(rlStore, t.quota())
		if err != nil {
			panic(err)
		}
	}

	return ctxhttp.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {

		var k string
		if t.VaryBy != nil {
			k = t.VaryBy.Key(r)
		}

		limited, context, err := t.RateLimiter.RateLimit(k, 1)

		if err != nil {
			return err
		}

		setRateLimitHeaders(w, context)

		if !limited {
			return h.ServeHTTPContext(ctx, w, r)
		}

		return t.DeniedHandler.ServeHTTPContext(ctx, w, r)
	})
}

// ExampleHTTPRateLimiter demonstrates the usage of HTTPRateLimiter
// for rate-limiting access to an http.Handler to 20 requests per path
// per minute with a maximum burst of 5 requests.
func ExampleHTTPRateLimiter() {
	store, err := memstore.New(65536)
	if err != nil {
		log.Fatal(err)
	}

	quota := throttled.RateQuota{throttled.PerMin(20), 5}
	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
	if err != nil {
		log.Fatal(err)
	}

	httpRateLimiter := throttled.HTTPRateLimiter{
		RateLimiter: rateLimiter,
		VaryBy:      &throttled.VaryBy{Path: true},
	}

	http.ListenAndServe(":8080", httpRateLimiter.RateLimit(myHandler))
}

func throttle(next http.Handler, o ServerOptions) http.Handler {
	store, err := memstore.New(65536)
	if err != nil {
		return throttleError(err)
	}

	quota := throttled.RateQuota{throttled.PerSec(o.Concurrency), o.Burst}
	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
	if err != nil {
		return throttleError(err)
	}

	httpRateLimiter := throttled.HTTPRateLimiter{
		RateLimiter: rateLimiter,
		VaryBy:      &throttled.VaryBy{Method: true},
	}

	return httpRateLimiter.RateLimit(next)
}

// Run is the main function
func Run(cmd *cobra.Command, args []string) {
	database.InitDB()
	address := viper.GetString("address")

	// Throttling control
	store, err := memstore.New(65536)
	if err != nil {
		logrus.Fatal(err)
	}

	quota := throttled.RateQuota{throttled.PerMin(20), 5}
	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
	if err != nil {
		logrus.Fatal(err)
	}

	httpRateLimiter := throttled.HTTPRateLimiter{
		RateLimiter: rateLimiter,
		VaryBy:      &throttled.VaryBy{Path: true},
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/authenticate/free", handleFreeCookie)
	mux.HandleFunc("/authenticate/verify", handleAuthenticate)
	mux.Handle("/", http.StripPrefix("/authenticate/", http.FileServer(http.Dir("./static"))))

	database.InitDB()

	if buildst == "none" {
		buildst = "[This Dev version is not compiled using regular procedure]"
	}

	logrus.Info("Starting App, ", Version)
	logrus.Infof("2FA HTTP layer listening on %s", address)
	logrus.Infof("Domain for cookies is %s", viper.GetString("domain"))
	logrus.Infof("Cookie max age is %d hour(s)", viper.GetInt("cookiemaxage"))
	logrus.Info("Starting instance on ", time.Now())

	if err := http.ListenAndServe(address, httpRateLimiter.RateLimit(mux)); err != nil {
		logrus.Fatal("Unable to create HTTP layer", err)
	}
}

func NewWallClockLimiter(burst, rate time.Duration, clock Clock) (Limiter, error) {
	var l wallClockLimiter

	l.estimates = make(map[string]time.Duration)

	store, err := memstore.New(65536)
	if err != nil {
		return nil, err
	}

	quota := throttled.RateQuota{
		throttled.PerSec(int(rate / time.Millisecond)),
		int(burst / time.Millisecond),
	}
	l.limiter, err = throttled.NewGCRARateLimiter(&clockedMemStore{store, clock}, quota)
	if err != nil {
		return nil, err
	}

	return &l, nil
}

func TestRateLimit(t *testing.T) {
	limit := 5
	rq := throttled.RateQuota{throttled.PerSec(1), limit - 1}
	start := time.Unix(0, 0)
	cases := []struct {
		now               time.Time
		volume, remaining int
		reset, retry      time.Duration
		limited           bool
	}{
		// You can never make a request larger than the maximum
		0: {start, 6, 5, 0, -1, true},
		// Rate limit normal requests appropriately
		1:  {start, 1, 4, time.Second, -1, false},
		2:  {start, 1, 3, 2 * time.Second, -1, false},
		3:  {start, 1, 2, 3 * time.Second, -1, false},
		4:  {start, 1, 1, 4 * time.Second, -1, false},
		5:  {start, 1, 0, 5 * time.Second, -1, false},
		6:  {start, 1, 0, 5 * time.Second, time.Second, true},
		7:  {start.Add(3000 * time.Millisecond), 1, 2, 3000 * time.Millisecond, -1, false},
		8:  {start.Add(3100 * time.Millisecond), 1, 1, 3900 * time.Millisecond, -1, false},
		9:  {start.Add(4000 * time.Millisecond), 1, 1, 4000 * time.Millisecond, -1, false},
		10: {start.Add(8000 * time.Millisecond), 1, 4, 1000 * time.Millisecond, -1, false},
		11: {start.Add(9500 * time.Millisecond), 1, 4, 1000 * time.Millisecond, -1, false},
		// Zero-volume request just peeks at the state
		12: {start.Add(9500 * time.Millisecond), 0, 4, time.Second, -1, false},
		// High-volume request uses up more of the limit
		13: {start.Add(9500 * time.Millisecond), 2, 2, 3 * time.Second, -1, false},
		// Large requests cannot exceed limits
		14: {start.Add(9500 * time.Millisecond), 5, 2, 3 * time.Second, 3 * time.Second, true},
	}

	mst, err := memstore.New(0)
	if err != nil {
		t.Fatal(err)
	}
	st := testStore{store: mst}

	rl, err := throttled.NewGCRARateLimiter(&st, rq)
	if err != nil {
		t.Fatal(err)
	}

	// Start the server
	for i, c := range cases {
		st.clock = c.now

		limited, context, err := rl.RateLimit("foo", c.volume)
		if err != nil {
			t.Fatalf("%d: %#v", i, err)
		}

		if limited != c.limited {
			t.Errorf("%d: expected Limited to be %t but got %t", i, c.limited, limited)
		}

		if have, want := context.Limit, limit; have != want {
			t.Errorf("%d: expected Limit to be %d but got %d", i, want, have)
		}

		if have, want := context.Remaining, c.remaining; have != want {
			t.Errorf("%d: expected Remaining to be %d but got %d", i, want, have)
		}

		if have, want := context.ResetAfter, c.reset; have != want {
			t.Errorf("%d: expected ResetAfter to be %s but got %s", i, want, have)
		}

		if have, want := context.RetryAfter, c.retry; have != want {
			t.Errorf("%d: expected RetryAfter to be %d but got %d", i, want, have)
		}
	}
}

func main() {

	maxProcs := runtime.NumCPU()
	runtime.GOMAXPROCS(maxProcs)

	var (
		listen_host        string
		fs_basepath        string
		fs_cachepath       string
		http_origin        string
		error_imgpath      string
		secret             string
		rate_limit_by_path int
		rate_limit_bursts  int
		concurrency        int
		newrelic_key       string
	)

	flag.StringVar(&listen_host, "listen", "localhost:3000", "HTTP host:port to listen for incoming requests")
	flag.StringVar(&fs_basepath, "basepath", "", "File system base path to lookup images")
	flag.StringVar(&fs_cachepath, "cachepath", "", "File system base path to cache images")
	flag.StringVar(&http_origin, "httporigin", "", "HTTP endpoint to lookup origin images")
	flag.StringVar(&error_imgpath, "errorimgpath", "./fixtures/error.png", "Error image path in local filesystem")
	flag.StringVar(&secret, "secret", "", "Secret used to check image signature")
	flag.StringVar(&newrelic_key, "newrelic", "", "NewRelic application key")
	flag.IntVar(&rate_limit_by_path, "ratelimit", 0, "Rate limit. Concurrent images per request path")
	flag.IntVar(&rate_limit_bursts, "bursts", 5, "Rate limits bursts")
	flag.IntVar(&concurrency, "concurrency", 0, "Max concurrent requests")

	flag.Parse()

	var store btcdn.Store
	var storeErr error
	var newrelic *gorelic.Agent

	if fs_basepath != "" {
		store, storeErr = btcdn.NewFsStore(fs_basepath)
	} else {
		store, storeErr = btcdn.NewHttpStore(http_origin)
	}

	if storeErr != nil {
		log.Fatal(storeErr)
	}

	pr, err := processor.New()
	if err != nil {
		log.Fatal(err)
	}

	imgSrv, err := btcdn.NewImageServer(store, pr, error_imgpath)
	if err != nil {
		log.Fatal(err)
	}

	router := mux.NewRouter()

	var srv http.Handler
	srv = router

	if secret == "" { // no security
		log.Println("No secret. Using ImageServer with no DoS protection")
		router.HandleFunc("/", imgSrv.ServeHTTP)
	} else { // secure server
		log.Println("Secret provided. Using ImageServer with DoS protection")
		srv = btcdn.NewSecureServer(imgSrv, secret, router)
	}

	// Only cache named sizes
	// to the filesystem for now
	if fs_cachepath != "" {
		log.Println("Caching named sizes to", fs_cachepath)
		cachedEndpoint := btcdn.NewCachedEndpoint(imgSrv, fs_cachepath)
		srv = btcdn.NewNamedSizesServer(cachedEndpoint, router)
	} else {
		srv = btcdn.NewNamedSizesServer(imgSrv, router)
	}

	// Setup rate-limited server
	if rate_limit_by_path > 0 {
		store, err := tstore.New(65536)
		if err != nil {
			log.Fatal(err)
		}
		quota := throttled.RateQuota{throttled.PerSec(rate_limit_by_path), rate_limit_bursts}
		rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
		if err != nil {
			log.Fatal(err)
		}

		httpRateLimiter := throttled.HTTPRateLimiter{
			RateLimiter: rateLimiter,
			VaryBy:      &throttled.VaryBy{Path: true},
		}

		log.Println("Rate limiting by path with ", rate_limit_by_path, "reqs. per second with bursts of ", rate_limit_bursts)
		srv = httpRateLimiter.RateLimit(srv)
	}

	// http.HandleFunc("/favicon.ico", FaviconHandler)
//.........这里部分代码省略.........