Merge "Revert "Refactor rate_limit.go for more clarify""

This commit is contained in:
Colin Cross 2017-08-17 23:56:00 +00:00 committed by Gerrit Code Review
commit 1115ddc838
2 changed files with 81 additions and 127 deletions

View file

@ -15,54 +15,71 @@
package main
import (
"fmt"
"runtime"
)
type RateLimit struct {
requests chan request
completions chan int64
stop chan struct{}
requests chan struct{}
finished chan int
released chan int
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{}),
// NewRateLimit starts a new rate limiter with maxExecs number of executions
// allowed to happen at a time. If maxExecs is <= 0, it will default to the
// number of logical CPUs on the system.
//
// With Finish and Release, we'll keep track of outstanding buffer sizes to be
// written. If that size goes above maxMem, we'll prevent starting new
// executions.
//
// The total memory use may be higher due to current executions. This just
// prevents runaway memory use due to slower writes.
func NewRateLimit(maxExecs int, maxMem int64) *RateLimit {
if maxExecs <= 0 {
maxExecs = runtime.NumCPU()
}
if maxMem <= 0 {
// Default to 512MB
maxMem = 512 * 1024 * 1024
}
go ret.monitorChannels(capacity)
ret := &RateLimit{
requests: make(chan struct{}),
// Let all of the pending executions to mark themselves as finished,
// even if our goroutine isn't processing input.
finished: make(chan int, maxExecs),
released: make(chan int),
stop: make(chan struct{}),
}
go ret.goFunc(maxExecs, maxMem)
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
// RequestExecution blocks until another execution can be allowed to run.
func (r *RateLimit) RequestExecution() Execution {
<-r.requests
return r.finished
}
// Finish declares the completion of an execution of size <size>
func (r *RateLimit) Finish(size int64) {
r.completions <- size
type Execution chan<- int
// Finish will mark your execution as finished, and allow another request to be
// approved.
//
// bufferSize may be specified to count memory buffer sizes, and must be
// matched with calls to RateLimit.Release to mark the buffers as released.
func (e Execution) Finish(bufferSize int) {
e <- bufferSize
}
// Call Release when finished with a buffer recorded with Finish.
func (r *RateLimit) Release(bufferSize int) {
r.released <- bufferSize
}
// Stop the background goroutine
@ -70,83 +87,29 @@ 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
func (r *RateLimit) goFunc(maxExecs int, maxMem int64) {
var curExecs int
var curMemory int64
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
var requests chan struct{}
if curExecs < maxExecs && curMemory < maxMem {
requests = r.requests
}
select {
case newRequest := <-requests:
currentRequest = &newRequest
break
case amountCompleted := <-r.completions:
usedCapacity -= amountCompleted
if usedCapacity < 0 {
panic(fmt.Sprintf("usedCapacity < 0: %v", usedCapacity))
case requests <- struct{}{}:
curExecs++
case amount := <-r.finished:
curExecs--
curMemory += int64(amount)
if curExecs < 0 {
panic("curExecs < 0")
}
case amount := <-r.released:
curMemory -= int64(amount)
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}
}

View file

@ -163,8 +163,7 @@ type zipWriter struct {
errors chan error
writeOps chan chan *zipEntry
cpuRateLimiter *CPURateLimiter
memoryRateLimiter *MemoryRateLimiter
rateLimit *RateLimit
compressorPool sync.Pool
compLevel int
@ -296,12 +295,8 @@ func (z *zipWriter) write(out string, pathMappings []pathMapping, manifest strin
// The RateLimit object will put the upper bounds on the number of
// parallel compressions and outstanding buffers.
z.writeOps = make(chan chan *zipEntry, 1000)
z.cpuRateLimiter = NewCPURateLimiter(int64(*parallelJobs))
z.memoryRateLimiter = NewMemoryRateLimiter(0)
defer func() {
z.cpuRateLimiter.Stop()
z.memoryRateLimiter.Stop()
}()
z.rateLimit = NewRateLimit(*parallelJobs, 0)
defer z.rateLimit.Stop()
go func() {
var err error
@ -391,7 +386,7 @@ func (z *zipWriter) write(out string, pathMappings []pathMapping, manifest strin
if err != nil {
return err
}
z.memoryRateLimiter.Finish(count)
z.rateLimit.Release(int(count))
currentReader = nil
@ -461,7 +456,7 @@ func (z *zipWriter) writeFile(dest, src string, method uint16) error {
return err
}
z.cpuRateLimiter.Request()
exec := z.rateLimit.RequestExecution()
if method == zip.Deflate && fileSize >= minParallelFileSize {
wg := new(sync.WaitGroup)
@ -478,14 +473,14 @@ func (z *zipWriter) writeFile(dest, src string, method uint16) error {
// know the result before we can begin writing the compressed
// data out to the zipfile.
wg.Add(1)
go z.crcFile(r, ze, compressChan, wg)
go z.crcFile(r, ze, exec, compressChan, wg)
for start := int64(0); start < fileSize; start += parallelBlockSize {
sr := io.NewSectionReader(r, start, parallelBlockSize)
resultChan := make(chan io.Reader, 1)
ze.futureReaders <- resultChan
z.cpuRateLimiter.Request()
exec := z.rateLimit.RequestExecution()
last := !(start+parallelBlockSize < fileSize)
var dict []byte
@ -494,7 +489,7 @@ func (z *zipWriter) writeFile(dest, src string, method uint16) error {
}
wg.Add(1)
go z.compressPartialFile(sr, dict, last, resultChan, wg)
go z.compressPartialFile(sr, dict, last, exec, resultChan, wg)
}
close(ze.futureReaders)
@ -505,15 +500,15 @@ func (z *zipWriter) writeFile(dest, src string, method uint16) error {
f.Close()
}(wg, r)
} else {
go z.compressWholeFile(ze, r, compressChan)
go z.compressWholeFile(ze, r, exec, compressChan)
}
return nil
}
func (z *zipWriter) crcFile(r io.Reader, ze *zipEntry, resultChan chan *zipEntry, wg *sync.WaitGroup) {
func (z *zipWriter) crcFile(r io.Reader, ze *zipEntry, exec Execution, resultChan chan *zipEntry, wg *sync.WaitGroup) {
defer wg.Done()
defer z.cpuRateLimiter.Finish()
defer exec.Finish(0)
crc := crc32.NewIEEE()
_, err := io.Copy(crc, r)
@ -527,7 +522,7 @@ func (z *zipWriter) crcFile(r io.Reader, ze *zipEntry, resultChan chan *zipEntry
close(resultChan)
}
func (z *zipWriter) compressPartialFile(r io.Reader, dict []byte, last bool, resultChan chan io.Reader, wg *sync.WaitGroup) {
func (z *zipWriter) compressPartialFile(r io.Reader, dict []byte, last bool, exec Execution, resultChan chan io.Reader, wg *sync.WaitGroup) {
defer wg.Done()
result, err := z.compressBlock(r, dict, last)
@ -536,9 +531,7 @@ func (z *zipWriter) compressPartialFile(r io.Reader, dict []byte, last bool, res
return
}
z.memoryRateLimiter.Request(int64(result.Len()))
z.cpuRateLimiter.Finish()
exec.Finish(result.Len())
resultChan <- result
}
@ -576,7 +569,7 @@ func (z *zipWriter) compressBlock(r io.Reader, dict []byte, last bool) (*bytes.B
return buf, nil
}
func (z *zipWriter) compressWholeFile(ze *zipEntry, r *os.File, compressChan chan *zipEntry) {
func (z *zipWriter) compressWholeFile(ze *zipEntry, r *os.File, exec Execution, compressChan chan *zipEntry) {
var bufSize int
defer r.Close()
@ -645,9 +638,7 @@ func (z *zipWriter) compressWholeFile(ze *zipEntry, r *os.File, compressChan cha
bufSize = int(ze.fh.UncompressedSize64)
}
z.memoryRateLimiter.Request(int64(bufSize))
z.cpuRateLimiter.Finish()
exec.Finish(bufSize)
close(futureReader)
compressChan <- ze
@ -718,7 +709,7 @@ func (z *zipWriter) writeSymlink(rel, file string) error {
// We didn't ask permission to execute, since this should be very short
// but we still need to increment the outstanding buffer sizes, since
// the read will decrement the buffer size.
z.memoryRateLimiter.Finish(int64(-len(dest)))
z.rateLimit.Release(-len(dest))
ze <- &zipEntry{
fh: fileHeader,