github.com/kyleu/dbaudit@v0.0.2-0.20240321155047-ff2f2c940496/app/util/async.go (about) 1 // Package util - Content managed by Project Forge, see [projectforge.md] for details. 2 package util 3 4 import ( 5 "sync" 6 "time" 7 8 "github.com/pkg/errors" 9 "github.com/samber/lo" 10 ) 11 12 func AsyncCollect[T any, R any](items []T, f func(x T) (R, error)) ([]R, []error) { 13 ret := make([]R, 0, len(items)) 14 var errs []error 15 mu := sync.Mutex{} 16 wg := sync.WaitGroup{} 17 wg.Add(len(items)) 18 lo.ForEach(items, func(x T, _ int) { 19 i := x 20 go func() { 21 r, err := f(i) 22 mu.Lock() 23 if err == nil { 24 ret = append(ret, r) 25 } else { 26 errs = append(errs, errors.Wrapf(err, "error running async function for item [%v]", i)) 27 } 28 mu.Unlock() 29 wg.Done() 30 }() 31 }) 32 wg.Wait() 33 return ret, errs 34 } 35 36 func AsyncCollectMap[T any, K comparable, R any](items []T, k func(x T) K, f func(x T) (R, error)) (map[K]R, map[K]error) { 37 ret := make(map[K]R, len(items)) 38 errs := map[K]error{} 39 mu := sync.Mutex{} 40 wg := sync.WaitGroup{} 41 wg.Add(len(items)) 42 lo.ForEach(items, func(x T, _ int) { 43 i := x 44 go func() { 45 key := k(i) 46 r, err := f(i) 47 mu.Lock() 48 if err == nil { 49 ret[key] = r 50 } else { 51 errs[key] = errors.Wrapf(err, "error running async function for item [%v]", key) 52 } 53 mu.Unlock() 54 wg.Done() 55 }() 56 }) 57 wg.Wait() 58 return ret, errs 59 } 60 61 func AsyncRateLimit[T any, R any](items []T, f func(x T) (R, error), maxConcurrent int, timeout time.Duration) ([]R, []error) { 62 ret := make([]R, 0, len(items)) 63 var errs []error 64 mu := sync.Mutex{} 65 idx := 0 66 67 limit := make(chan struct{}, maxConcurrent) 68 defer close(limit) 69 70 for { 71 select { 72 case limit <- struct{}{}: 73 idx++ 74 item := items[idx] 75 go func() { 76 r, err := f(item) 77 mu.Lock() 78 if err == nil { 79 ret = append(ret, r) 80 } else { 81 errs = append(errs, errors.Wrapf(err, "error running async function for item [%v]", item)) 82 } 83 mu.Unlock() 84 }() 85 case <-time.After(timeout): 86 errs = append(errs, errors.Errorf("job timed out after [%v]", timeout)) 87 return ret, errs 88 default: 89 return ret, errs 90 } 91 } 92 }