116 lines
1.8 KiB
Go
116 lines
1.8 KiB
Go
package errgroups
|
|
|
|
import (
|
|
"context"
|
|
"dashboard/utils"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"golang.org/x/sync/errgroup"
|
|
)
|
|
|
|
type ErrGroupFunc struct {
|
|
wg *errgroup.Group
|
|
mu sync.RWMutex
|
|
funs map[string]FuncWithErrFunc
|
|
}
|
|
|
|
type FuncWithErrFunc func(context.Context) error
|
|
|
|
func NewErrGroupFunc() *ErrGroupFunc {
|
|
return &ErrGroupFunc{
|
|
funs: make(map[string]FuncWithErrFunc),
|
|
}
|
|
}
|
|
|
|
func (e *ErrGroupFunc) Run(ctx context.Context) error {
|
|
e.mu.RLock()
|
|
defer e.mu.RUnlock()
|
|
|
|
e.wg, ctx = errgroup.WithContext(ctx)
|
|
|
|
for name, fn := range e.funs {
|
|
fn := fn
|
|
name := name
|
|
e.wg.Go(func() error {
|
|
defer fmt.Printf("func: %s stop\n", name)
|
|
fmt.Printf("func: %s runing...\n", name)
|
|
return fn(ctx)
|
|
})
|
|
}
|
|
|
|
if err := e.wg.Wait(); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (e *ErrGroupFunc) Add(fn FuncWithErrFunc) {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
|
|
fnStr := utils.GetFuncName(fn)
|
|
|
|
e.funs[fnStr] = fn
|
|
}
|
|
|
|
func (e *ErrGroupFunc) Del(fn FuncWithErrFunc) {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
|
|
fnStr := utils.GetFuncName(fn)
|
|
|
|
delete(e.funs, fnStr)
|
|
}
|
|
|
|
type ErrGroup struct {
|
|
wg *errgroup.Group
|
|
mu sync.RWMutex
|
|
funs map[FuncWithErr]struct{}
|
|
}
|
|
|
|
type FuncWithErr interface {
|
|
Run(context.Context) error
|
|
}
|
|
|
|
func NewErrGroup() *ErrGroup {
|
|
return &ErrGroup{
|
|
funs: make(map[FuncWithErr]struct{}),
|
|
}
|
|
}
|
|
|
|
func (e *ErrGroup) Run(ctx context.Context) error {
|
|
e.mu.RLock()
|
|
defer e.mu.RUnlock()
|
|
|
|
e.wg, ctx = errgroup.WithContext(ctx)
|
|
|
|
for fn := range e.funs {
|
|
fn := fn
|
|
e.wg.Go(func() error {
|
|
return fn.Run(ctx)
|
|
})
|
|
}
|
|
|
|
if err := e.wg.Wait(); err != nil {
|
|
return err
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (e *ErrGroup) Add(fn FuncWithErr) {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
|
|
e.funs[fn] = struct{}{}
|
|
}
|
|
|
|
func (e *ErrGroup) Del(fn FuncWithErr) {
|
|
e.mu.Lock()
|
|
defer e.mu.Unlock()
|
|
|
|
delete(e.funs, fn)
|
|
}
|