rate_limit.go 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  1. // Copyright 2016 Google Inc. All rights reserved.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package zip
  15. import (
  16. "fmt"
  17. "runtime"
  18. )
  19. type RateLimit struct {
  20. requests chan request
  21. completions chan int64
  22. stop chan struct{}
  23. }
  24. type request struct {
  25. size int64
  26. serviced chan struct{}
  27. }
  28. // NewRateLimit starts a new rate limiter that permits the usage of up to <capacity> at once,
  29. // except when no capacity is in use, in which case the first caller is always permitted
  30. func NewRateLimit(capacity int64) *RateLimit {
  31. ret := &RateLimit{
  32. requests: make(chan request),
  33. completions: make(chan int64),
  34. stop: make(chan struct{}),
  35. }
  36. go ret.monitorChannels(capacity)
  37. return ret
  38. }
  39. // RequestExecution blocks until another execution of size <size> can be allowed to run.
  40. func (r *RateLimit) Request(size int64) {
  41. request := request{
  42. size: size,
  43. serviced: make(chan struct{}, 1),
  44. }
  45. // wait for the request to be received
  46. r.requests <- request
  47. // wait for the request to be accepted
  48. <-request.serviced
  49. }
  50. // Finish declares the completion of an execution of size <size>
  51. func (r *RateLimit) Finish(size int64) {
  52. r.completions <- size
  53. }
  54. // Stop the background goroutine
  55. func (r *RateLimit) Stop() {
  56. close(r.stop)
  57. }
  58. // monitorChannels processes incoming requests from channels
  59. func (r *RateLimit) monitorChannels(capacity int64) {
  60. var usedCapacity int64
  61. var currentRequest *request
  62. for {
  63. var requests chan request
  64. if currentRequest == nil {
  65. // If we don't already have a queued request, then we should check for a new request
  66. requests = r.requests
  67. }
  68. select {
  69. case newRequest := <-requests:
  70. currentRequest = &newRequest
  71. case amountCompleted := <-r.completions:
  72. usedCapacity -= amountCompleted
  73. if usedCapacity < 0 {
  74. panic(fmt.Sprintf("usedCapacity < 0: %v (decreased by %v)", usedCapacity, amountCompleted))
  75. }
  76. case <-r.stop:
  77. return
  78. }
  79. if currentRequest != nil {
  80. accepted := false
  81. if usedCapacity == 0 {
  82. accepted = true
  83. } else {
  84. if capacity >= usedCapacity+currentRequest.size {
  85. accepted = true
  86. }
  87. }
  88. if accepted {
  89. usedCapacity += currentRequest.size
  90. currentRequest.serviced <- struct{}{}
  91. currentRequest = nil
  92. }
  93. }
  94. }
  95. }
  96. // A CPURateLimiter limits the number of active calls based on CPU requirements
  97. type CPURateLimiter struct {
  98. impl *RateLimit
  99. }
  100. func NewCPURateLimiter(capacity int64) *CPURateLimiter {
  101. if capacity <= 0 {
  102. capacity = int64(runtime.NumCPU())
  103. }
  104. impl := NewRateLimit(capacity)
  105. return &CPURateLimiter{impl: impl}
  106. }
  107. func (e CPURateLimiter) Request() {
  108. e.impl.Request(1)
  109. }
  110. func (e CPURateLimiter) Finish() {
  111. e.impl.Finish(1)
  112. }
  113. func (e CPURateLimiter) Stop() {
  114. e.impl.Stop()
  115. }
  116. // A MemoryRateLimiter limits the number of active calls based on Memory requirements
  117. type MemoryRateLimiter struct {
  118. *RateLimit
  119. }
  120. func NewMemoryRateLimiter(capacity int64) *MemoryRateLimiter {
  121. if capacity <= 0 {
  122. capacity = 512 * 1024 * 1024 // 512MB
  123. }
  124. impl := NewRateLimit(capacity)
  125. return &MemoryRateLimiter{RateLimit: impl}
  126. }