-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Expand file tree
/
Copy pathusage.go
More file actions
76 lines (64 loc) · 1.93 KB
/
usage.go
File metadata and controls
76 lines (64 loc) · 1.93 KB
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
package coderdtest
import (
"context"
"sync"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/usage"
"github.com/coder/coder/v2/coderd/usage/usagetypes"
)
var _ usage.Inserter = (*UsageInserter)(nil)
type UsageInserter struct {
sync.Mutex
discreteEvents []usagetypes.DiscreteEvent
heartbeatEvents []usagetypes.HeartbeatEvent
seenHeartbeats map[string]struct{}
}
func NewUsageInserter() *UsageInserter {
return &UsageInserter{
discreteEvents: []usagetypes.DiscreteEvent{},
seenHeartbeats: map[string]struct{}{},
heartbeatEvents: []usagetypes.HeartbeatEvent{},
}
}
func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
u.Lock()
defer u.Unlock()
u.discreteEvents = append(u.discreteEvents, event)
return nil
}
func (u *UsageInserter) InsertHeartbeatUsageEvent(_ context.Context, _ database.Store, id string, event usagetypes.HeartbeatEvent) error {
u.Lock()
defer u.Unlock()
if _, seen := u.seenHeartbeats[id]; seen {
return nil
}
u.seenHeartbeats[id] = struct{}{}
u.heartbeatEvents = append(u.heartbeatEvents, event)
return nil
}
func (u *UsageInserter) GetHeartbeatEvents() []usagetypes.HeartbeatEvent {
u.Lock()
defer u.Unlock()
eventsCopy := make([]usagetypes.HeartbeatEvent, len(u.heartbeatEvents))
copy(eventsCopy, u.heartbeatEvents)
return eventsCopy
}
func (u *UsageInserter) GetDiscreteEvents() []usagetypes.DiscreteEvent {
u.Lock()
defer u.Unlock()
eventsCopy := make([]usagetypes.DiscreteEvent, len(u.discreteEvents))
copy(eventsCopy, u.discreteEvents)
return eventsCopy
}
func (u *UsageInserter) TotalEventCount() int {
u.Lock()
defer u.Unlock()
return len(u.discreteEvents) + len(u.heartbeatEvents)
}
func (u *UsageInserter) Reset() {
u.Lock()
defer u.Unlock()
u.seenHeartbeats = map[string]struct{}{}
u.discreteEvents = []usagetypes.DiscreteEvent{}
u.heartbeatEvents = []usagetypes.HeartbeatEvent{}
}