59 lines
1.4 KiB
Go
59 lines
1.4 KiB
Go
package goatomic
|
|
|
|
import (
|
|
"runtime"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
const waitTime = time.Nanosecond * 10
|
|
|
|
// WorkerGroup acts as a sync.WaitGroup but with a max worker count.
|
|
type WorkerGroup struct {
|
|
maxWorkers uint32
|
|
workers uint32
|
|
}
|
|
|
|
// MaxWorkers sets max workers to `workers` or returns current setting if < 0
|
|
func (wg *WorkerGroup) MaxWorkers(workers uint32) uint32 {
|
|
if workers == 0 {
|
|
// update max workers to GOMAXPROCS if set to 0
|
|
atomic.CompareAndSwapUint32(&wg.maxWorkers, 0, uint32(runtime.GOMAXPROCS(0)))
|
|
|
|
return atomic.LoadUint32(&wg.maxWorkers)
|
|
}
|
|
|
|
atomic.StoreUint32(&wg.maxWorkers, workers)
|
|
|
|
return workers
|
|
}
|
|
|
|
// Add delta to the worker count
|
|
func (wg *WorkerGroup) Add(delta uint32) {
|
|
// update max workers to GOMAXPROCES if set to 0
|
|
atomic.CompareAndSwapUint32(&wg.maxWorkers, 0, uint32(runtime.GOMAXPROCS(0)))
|
|
|
|
var oldCount uint32
|
|
for oldCount = atomic.LoadUint32(&wg.workers); oldCount+delta > atomic.LoadUint32(&wg.maxWorkers); oldCount = atomic.LoadUint32(&wg.workers) {
|
|
time.Sleep(waitTime)
|
|
}
|
|
|
|
if atomic.CompareAndSwapUint32(&wg.workers, oldCount, oldCount+delta) {
|
|
return
|
|
}
|
|
|
|
wg.Add(delta)
|
|
}
|
|
|
|
// Done decrement the worker counter
|
|
func (wg *WorkerGroup) Done() {
|
|
atomic.AddUint32(&wg.workers, ^uint32(0))
|
|
}
|
|
|
|
// Wait until worker count is zero
|
|
func (wg *WorkerGroup) Wait() {
|
|
for atomic.LoadUint32(&wg.workers) != 0 {
|
|
time.Sleep(waitTime)
|
|
}
|
|
}
|