github.com/lingyao2333/mo-zero@v1.4.1/core/collection/rollingwindow.go (about)

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