github.com/TrueBlocks/trueblocks-core/src/apps/chifra@v0.0.0-20241022031540-b362680128f7/pkg/utils/iteration.go (about)

     1  package utils
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"runtime"
     7  	"sync"
     8  
     9  	"github.com/panjf2000/ants/v2"
    10  )
    11  
    12  // stepFunc represents single iteration step
    13  type stepFunc[Key, Value any] func(key Key, value Value) error
    14  
    15  // IterateOverMap distributes batches of `target` items to a pool of goroutines which execute
    16  // `step` for every `target` item. Use benchmarks to make sure that concurrent iteration
    17  // is faster than synchronous one.
    18  func IterateOverMap[Key comparable, Value any](ctx context.Context, errorChan chan error, target map[Key]Value, step stepFunc[Key, Value]) {
    19  	var wg sync.WaitGroup
    20  	defer close(errorChan)
    21  
    22  	max := func(x, y int) int {
    23  		if x > y {
    24  			return x
    25  		}
    26  		return y
    27  	}
    28  
    29  	nRoutines := max(1, runtime.GOMAXPROCS(0))
    30  	itemsPerPool := max(1, len(target)/nRoutines)
    31  
    32  	type stepArguments struct {
    33  		key   Key
    34  		value Value
    35  	}
    36  
    37  	pool, err := ants.NewPoolWithFunc(nRoutines, func(i interface{}) {
    38  		defer wg.Done()
    39  		select {
    40  		case <-ctx.Done():
    41  			return
    42  		default:
    43  			args, ok := i.([]stepArguments)
    44  			if !ok {
    45  				errorChan <- fmt.Errorf("invalid worker argument: %v", i)
    46  				return
    47  			}
    48  			for _, arg := range args {
    49  				if err := step(arg.key, arg.value); err != nil {
    50  					errorChan <- err
    51  					return
    52  				}
    53  			}
    54  		}
    55  	})
    56  	if err != nil {
    57  		errorChan <- err
    58  		return
    59  	}
    60  	defer pool.Release()
    61  
    62  	// payload is a batch of map items that we distribute to the pool
    63  	payload := make([]stepArguments, 0, itemsPerPool)
    64  	count := 0
    65  
    66  	// distribute batches of map items
    67  	for key, value := range target {
    68  		count++
    69  		payload = append(payload, stepArguments{key, value})
    70  
    71  		if count%itemsPerPool == 0 {
    72  			wg.Add(1)
    73  			err = pool.Invoke(payload)
    74  			if err != nil {
    75  				errorChan <- err
    76  				return
    77  			}
    78  			// new batch
    79  			payload = make([]stepArguments, 0, itemsPerPool)
    80  		}
    81  	}
    82  
    83  	// don't forget about the remainder...
    84  	if len(payload) > 0 {
    85  		wg.Add(1)
    86  		err = pool.Invoke(payload)
    87  		if err != nil {
    88  			errorChan <- err
    89  			return
    90  		}
    91  	}
    92  
    93  	wg.Wait()
    94  }