Rework batch package
Use a buffered channel to limit the number of active goroutines without needing to wait for all active goroutines to finish
This commit is contained in:
@@ -1,22 +1,17 @@
|
||||
|
||||
package batch
|
||||
|
||||
import (
|
||||
"time"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type BatchPool struct {
|
||||
active int
|
||||
mintime int
|
||||
jobcount int
|
||||
jobs []func()
|
||||
jobs []func()
|
||||
}
|
||||
|
||||
func New(active int, mintime int) *BatchPool {
|
||||
func New(active int) *BatchPool {
|
||||
pool := &BatchPool{
|
||||
active: active,
|
||||
mintime: mintime,
|
||||
}
|
||||
pool.jobs = []func(){}
|
||||
return pool
|
||||
@@ -30,28 +25,24 @@ func Min(x, y int) int {
|
||||
}
|
||||
|
||||
func (pool *BatchPool) Submit(f func()) {
|
||||
pool.jobcount++
|
||||
pool.jobs = append(pool.jobs, f)
|
||||
}
|
||||
|
||||
func wrapper(wg *sync.WaitGroup, f func()) {
|
||||
defer wg.Done()
|
||||
func wrapper(wg *sync.WaitGroup, limiter chan struct{}, f func()) {
|
||||
f()
|
||||
<-limiter
|
||||
wg.Done()
|
||||
}
|
||||
|
||||
func (pool *BatchPool) Run() {
|
||||
var wg sync.WaitGroup
|
||||
count := pool.jobcount
|
||||
for count > 0 {
|
||||
for i:=0; i<Min(count, pool.active); i++ {
|
||||
wg.Add(1)
|
||||
go wrapper(&wg, pool.jobs[pool.jobcount-count])
|
||||
count--
|
||||
}
|
||||
if (pool.mintime > 0) {
|
||||
time.Sleep(time.Second * time.Duration(pool.mintime))
|
||||
}
|
||||
wg.Wait()
|
||||
limiter := make(chan struct{}, pool.active)
|
||||
for i := range pool.jobs {
|
||||
limiter <- struct{}{}
|
||||
wg.Add(1)
|
||||
go wrapper(&wg, limiter, pool.jobs[i])
|
||||
}
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
close(limiter)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user