124 lines
2.6 KiB
Go
124 lines
2.6 KiB
Go
package core
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"golang.org/x/sync/errgroup"
|
|
)
|
|
|
|
type WorkGroup interface {
|
|
// Queue queues a function to run. It may be invoked immediately, or deferred until RunAndWait.
|
|
// It is not safe to call Queue after RunAndWait has returned.
|
|
Queue(fn func())
|
|
|
|
// RunAndWait runs all queued functions, blocking until they have all completed.
|
|
RunAndWait()
|
|
}
|
|
|
|
func NewWorkGroup(singleThreaded bool) WorkGroup {
|
|
if singleThreaded {
|
|
return &singleThreadedWorkGroup{}
|
|
}
|
|
return ¶llelWorkGroup{}
|
|
}
|
|
|
|
type parallelWorkGroup struct {
|
|
done atomic.Bool
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
var _ WorkGroup = (*parallelWorkGroup)(nil)
|
|
|
|
func (w *parallelWorkGroup) Queue(fn func()) {
|
|
if w.done.Load() {
|
|
panic("Queue called after RunAndWait returned")
|
|
}
|
|
|
|
w.wg.Go(func() {
|
|
fn()
|
|
})
|
|
}
|
|
|
|
func (w *parallelWorkGroup) RunAndWait() {
|
|
defer w.done.Store(true)
|
|
w.wg.Wait()
|
|
}
|
|
|
|
type singleThreadedWorkGroup struct {
|
|
done atomic.Bool
|
|
fnsMu sync.Mutex
|
|
fns []func()
|
|
}
|
|
|
|
var _ WorkGroup = (*singleThreadedWorkGroup)(nil)
|
|
|
|
func (w *singleThreadedWorkGroup) Queue(fn func()) {
|
|
if w.done.Load() {
|
|
panic("Queue called after RunAndWait returned")
|
|
}
|
|
|
|
w.fnsMu.Lock()
|
|
defer w.fnsMu.Unlock()
|
|
w.fns = append(w.fns, fn)
|
|
}
|
|
|
|
func (w *singleThreadedWorkGroup) RunAndWait() {
|
|
defer w.done.Store(true)
|
|
for {
|
|
fn := w.pop()
|
|
if fn == nil {
|
|
return
|
|
}
|
|
fn()
|
|
}
|
|
}
|
|
|
|
func (w *singleThreadedWorkGroup) pop() func() {
|
|
w.fnsMu.Lock()
|
|
defer w.fnsMu.Unlock()
|
|
if len(w.fns) == 0 {
|
|
return nil
|
|
}
|
|
end := len(w.fns) - 1
|
|
fn := w.fns[end]
|
|
w.fns[end] = nil // Allow GC
|
|
w.fns = w.fns[:end]
|
|
return fn
|
|
}
|
|
|
|
// ThrottleGroup is like errgroup.Group but with global concurrency limiting via a semaphore.
|
|
type ThrottleGroup struct {
|
|
semaphore chan struct{}
|
|
group *errgroup.Group
|
|
}
|
|
|
|
// NewThrottleGroup creates a new ThrottleGroup with the given context and semaphore for concurrency limiting.
|
|
func NewThrottleGroup(ctx context.Context, semaphore chan struct{}) *ThrottleGroup {
|
|
g, _ := errgroup.WithContext(ctx)
|
|
return &ThrottleGroup{
|
|
semaphore: semaphore,
|
|
group: g,
|
|
}
|
|
}
|
|
|
|
// Go runs the given function in a new goroutine, but first acquires a slot from the semaphore.
|
|
// The semaphore slot is released when the function completes.
|
|
func (tg *ThrottleGroup) Go(fn func() error) {
|
|
tg.group.Go(func() error {
|
|
// Acquire semaphore slot - this will block until a slot is available
|
|
tg.semaphore <- struct{}{}
|
|
defer func() {
|
|
// Release semaphore slot when done
|
|
<-tg.semaphore
|
|
}()
|
|
return fn()
|
|
})
|
|
}
|
|
|
|
// Wait waits for all goroutines to complete and returns the first error encountered, if any.
|
|
func (tg *ThrottleGroup) Wait() error {
|
|
return tg.group.Wait()
|
|
}
|