timed_queue.go 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. package collect
  2. import (
  3. "container/heap"
  4. "sync"
  5. "time"
  6. )
  7. type timedQueueEntry struct {
  8. timeSec int64
  9. value interface{}
  10. }
  11. type timedQueueImpl []*timedQueueEntry
  12. func (queue timedQueueImpl) Len() int {
  13. return len(queue)
  14. }
  15. func (queue timedQueueImpl) Less(i, j int) bool {
  16. return queue[i].timeSec < queue[j].timeSec
  17. }
  18. func (queue timedQueueImpl) Swap(i, j int) {
  19. tmp := queue[i]
  20. queue[i] = queue[j]
  21. queue[j] = tmp
  22. }
  23. func (queue *timedQueueImpl) Push(value interface{}) {
  24. entry := value.(*timedQueueEntry)
  25. *queue = append(*queue, entry)
  26. }
  27. func (queue *timedQueueImpl) Pop() interface{} {
  28. old := *queue
  29. n := len(old)
  30. v := old[n-1]
  31. *queue = old[:n-1]
  32. return v
  33. }
  34. type TimedQueue struct {
  35. queue timedQueueImpl
  36. access sync.RWMutex
  37. removed chan interface{}
  38. }
  39. func NewTimedQueue(updateInterval int) *TimedQueue {
  40. queue := &TimedQueue{
  41. queue: make([]*timedQueueEntry, 0, 256),
  42. removed: make(chan interface{}, 16),
  43. access: sync.RWMutex{},
  44. }
  45. go queue.cleanup(time.Tick(time.Duration(updateInterval) * time.Second))
  46. return queue
  47. }
  48. func (queue *TimedQueue) Add(value interface{}, time2Remove int64) {
  49. queue.access.Lock()
  50. heap.Push(&queue.queue, &timedQueueEntry{
  51. timeSec: time2Remove,
  52. value: value,
  53. })
  54. queue.access.Unlock()
  55. }
  56. func (queue *TimedQueue) RemovedEntries() <-chan interface{} {
  57. return queue.removed
  58. }
  59. func (queue *TimedQueue) cleanup(tick <-chan time.Time) {
  60. for {
  61. now := <-tick
  62. queue.access.RLock()
  63. queueLen := queue.queue.Len()
  64. queue.access.RUnlock()
  65. if queueLen == 0 {
  66. continue
  67. }
  68. nowSec := now.UTC().Unix()
  69. queue.access.RLock()
  70. firstEntryTime := queue.queue[0].timeSec
  71. queue.access.RUnlock()
  72. if firstEntryTime > nowSec {
  73. continue
  74. }
  75. queue.access.Lock()
  76. firstEntryValue := heap.Pop(&queue.queue).(*timedQueueEntry).value
  77. queue.access.Unlock()
  78. queue.removed <- firstEntryValue
  79. }
  80. }