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 15 package main 16 17 import ( 18 "runtime" 19 ) 20 21 type RateLimit struct { 22 requests chan struct{} 23 finished chan int 24 released chan int 25 stop chan struct{} 26 } 27 28 // NewRateLimit starts a new rate limiter with maxExecs number of executions 29 // allowed to happen at a time. If maxExecs is <= 0, it will default to the 30 // number of logical CPUs on the system. 31 // 32 // With Finish and Release, we'll keep track of outstanding buffer sizes to be 33 // written. If that size goes above maxMem, we'll prevent starting new 34 // executions. 35 // 36 // The total memory use may be higher due to current executions. This just 37 // prevents runaway memory use due to slower writes. 38 func NewRateLimit(maxExecs int, maxMem int64) *RateLimit { 39 if maxExecs <= 0 { 40 maxExecs = runtime.NumCPU() 41 } 42 if maxMem <= 0 { 43 // Default to 512MB 44 maxMem = 512 * 1024 * 1024 45 } 46 47 ret := &RateLimit{ 48 requests: make(chan struct{}), 49 50 // Let all of the pending executions to mark themselves as finished, 51 // even if our goroutine isn't processing input. 52 finished: make(chan int, maxExecs), 53 54 released: make(chan int), 55 stop: make(chan struct{}), 56 } 57 58 go ret.goFunc(maxExecs, maxMem) 59 60 return ret 61 } 62 63 // RequestExecution blocks until another execution can be allowed to run. 64 func (r *RateLimit) RequestExecution() Execution { 65 <-r.requests 66 return r.finished 67 } 68 69 type Execution chan<- int 70 71 // Finish will mark your execution as finished, and allow another request to be 72 // approved. 73 // 74 // bufferSize may be specified to count memory buffer sizes, and must be 75 // matched with calls to RateLimit.Release to mark the buffers as released. 76 func (e Execution) Finish(bufferSize int) { 77 e <- bufferSize 78 } 79 80 // Call Release when finished with a buffer recorded with Finish. 81 func (r *RateLimit) Release(bufferSize int) { 82 r.released <- bufferSize 83 } 84 85 // Stop the background goroutine 86 func (r *RateLimit) Stop() { 87 close(r.stop) 88 } 89 90 func (r *RateLimit) goFunc(maxExecs int, maxMem int64) { 91 var curExecs int 92 var curMemory int64 93 94 for { 95 var requests chan struct{} 96 if curExecs < maxExecs && curMemory < maxMem { 97 requests = r.requests 98 } 99 100 select { 101 case requests <- struct{}{}: 102 curExecs++ 103 case amount := <-r.finished: 104 curExecs-- 105 curMemory += int64(amount) 106 if curExecs < 0 { 107 panic("curExecs < 0") 108 } 109 case amount := <-r.released: 110 curMemory -= int64(amount) 111 case <-r.stop: 112 return 113 } 114 } 115 } 116
