Golang ratelimit.Bucket類代碼示例

func publish(bucket *ratelimit.Bucket) {
	conn, err := amqp.Dial("amqp://test:[email protected]:5672/")
	failOnError(err, "Failed to connect to RabbitMQ")
	defer conn.Close()
	ch, err := conn.Channel()
	failOnError(err, "Failed to open a channel")
	defer ch.Close()
	q, err := ch.QueueDeclare(
		"hello", // name
		false,   // durable
		false,   // delete when unused
		false,   // exclusive
		false,   // no-wait
		nil,     // arguments
	)
	failOnError(err, "Failed to declare a queue")
	for {
		bucket.Wait(1)
		body := "hello"
		err = ch.Publish(
			"",     // exchange
			q.Name, // routing key
			false,  // mandatory
			false,  // immediate
			amqp.Publishing{
				ContentType: "text/plain",
				Body:        []byte(body),
			})
		failOnError(err, "Failed to publish a message")
	}

}

// NewTokenBucketThrottler returns an endpoint.Middleware that acts as a
// request throttler based on a token-bucket algorithm. Requests that would
// exceed the maximum request rate are delayed via the parameterized sleep
// function. By default you may pass time.Sleep.
func NewTokenBucketThrottler(tb *ratelimit.Bucket, sleep func(time.Duration)) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func() error {
			sleep(tb.Take(1))
			return next()
		}
	}
}

// NewTokenBucketThrottler returns an endpoint.Middleware that acts as a
// request throttler based on a token-bucket algorithm. Requests that would
// exceed the maximum request rate are delayed via the parameterized sleep
// function. By default you may pass time.Sleep.
func NewTokenBucketThrottler(tb *ratelimit.Bucket, sleep func(time.Duration)) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (interface{}, error) {
			sleep(tb.Take(1))
			return next(ctx, request)
		}
	}
}

// NewTokenBucketLimiter returns an endpoint.Middleware that acts as a rate
// limiter based on a token-bucket algorithm. Requests that would exceed the
// maximum request rate are simply rejected with an error.
func NewTokenBucketLimiter(tb *ratelimit.Bucket) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func() error {
			if tb.TakeAvailable(1) == 0 {
				return ErrLimited
			}
			return next()
		}
	}
}

func limit(b *ratelimit.Bucket, wait bool, errId string) func() error {
	return func() error {
		if wait {
			time.Sleep(b.Take(1))
		} else if b.TakeAvailable(1) == 0 {
			return errors.New(errId, "too many request", 429)
		}
		return nil
	}
}

func NewTokenBucketLimiter(tb *ratelimit.Bucket) endpoint.Middleware {
	return func(next endpoint.Endpoint) endpoint.Endpoint {
		return func(ctx context.Context, request interface{}) (interface{}, error) {
			if tb.TakeAvailable(1) == 0 { // HL
				return nil, ErrLimited // HL
			} // HL
			return next(ctx, request) // HL
		}
	}
}

func getTestServerThrottled(body string) *httptest.Server {
	var rateLimiter *ratelimit.Bucket
	// Rate limit: 2 req/s, capacity 2
	rateLimiter = ratelimit.NewBucket(500*time.Millisecond, 2)
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		code := 200
		if tokens := rateLimiter.TakeAvailable(1); tokens == 0 {
			code = 503
		}

		w.WriteHeader(code)
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, body)
	}))
}

func makeRateLimitFunc(sessionRateLimit, globalRateLimit *ratelimit.Bucket) func(int64) {
	// This may be a case of super duper premature optimization... We build an
	// optimized function to do the rate limiting here based on what we need
	// to do and then use it in the loop.

	if sessionRateLimit == nil && globalRateLimit == nil {
		// No limiting needed. We could equally well return a func(int64){} and
		// not do a nil check were we use it, but I think the nil check there
		// makes it clear that there will be no limiting if none is
		// configured...
		return nil
	}

	if sessionRateLimit == nil {
		// We only have a global limiter
		return func(bytes int64) {
			globalRateLimit.Wait(bytes)
		}
	}

	if globalRateLimit == nil {
		// We only have a session limiter
		return func(bytes int64) {
			sessionRateLimit.Wait(bytes)
		}
	}

	// We have both. Queue the bytes on both the global and session specific
	// rate limiters. Wait for both in parallell, so that the actual send
	// happens when both conditions are satisfied. In practice this just means
	// wait the longer of the two times.
	return func(bytes int64) {
		t0 := sessionRateLimit.Take(bytes)
		t1 := globalRateLimit.Take(bytes)
		if t0 > t1 {
			time.Sleep(t0)
		} else {
			time.Sleep(t1)
		}
	}
}