-
Notifications
You must be signed in to change notification settings - Fork 10
/
rate.go
106 lines (84 loc) · 2.15 KB
/
rate.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
// Copyright (c) 2022-present Mattermost, Inc. All Rights Reserved.
// See LICENSE.txt for license information.
package rtc
import (
"fmt"
"math"
"sync"
"time"
)
type sample struct {
ts time.Time
size int
}
type RateMonitor struct {
samples []sample
samplesPtr int
samplingSize time.Duration
filled bool
now func() time.Time
mut sync.RWMutex
}
func NewRateMonitor(samplingSize time.Duration, now func() time.Time) (*RateMonitor, error) {
if samplingSize <= 0 {
return nil, fmt.Errorf("invalid sampling size")
}
if now == nil {
now = time.Now
}
return &RateMonitor{
now: now,
samplingSize: samplingSize,
samples: make([]sample, 0),
}, nil
}
func (m *RateMonitor) PushSample(size int) {
m.mut.Lock()
defer m.mut.Unlock()
// Filling up to double the sampling size to make sure we have enough samples
// to calculate the desired duration since at the beginning it's likely we get
// a burst of packets.
if !m.filled {
m.samples = append(m.samples, sample{ts: m.now(), size: size})
m.samplesPtr++
if m.getSamplesDuration() >= m.samplingSize*2 {
m.filled = true
}
return
}
m.samples[m.samplesPtr%len(m.samples)] = sample{ts: m.now(), size: size}
m.samplesPtr++
}
func (m *RateMonitor) GetSamplesDuration() time.Duration {
m.mut.RLock()
defer m.mut.RUnlock()
return m.getSamplesDuration()
}
func (m *RateMonitor) getSamplesDuration() time.Duration {
if len(m.samples) == 0 {
return 0
}
lastTS := m.samples[(m.samplesPtr-1)%len(m.samples)].ts
firstTS := m.samples[m.samplesPtr%len(m.samples)].ts
return lastTS.Sub(firstTS)
}
func (m *RateMonitor) GetRate() (int, time.Duration) {
m.mut.RLock()
defer m.mut.RUnlock()
if !m.filled {
return -1, 0
}
now := m.now()
var totalBytes int
var samplesDuration time.Duration
for i := m.samplesPtr - 1; i >= m.samplesPtr-len(m.samples); i-- {
sample := m.samples[i%len(m.samples)]
samplesDuration = now.Sub(sample.ts)
totalBytes += sample.size
if samplesDuration >= m.samplingSize {
break
}
}
bitsPerSec := math.Round((float64(totalBytes) / samplesDuration.Seconds()) * 8)
return int(bitsPerSec), samplesDuration
}