1
// MIT License
2
 
3
// Copyright (c) 2022 Enterprize Software
4
// https://github.com/enterprizesoftware/rate-counter
5
 
6
// Permission is hereby granted, free of charge, to any person obtaining a copy
7
// of this software and associated documentation files (the "Software"), to deal
8
// in the Software without restriction, including without limitation the rights
9
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
10
// copies of the Software, and to permit persons to whom the Software is
11
// furnished to do so, subject to the following conditions:
12
 
13
// The above copyright notice and this permission notice shall be included in all
14
// copies or substantial portions of the Software.
15
 
16
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
19
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
21
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
22
// SOFTWARE.
23
package ratecounter
24
 
25
import (
26
	"math"
27
	"sync"
28
	"time"
29
 
30
	"go.uber.org/atomic"
31
)
32
 
33
type Rate struct {
34
	lock        sync.RWMutex
35
	sampleCount *atomic.Uint64
36
	total       *atomic.Uint64
37
	count       *atomic.Uint64
38
	opened      *atomic.Time
39
	samples     []*sample
40
	interval    time.Duration
41
	observe     time.Duration
42
}
43
 
44
type sample struct {
45
	stored   time.Time
46
	count    uint64
47
	duration time.Duration
48
}
49
 
50
func New(interval time.Duration, observe time.Duration) *Rate {
51
	num := uint64(math.Ceil(float64(observe) / float64(interval)))
52
	samples := make([]*sample, num)
53
	for i := range samples {
54
		samples[i] = &sample{}
55
	}
56
	return &Rate{
57
		interval:    interval,
58
		observe:     observe,
59
		samples:     samples,
60
		sampleCount: atomic.NewUint64(0),
61
		total:       atomic.NewUint64(0),
62
		count:       atomic.NewUint64(0),
63
		opened:      atomic.NewTime(time.Now()),
64
	}
65
}
66
 
67
func (r *Rate) RatePerInterval() float64 {
68
	return float64(r.interval) * r.rate()
69
}
70
 
71
func (r *Rate) RatePer(interval time.Duration) float64 {
72
	return float64(interval) * r.rate()
73
}
74
 
75
func (r *Rate) rate() float64 {
76
	r.lock.RLock()
77
	defer r.lock.RUnlock()
78
 
79
	var c uint64
80
	var d time.Duration
81
	var num = min(len(r.samples), r.sampleCount.Load())
82
	var now = time.Now()
83
 
84
	for i := range num {
85
		if now.Sub(r.samples[i].stored) <= r.observe {
86
			c += r.samples[i].count
87
			d += r.samples[i].duration
88
		}
89
	}
90
 
91
	if d == 0 {
92
		return 0
93
	}
94
 
95
	return float64(c) / float64(d.Nanoseconds())
96
}
97
 
98
func (r *Rate) Total() uint64 {
99
	return r.total.Load()
100
}
101
 
102
func (r *Rate) Increment() {
103
	r.IncrementBy(1)
104
}
105
 
106
func (r *Rate) IncrementBy(i int) {
107
	r.count.Add(uint64(i))
108
	r.total.Add(uint64(i))
109
 
110
	r.lock.RLock()
111
	if time.Since(r.opened.Load()) > r.interval {
112
		r.lock.RUnlock()
113
		r.captureSample()
114
	} else {
115
		r.lock.RUnlock()
116
	}
117
}
118
 
119
func (r *Rate) captureSample() {
120
	r.lock.Lock()
121
 
122
	now := time.Now()
123
	sc := r.sampleCount.Load()
124
 
125
	index := int(sc+1) % len(r.samples)
126
	r.samples[index].count = r.count.Load()
127
	r.samples[index].duration = now.Sub(r.opened.Load())
128
	r.samples[index].stored = now
129
 
130
	r.sampleCount.Inc()
131
 
132
	r.opened.Store(now)
133
	r.count.Store(0)
134
 
135
	r.lock.Unlock()
136
}
137
 
138
func min(x int, y uint64) int {
139
	if uint64(x) < y {
140
		return x
141
	}
142
	return int(y)
143
}