github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/structs/merge/merge.go (about)

     1  package merge
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  type Merger interface {
     8  	Merge(src Merger)
     9  }
    10  
    11  func MergeTriesConcurrently(concurrency int, tries ...Merger) Merger {
    12  	// mutex := sync.Mutex{}
    13  	if len(tries) == 0 {
    14  		return nil
    15  	}
    16  	pool := tries
    17  	jobs := make(chan []Merger)
    18  	done := make(chan Merger, len(tries)-1)
    19  	wg := sync.WaitGroup{}
    20  	wg.Add(concurrency)
    21  	for i := 0; i < concurrency; i++ {
    22  		go func() {
    23  			for j := range jobs {
    24  				j[0].Merge(j[1])
    25  				done <- j[0]
    26  			}
    27  			wg.Done()
    28  		}()
    29  	}
    30  
    31  	merges := 0
    32  	inProgress := 0
    33  	for merges < len(tries)-1 {
    34  		// first, queue all possible jobs
    35  		for len(pool) >= 2 && inProgress < concurrency {
    36  			inProgress++
    37  			j := pool[:2]
    38  			jobs <- j
    39  			pool = pool[2:]
    40  		}
    41  		// then block until there's a job done notification
    42  		t := <-done
    43  		pool = append([]Merger{t}, pool...)
    44  		merges++
    45  		inProgress--
    46  	}
    47  	close(jobs)
    48  	wg.Wait()
    49  	return pool[0]
    50  }
    51  
    52  func MergeTriesSerially(_ int, tries ...Merger) Merger {
    53  	// rand.Shuffle(len(tries), func(i, j int) {
    54  	// 	tries[i], tries[j] = tries[j], tries[i]
    55  	// })
    56  	// mutex := sync.Mutex{}
    57  	if len(tries) == 0 {
    58  		return nil
    59  	}
    60  	resultTrie := tries[0]
    61  	for i := 1; i < len(tries); i++ {
    62  		resultTrie.Merge(tries[i])
    63  	}
    64  	return resultTrie
    65  }