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 }