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  }