github.com/kaiya/goutils@v1.0.1-0.20230226104005-4ae4a4dc3688/routine/map.go (about)

     1  package routine
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	multierror "github.com/hashicorp/go-multierror"
     9  	"github.com/pkg/errors"
    10  	"gitlab.momoso.com/cm/kit/third_party/lg"
    11  	"golang.org/x/sync/errgroup"
    12  )
    13  
    14  // Map is a helper function to parallel operate on slice.
    15  // Example usage:
    16  // ```
    17  //   itemIDs := []int{1, 2, 3}
    18  //   items := make([]*Item, len(itemIDs))
    19  //   routine.Map(len(items), 10, func(i int) {
    20  //     items[i] = getItemByID(itemIDs[i])
    21  //   })
    22  // ```
    23  func Map(size, numWorker int, f func(i int) error) error {
    24  	return MapWithTimeout(size, numWorker, time.Hour*24*365, func(ctx context.Context, i int) error {
    25  		err := f(i)
    26  		return err
    27  	})
    28  }
    29  
    30  // MapWithTimeout is the same as Map, except each function need to be executed within timeout.
    31  // Important reminder to handle the context provided:
    32  // It is better to check the `ctx.Err() == nil` before any write operation, because the
    33  // function may be canceled at any momement, and the underlying resource refering may have already been
    34  // discarded.
    35  func MapWithTimeout(size, numWorker int, timeout time.Duration, f func(ctx context.Context, i int) error) error {
    36  	return MapWithCtxTimeout(context.Background(), size, numWorker, timeout, f)
    37  }
    38  
    39  func MapWithContext(ctx context.Context, size, numWorker int, f func(ctx context.Context, i int) error) error {
    40  	return MapWithCtxTimeout(ctx, size, numWorker, time.Hour*24*365, f)
    41  }
    42  
    43  func MapWithCtxTimeout(baseCtx context.Context, size, numWorker int, timeout time.Duration, f func(ctx context.Context, i int) error) error {
    44  	if numWorker <= 0 {
    45  		return errors.New("no num worker has been specified")
    46  	}
    47  	if numWorker == 1 {
    48  		return doSimpleMap(baseCtx, size, f)
    49  	}
    50  	var allerror error
    51  	ch := make(chan int)
    52  	var mutex sync.Mutex
    53  	go func() {
    54  		defer close(ch)
    55  		for i := 0; i < size; i++ {
    56  			ch <- i
    57  		}
    58  	}()
    59  
    60  	delayTerminate := false
    61  	grp := errgroup.Group{}
    62  	for i := 0; i < numWorker; i++ {
    63  		grp.Go((func() error {
    64  			for idx := range ch {
    65  				if baseCtx.Err() != nil {
    66  					mutex.Lock()
    67  					allerror = multierror.Append(allerror, errors.Wrapf(baseCtx.Err(), "Task[%d]", idx))
    68  					mutex.Unlock()
    69  					break
    70  				}
    71  				ctx, cancel := context.WithTimeout(baseCtx, timeout)
    72  				start := time.Now()
    73  				err := f(ctx, idx)
    74  				if err != nil {
    75  					mutex.Lock()
    76  					allerror = multierror.Append(allerror, errors.Wrapf(baseCtx.Err(), "Task[%d]", idx))
    77  					mutex.Unlock()
    78  				}
    79  				if ctx.Err() != nil && time.Since(start) > timeout+time.Second {
    80  					delayTerminate = true
    81  				}
    82  				if ctx.Err() != nil {
    83  					mutex.Lock()
    84  					allerror = multierror.Append(allerror, errors.Wrapf(baseCtx.Err(), "Task[%d]", idx))
    85  					mutex.Unlock()
    86  
    87  				}
    88  				cancel()
    89  			}
    90  			return nil
    91  		}))
    92  	}
    93  	grp.Wait()
    94  
    95  	for range ch {
    96  		//pass
    97  	}
    98  	if delayTerminate {
    99  		lg.Warn("worker dont respect the ctx", lg.GetFuncName(f))
   100  	}
   101  	return allerror
   102  }
   103  
   104  func doSimpleMap(ctx context.Context, size int, f func(ctx context.Context, i int) error) error {
   105  	var allerror error
   106  	for i := 0; i < size; i++ {
   107  		if ctx.Err() != nil {
   108  			return ctx.Err()
   109  		}
   110  		err := f(ctx, i)
   111  		if err != nil {
   112  			allerror = multierror.Append(allerror, err)
   113  		}
   114  	}
   115  	return allerror
   116  }