rollingwindow.go 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  1. package collection
  2. import (
  3. "sync"
  4. "time"
  5. "github.com/tal-tech/go-zero/core/timex"
  6. )
  7. type (
  8. // RollingWindowOption let callers customize the RollingWindow.
  9. RollingWindowOption func(rollingWindow *RollingWindow)
  10. // RollingWindow defines a rolling window to calculate the events in buckets with time interval.
  11. RollingWindow struct {
  12. lock sync.RWMutex
  13. size int
  14. win *window
  15. interval time.Duration
  16. offset int
  17. ignoreCurrent bool
  18. lastTime time.Duration // start time of the last bucket
  19. }
  20. )
  21. // NewRollingWindow returns a RollingWindow that with size buckets and time interval,
  22. // use opts to customize the RollingWindow.
  23. func NewRollingWindow(size int, interval time.Duration, opts ...RollingWindowOption) *RollingWindow {
  24. if size < 1 {
  25. panic("size must be greater than 0")
  26. }
  27. w := &RollingWindow{
  28. size: size,
  29. win: newWindow(size),
  30. interval: interval,
  31. lastTime: timex.Now(),
  32. }
  33. for _, opt := range opts {
  34. opt(w)
  35. }
  36. return w
  37. }
  38. // Add adds value to current bucket.
  39. func (rw *RollingWindow) Add(v float64) {
  40. rw.lock.Lock()
  41. defer rw.lock.Unlock()
  42. rw.updateOffset()
  43. rw.win.add(rw.offset, v)
  44. }
  45. // Reduce runs fn on all buckets, ignore current bucket if ignoreCurrent was set.
  46. func (rw *RollingWindow) Reduce(fn func(b *Bucket)) {
  47. rw.lock.RLock()
  48. defer rw.lock.RUnlock()
  49. var diff int
  50. span := rw.span()
  51. // ignore current bucket, because of partial data
  52. if span == 0 && rw.ignoreCurrent {
  53. diff = rw.size - 1
  54. } else {
  55. diff = rw.size - span
  56. }
  57. if diff > 0 {
  58. offset := (rw.offset + span + 1) % rw.size
  59. rw.win.reduce(offset, diff, fn)
  60. }
  61. }
  62. func (rw *RollingWindow) span() int {
  63. offset := int(timex.Since(rw.lastTime) / rw.interval)
  64. if 0 <= offset && offset < rw.size {
  65. return offset
  66. }
  67. return rw.size
  68. }
  69. func (rw *RollingWindow) updateOffset() {
  70. span := rw.span()
  71. if span <= 0 {
  72. return
  73. }
  74. offset := rw.offset
  75. // reset expired buckets
  76. for i := 0; i < span; i++ {
  77. rw.win.resetBucket((offset + i + 1) % rw.size)
  78. }
  79. rw.offset = (offset + span) % rw.size
  80. now := timex.Now()
  81. // align to interval time boundary
  82. rw.lastTime = now - (now-rw.lastTime)%rw.interval
  83. }
  84. // Bucket defines the bucket that holds sum and num of additions.
  85. type Bucket struct {
  86. Sum float64
  87. Count int64
  88. }
  89. func (b *Bucket) add(v float64) {
  90. b.Sum += v
  91. b.Count++
  92. }
  93. func (b *Bucket) reset() {
  94. b.Sum = 0
  95. b.Count = 0
  96. }
  97. type window struct {
  98. buckets []*Bucket
  99. size int
  100. }
  101. func newWindow(size int) *window {
  102. buckets := make([]*Bucket, size)
  103. for i := 0; i < size; i++ {
  104. buckets[i] = new(Bucket)
  105. }
  106. return &window{
  107. buckets: buckets,
  108. size: size,
  109. }
  110. }
  111. func (w *window) add(offset int, v float64) {
  112. w.buckets[offset%w.size].add(v)
  113. }
  114. func (w *window) reduce(start, count int, fn func(b *Bucket)) {
  115. for i := 0; i < count; i++ {
  116. fn(w.buckets[(start+i)%w.size])
  117. }
  118. }
  119. func (w *window) resetBucket(offset int) {
  120. w.buckets[offset%w.size].reset()
  121. }
  122. // IgnoreCurrentBucket lets the Reduce call ignore current bucket.
  123. func IgnoreCurrentBucket() RollingWindowOption {
  124. return func(w *RollingWindow) {
  125. w.ignoreCurrent = true
  126. }
  127. }