120 lines
2.4 KiB
Go
120 lines
2.4 KiB
Go
package execute
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"runtime/debug"
|
|
"sync"
|
|
)
|
|
|
|
// Dispatcher schedules work for a query.
|
|
// Each transformation submits work to be done to the dispatcher.
|
|
// Then the dispatcher schedules to work based on the available resources.
|
|
type Dispatcher interface {
|
|
// Schedule fn to be executed
|
|
Schedule(fn ScheduleFunc)
|
|
}
|
|
|
|
// ScheduleFunc is a function that represents work to do.
|
|
// The throughput is the maximum number of messages to process for this scheduling.
|
|
type ScheduleFunc func(throughput int)
|
|
|
|
// poolDispatcher implements Dispatcher using a pool of goroutines.
|
|
type poolDispatcher struct {
|
|
work chan ScheduleFunc
|
|
|
|
throughput int
|
|
|
|
mu sync.Mutex
|
|
closed bool
|
|
closing chan struct{}
|
|
wg sync.WaitGroup
|
|
err error
|
|
errC chan error
|
|
}
|
|
|
|
func newPoolDispatcher(throughput int) *poolDispatcher {
|
|
return &poolDispatcher{
|
|
throughput: throughput,
|
|
work: make(chan ScheduleFunc, 100),
|
|
closing: make(chan struct{}),
|
|
errC: make(chan error, 1),
|
|
}
|
|
}
|
|
|
|
func (d *poolDispatcher) Schedule(fn ScheduleFunc) {
|
|
select {
|
|
case d.work <- fn:
|
|
case <-d.closing:
|
|
}
|
|
}
|
|
|
|
func (d *poolDispatcher) Start(n int, ctx context.Context) {
|
|
d.wg.Add(n)
|
|
for i := 0; i < n; i++ {
|
|
go func() {
|
|
defer d.wg.Done()
|
|
// Setup panic handling on the worker goroutines
|
|
defer func() {
|
|
if e := recover(); e != nil {
|
|
var err error
|
|
switch e := e.(type) {
|
|
case error:
|
|
err = e
|
|
default:
|
|
err = fmt.Errorf("%v", e)
|
|
}
|
|
d.setErr(fmt.Errorf("panic: %v\n%s", err, debug.Stack()))
|
|
}
|
|
}()
|
|
d.run(ctx)
|
|
}()
|
|
}
|
|
}
|
|
|
|
// Err returns a channel with will produce an error if encountered.
|
|
func (d *poolDispatcher) Err() <-chan error {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
return d.errC
|
|
}
|
|
|
|
func (d *poolDispatcher) setErr(err error) {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
// TODO(nathanielc): Collect all error information.
|
|
if d.err == nil {
|
|
d.err = err
|
|
d.errC <- err
|
|
}
|
|
}
|
|
|
|
//Stop the dispatcher.
|
|
func (d *poolDispatcher) Stop() error {
|
|
d.mu.Lock()
|
|
defer d.mu.Unlock()
|
|
if d.closed {
|
|
return d.err
|
|
}
|
|
d.closed = true
|
|
close(d.closing)
|
|
d.wg.Wait()
|
|
return d.err
|
|
}
|
|
|
|
// run is the logic executed by each worker goroutine in the pool.
|
|
func (d *poolDispatcher) run(ctx context.Context) {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
// Immediately return, do not process any more work
|
|
return
|
|
case <-d.closing:
|
|
// We are done, nothing left to do.
|
|
return
|
|
case fn := <-d.work:
|
|
fn(d.throughput)
|
|
}
|
|
}
|
|
}
|