97 lines
1.7 KiB
Go
97 lines
1.7 KiB
Go
package runnergroup
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// RunnerGroup is like sync.WaitGroup,
|
|
// the diffrence is if one task stops, all will be stopped.
|
|
type RunnerGroup struct {
|
|
runners []*Runner
|
|
done chan byte
|
|
}
|
|
|
|
type Runner struct {
|
|
// Start is a blocking function.
|
|
Start func() error
|
|
// Stop is not a blocking function, if Stop called, must let Start return.
|
|
// Notice: Stop maybe called multi times even if before start.
|
|
Stop func() error
|
|
lock sync.Mutex
|
|
status int
|
|
}
|
|
|
|
func New() *RunnerGroup {
|
|
g := &RunnerGroup{}
|
|
g.runners = make([]*Runner, 0)
|
|
g.done = make(chan byte)
|
|
return g
|
|
}
|
|
|
|
func (g *RunnerGroup) Add(r *Runner) {
|
|
g.runners = append(g.runners, r)
|
|
}
|
|
|
|
// Call Wait after all task have been added,
|
|
// Return the first ended start's result.
|
|
func (g *RunnerGroup) Wait() error {
|
|
e := make(chan error)
|
|
for _, v := range g.runners {
|
|
v.status = 1
|
|
go func(v *Runner) {
|
|
err := v.Start()
|
|
v.lock.Lock()
|
|
v.status = 0
|
|
v.lock.Unlock()
|
|
select {
|
|
case <-g.done:
|
|
case e <- err:
|
|
}
|
|
}(v)
|
|
}
|
|
err := <-e
|
|
for _, v := range g.runners {
|
|
for {
|
|
v.lock.Lock()
|
|
if v.status == 0 {
|
|
v.lock.Unlock()
|
|
break
|
|
}
|
|
v.lock.Unlock()
|
|
_ = v.Stop()
|
|
time.Sleep(300 * time.Millisecond)
|
|
}
|
|
}
|
|
close(g.done)
|
|
return err
|
|
}
|
|
|
|
// Call Done if you want to stop all.
|
|
// return the stop's return which is not nil, do not guarantee,
|
|
// because starts may ended caused by itself.
|
|
func (g *RunnerGroup) Done() error {
|
|
if len(g.runners) == 0 {
|
|
return nil
|
|
}
|
|
var e error
|
|
for _, v := range g.runners {
|
|
for {
|
|
v.lock.Lock()
|
|
if v.status == 0 {
|
|
v.lock.Unlock()
|
|
break
|
|
}
|
|
v.lock.Unlock()
|
|
if err := v.Stop(); err != nil {
|
|
if e == nil {
|
|
e = err
|
|
}
|
|
}
|
|
time.Sleep(300 * time.Millisecond)
|
|
}
|
|
}
|
|
<-g.done
|
|
return e
|
|
}
|