forked from gravitational/teleport
-
Notifications
You must be signed in to change notification settings - Fork 0
/
servicecounter.go
71 lines (61 loc) · 2.19 KB
/
servicecounter.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
/*
* Teleport
* Copyright (C) 2023 Gravitational, Inc.
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package inventory
import (
"sync"
"sync/atomic"
"github.com/gravitational/teleport/api/types"
)
// serviceCounter will count services seen in the inventory.
type serviceCounter struct {
countMap sync.Map
}
// counts returns the count of each service seen in the counter.
func (s *serviceCounter) counts() map[types.SystemRole]uint64 {
counts := map[types.SystemRole]uint64{}
s.countMap.Range(func(key, value any) bool {
counts[key.(types.SystemRole)] = value.(*atomic.Uint64).Load()
return true
})
return counts
}
// increment will increment the counter for a service.
func (s *serviceCounter) increment(service types.SystemRole) {
s.load(service).Add(1)
}
// decrement will decrement the counter for a service.
func (s *serviceCounter) decrement(service types.SystemRole) {
// refer to the docs for atomic.AddUint64 for why this works as a decrement.
s.load(service).Add(^uint64(0))
}
// get will return the value of a counter for a service.
func (s *serviceCounter) get(service types.SystemRole) uint64 {
if result, ok := s.countMap.Load(service); ok {
return result.(*atomic.Uint64).Load()
}
return 0
}
// load will load the underlying atomic value in the sync map. This should
// only be used within the service counter.
func (s *serviceCounter) load(service types.SystemRole) *atomic.Uint64 {
if result, ok := s.countMap.Load(service); ok {
return result.(*atomic.Uint64)
}
result, _ := s.countMap.LoadOrStore(service, &atomic.Uint64{})
return result.(*atomic.Uint64)
}