github.com/unicornultrafoundation/go-u2u@v1.0.0-rc1.0.20240205080301-e74a83d3fadc/utils/num_queue.go (about)

     1  package utils
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  type NumQueue struct {
     8  	mu       sync.Mutex
     9  	lastDone uint64
    10  	waiters  []chan struct{}
    11  }
    12  
    13  func NewNumQueue(init uint64) *NumQueue {
    14  	return &NumQueue{
    15  		lastDone: init,
    16  	}
    17  }
    18  
    19  func (q *NumQueue) Done(n uint64) {
    20  	q.mu.Lock()
    21  	defer q.mu.Unlock()
    22  
    23  	if n <= q.lastDone {
    24  		panic("Already done!")
    25  	}
    26  
    27  	pos := int(n - q.lastDone - 1)
    28  	for i := 0; i < len(q.waiters) && i <= pos; i++ {
    29  		close(q.waiters[i])
    30  	}
    31  	if pos < len(q.waiters) {
    32  		q.waiters = q.waiters[pos+1:]
    33  	} else {
    34  		q.waiters = make([]chan struct{}, 0, 1000)
    35  	}
    36  
    37  	q.lastDone = n
    38  }
    39  
    40  func (q *NumQueue) WaitFor(n uint64) {
    41  	q.mu.Lock()
    42  
    43  	if n <= q.lastDone {
    44  		q.mu.Unlock()
    45  		return
    46  	}
    47  
    48  	count := int(n - q.lastDone)
    49  	for i := len(q.waiters); i < count; i++ {
    50  		q.waiters = append(q.waiters, make(chan struct{}))
    51  	}
    52  	ch := q.waiters[count-1]
    53  	q.mu.Unlock()
    54  	<-ch
    55  }