gitlab.com/beacon-software/gadget@v0.0.0-20181217202115-54565ea1ed5e/collection/specialized/ratehashpriorityqueue.go (about)

     1  package specialized
     2  
     3  import (
     4  	"sync/atomic"
     5  	"time"
     6  
     7  	"gitlab.com/beacon-software/gadget/timeutil"
     8  )
     9  
    10  // RateHashPriorityQueue allows prioritized unique elements to be emitted
    11  // at a specific rate.
    12  type RateHashPriorityQueue interface {
    13  	HashPriorityQueue
    14  	// NoLimitPop the highest priority element off the queue ignoring the rate limit
    15  	// for the purpose of batching commands
    16  	NoLimitPop() (HashPriority, bool)
    17  	// Channel that can be used instead of Pop
    18  	Channel() <-chan HashPriority
    19  	// Stop this RateHashPriorityQueue so that it can be garbage collected.
    20  	Stop()
    21  }
    22  
    23  // NewRateHashPriorityQueue that will return at max 'n' elements per 'rate' duration.
    24  func NewRateHashPriorityQueue(n int, rate time.Duration) RateHashPriorityQueue {
    25  	q := &rhpQueue{
    26  		rate:    rate,
    27  		queue:   NewHashPriorityQueue(),
    28  		channel: make(chan HashPriority, n),
    29  		stop:    make(chan bool),
    30  	}
    31  	go q.run()
    32  	return q
    33  }
    34  
    35  // rhpQueue implementes the RateHashPriorityQueue interface
    36  type rhpQueue struct {
    37  	rate    time.Duration
    38  	queue   HashPriorityQueue
    39  	size    int32
    40  	channel chan HashPriority
    41  	stop    chan bool
    42  }
    43  
    44  func (q *rhpQueue) run() {
    45  	ticker := timeutil.NewTicker(q.rate).Start()
    46  	for {
    47  		select {
    48  		case <-ticker.Channel():
    49  			if elm, ok := q.queue.Pop(); ok {
    50  				// if this blocks either the rate has been reached
    51  				// or no one is listening. either way we want to block this thread
    52  				// on it
    53  				q.channel <- elm
    54  				atomic.AddInt32(&q.size, -1)
    55  			}
    56  			ticker.Reset()
    57  		case <-q.stop:
    58  			ticker.Stop()
    59  			close(q.channel)
    60  			return
    61  		}
    62  	}
    63  }
    64  
    65  func (q *rhpQueue) Size() int {
    66  	return int(atomic.LoadInt32(&q.size))
    67  }
    68  
    69  func (q *rhpQueue) Push(element HashPriority) {
    70  	q.queue.Push(element)
    71  	atomic.StoreInt32(&q.size, int32(q.queue.Size()))
    72  }
    73  
    74  func (q *rhpQueue) Pop() (HashPriority, bool) {
    75  	return <-q.channel, true
    76  }
    77  
    78  func (q *rhpQueue) Channel() <-chan HashPriority {
    79  	return q.channel
    80  }
    81  
    82  func (q *rhpQueue) Peek() (HashPriority, bool) {
    83  	return q.queue.Peek()
    84  }
    85  
    86  func (q *rhpQueue) NoLimitPop() (HashPriority, bool) {
    87  	select {
    88  	case elm := <-q.channel:
    89  		return elm, true
    90  	default:
    91  		if elm, ok := q.queue.Pop(); ok {
    92  			atomic.StoreInt32(&q.size, int32(q.queue.Size()))
    93  			return elm, true
    94  		}
    95  		return nil, false
    96  	}
    97  }
    98  
    99  func (q *rhpQueue) Stop() {
   100  	// non-blocking so that this is reentrant
   101  	select {
   102  	case q.stop <- true:
   103  		return
   104  	default:
   105  		return
   106  	}
   107  }