123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152 |
- // Copyright 2016 Google Inc. All rights reserved.
- //
- // Licensed under the Apache License, Version 2.0 (the "License");
- // you may not use this file except in compliance with the License.
- // You may obtain a copy of the License at
- //
- // http://www.apache.org/licenses/LICENSE-2.0
- //
- // Unless required by applicable law or agreed to in writing, software
- // distributed under the License is distributed on an "AS IS" BASIS,
- // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- // See the License for the specific language governing permissions and
- // limitations under the License.
- package zip
- import (
- "fmt"
- "runtime"
- )
- type RateLimit struct {
- requests chan request
- completions chan int64
- stop chan struct{}
- }
- type request struct {
- size int64
- serviced chan struct{}
- }
- // NewRateLimit starts a new rate limiter that permits the usage of up to <capacity> at once,
- // except when no capacity is in use, in which case the first caller is always permitted
- func NewRateLimit(capacity int64) *RateLimit {
- ret := &RateLimit{
- requests: make(chan request),
- completions: make(chan int64),
- stop: make(chan struct{}),
- }
- go ret.monitorChannels(capacity)
- return ret
- }
- // RequestExecution blocks until another execution of size <size> can be allowed to run.
- func (r *RateLimit) Request(size int64) {
- request := request{
- size: size,
- serviced: make(chan struct{}, 1),
- }
- // wait for the request to be received
- r.requests <- request
- // wait for the request to be accepted
- <-request.serviced
- }
- // Finish declares the completion of an execution of size <size>
- func (r *RateLimit) Finish(size int64) {
- r.completions <- size
- }
- // Stop the background goroutine
- func (r *RateLimit) Stop() {
- close(r.stop)
- }
- // monitorChannels processes incoming requests from channels
- func (r *RateLimit) monitorChannels(capacity int64) {
- var usedCapacity int64
- var currentRequest *request
- for {
- var requests chan request
- if currentRequest == nil {
- // If we don't already have a queued request, then we should check for a new request
- requests = r.requests
- }
- select {
- case newRequest := <-requests:
- currentRequest = &newRequest
- case amountCompleted := <-r.completions:
- usedCapacity -= amountCompleted
- if usedCapacity < 0 {
- panic(fmt.Sprintf("usedCapacity < 0: %v (decreased by %v)", usedCapacity, amountCompleted))
- }
- case <-r.stop:
- return
- }
- if currentRequest != nil {
- accepted := false
- if usedCapacity == 0 {
- accepted = true
- } else {
- if capacity >= usedCapacity+currentRequest.size {
- accepted = true
- }
- }
- if accepted {
- usedCapacity += currentRequest.size
- currentRequest.serviced <- struct{}{}
- currentRequest = nil
- }
- }
- }
- }
- // A CPURateLimiter limits the number of active calls based on CPU requirements
- type CPURateLimiter struct {
- impl *RateLimit
- }
- func NewCPURateLimiter(capacity int64) *CPURateLimiter {
- if capacity <= 0 {
- capacity = int64(runtime.NumCPU())
- }
- impl := NewRateLimit(capacity)
- return &CPURateLimiter{impl: impl}
- }
- func (e CPURateLimiter) Request() {
- e.impl.Request(1)
- }
- func (e CPURateLimiter) Finish() {
- e.impl.Finish(1)
- }
- func (e CPURateLimiter) Stop() {
- e.impl.Stop()
- }
- // A MemoryRateLimiter limits the number of active calls based on Memory requirements
- type MemoryRateLimiter struct {
- *RateLimit
- }
- func NewMemoryRateLimiter(capacity int64) *MemoryRateLimiter {
- if capacity <= 0 {
- capacity = 512 * 1024 * 1024 // 512MB
- }
- impl := NewRateLimit(capacity)
- return &MemoryRateLimiter{RateLimit: impl}
- }
|