github.com/aaabigfish/gopkg@v1.1.0/stat/counter/rolling.go (about)

     1  package counter
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  )
     7  
     8  type bucket struct {
     9  	val  int64
    10  	next *bucket
    11  }
    12  
    13  func (b *bucket) Add(val int64) {
    14  	b.val += val
    15  }
    16  
    17  func (b *bucket) Value() int64 {
    18  	return b.val
    19  }
    20  
    21  func (b *bucket) Reset() {
    22  	b.val = 0
    23  }
    24  
    25  var _ Counter = new(rollingCounter)
    26  
    27  type rollingCounter struct {
    28  	mu         sync.RWMutex
    29  	buckets    []bucket
    30  	bucketTime int64
    31  	lastAccess int64
    32  	cur        *bucket
    33  }
    34  
    35  // NewRolling creates a new window. windowTime is the time covering the entire
    36  // window. windowBuckets is the number of buckets the window is divided into.
    37  // An example: a 10 second window with 10 buckets will have 10 buckets covering
    38  // 1 second each.
    39  func NewRolling(window time.Duration, winBucket int) Counter {
    40  	buckets := make([]bucket, winBucket)
    41  	bucket := &buckets[0]
    42  	for i := 1; i < winBucket; i++ {
    43  		bucket.next = &buckets[i]
    44  		bucket = bucket.next
    45  	}
    46  	bucket.next = &buckets[0]
    47  	bucketTime := time.Duration(window.Nanoseconds() / int64(winBucket))
    48  	return &rollingCounter{
    49  		cur:        &buckets[0],
    50  		buckets:    buckets,
    51  		bucketTime: int64(bucketTime),
    52  		lastAccess: time.Now().UnixNano(),
    53  	}
    54  }
    55  
    56  // Add increments the counter by value and return new value.
    57  func (r *rollingCounter) Add(val int64) {
    58  	r.mu.Lock()
    59  	r.lastBucket().Add(val)
    60  	r.mu.Unlock()
    61  }
    62  
    63  // Value get the counter value.
    64  func (r *rollingCounter) Value() (sum int64) {
    65  	now := time.Now().UnixNano()
    66  	r.mu.RLock()
    67  	b := r.cur
    68  	i := r.elapsed(now)
    69  	for j := 0; j < len(r.buckets); j++ {
    70  		// skip all future reset bucket.
    71  		if i > 0 {
    72  			i--
    73  		} else {
    74  			sum += b.Value()
    75  		}
    76  		b = b.next
    77  	}
    78  	r.mu.RUnlock()
    79  	return
    80  }
    81  
    82  //  Reset reset the counter.
    83  func (r *rollingCounter) Reset() {
    84  	r.mu.Lock()
    85  	for i := range r.buckets {
    86  		r.buckets[i].Reset()
    87  	}
    88  	r.mu.Unlock()
    89  }
    90  
    91  func (r *rollingCounter) elapsed(now int64) (i int) {
    92  	var e int64
    93  	if e = now - r.lastAccess; e <= r.bucketTime {
    94  		return
    95  	}
    96  	if i = int(e / r.bucketTime); i > len(r.buckets) {
    97  		i = len(r.buckets)
    98  	}
    99  	return
   100  }
   101  
   102  func (r *rollingCounter) lastBucket() (b *bucket) {
   103  	now := time.Now().UnixNano()
   104  	b = r.cur
   105  	// reset the buckets between now and number of buckets ago. If
   106  	// that is more that the existing buckets, reset all.
   107  	if i := r.elapsed(now); i > 0 {
   108  		r.lastAccess = now
   109  		for ; i > 0; i-- {
   110  			// replace the next used bucket.
   111  			b = b.next
   112  			b.Reset()
   113  		}
   114  	}
   115  	r.cur = b
   116  	return
   117  }