jackhammer

Utilities for Go
Log | Files | Refs | README | LICENSE

para.go (1518B)

      1 package para
      2 
      3 import (
      4 	"context"
      5 	"fmt"
      6 	"iter"
      7 	"sync"
      8 
      9 	"golang.org/x/sync/errgroup"
     10 )
     11 
     12 // Fanout processes a sequence concurrently.
     13 func Fanout[V any](seq iter.Seq[V], fn func(value V)) {
     14 	wg := sync.WaitGroup{}
     15 	defer wg.Wait()
     16 
     17 	for v := range seq {
     18 		wg.Add(1)
     19 		go func() {
     20 			defer wg.Done()
     21 			fn(v)
     22 		}()
     23 	}
     24 }
     25 
     26 // Batch processes the sequence concurrently. Any error cancels the
     27 // entire batch. No results are generated. This is helpful for a fanout of
     28 // side-effects.
     29 func Batch[T any](ctx context.Context, seq iter.Seq[T], work func(ctx context.Context, ii int, v T) error) error {
     30 	ctx, cancel := context.WithCancel(ctx)
     31 	defer cancel()
     32 
     33 	group, ctx := errgroup.WithContext(ctx)
     34 
     35 	ii := 0
     36 	for v := range seq {
     37 		func(ii int) {
     38 			group.Go(func() error {
     39 				return work(ctx, ii, v)
     40 			})
     41 		}(ii)
     42 		ii += 1
     43 	}
     44 
     45 	return group.Wait()
     46 }
     47 
     48 // MapErr processes the sequence concurrently, transforming [In] to [Out].
     49 // Any error cancels the entire batch. This is helpful for a fanout of
     50 // fallible transformations.
     51 func MapErr[In any, Out any](ctx context.Context, seq []In, work func(ctx context.Context, ii int, v In) (Out, error)) ([]Out, error) {
     52 	ctx, cancel := context.WithCancel(ctx)
     53 	defer cancel()
     54 
     55 	group, ctx := errgroup.WithContext(ctx)
     56 
     57 	res := make([]Out, len(seq))
     58 
     59 	for ii, v := range seq {
     60 		group.Go(func() error {
     61 			out, err := work(ctx, ii, v)
     62 			res[ii] = out
     63 			return err
     64 		})
     65 	}
     66 
     67 	if err := group.Wait(); err != nil {
     68 		return nil, fmt.Errorf("map: %w", err)
     69 	}
     70 
     71 	return res, nil
     72 }