github.com/whoyao/protocol@v0.0.0-20230519045905-2d8ace718ca5/utils/parallel.go (about)

     1  package utils
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  
     7  	"go.uber.org/atomic"
     8  )
     9  
    10  // ParallelExec will executes the given function with each element of vals, if len(vals) >= parallelThreshold,
    11  // will execute them in parallel, with the given step size. So fn must be thread-safe.
    12  func ParallelExec[T any](vals []T, parallelThreshold, step uint64, fn func(T)) {
    13  	if uint64(len(vals)) < parallelThreshold {
    14  		for _, v := range vals {
    15  			fn(v)
    16  		}
    17  		return
    18  	}
    19  
    20  	// parallel - enables much more efficient multi-core utilization
    21  	start := atomic.NewUint64(0)
    22  	end := uint64(len(vals))
    23  
    24  	var wg sync.WaitGroup
    25  	numCPU := runtime.NumCPU()
    26  	wg.Add(numCPU)
    27  	for p := 0; p < numCPU; p++ {
    28  		go func() {
    29  			defer wg.Done()
    30  			for {
    31  				n := start.Add(step)
    32  				if n >= end+step {
    33  					return
    34  				}
    35  
    36  				for i := n - step; i < n && i < end; i++ {
    37  					fn(vals[i])
    38  				}
    39  			}
    40  		}()
    41  	}
    42  	wg.Wait()
    43  }