...
1
16
17 package state
18
19 import (
20 "sync"
21
22 "k8s.io/klog/v2"
23 "k8s.io/utils/cpuset"
24 )
25
26 type stateMemory struct {
27 sync.RWMutex
28 assignments ContainerCPUAssignments
29 defaultCPUSet cpuset.CPUSet
30 }
31
32 var _ State = &stateMemory{}
33
34
35 func NewMemoryState() State {
36 klog.InfoS("Initialized new in-memory state store")
37 return &stateMemory{
38 assignments: ContainerCPUAssignments{},
39 defaultCPUSet: cpuset.New(),
40 }
41 }
42
43 func (s *stateMemory) GetCPUSet(podUID string, containerName string) (cpuset.CPUSet, bool) {
44 s.RLock()
45 defer s.RUnlock()
46
47 res, ok := s.assignments[podUID][containerName]
48 return res.Clone(), ok
49 }
50
51 func (s *stateMemory) GetDefaultCPUSet() cpuset.CPUSet {
52 s.RLock()
53 defer s.RUnlock()
54
55 return s.defaultCPUSet.Clone()
56 }
57
58 func (s *stateMemory) GetCPUSetOrDefault(podUID string, containerName string) cpuset.CPUSet {
59 if res, ok := s.GetCPUSet(podUID, containerName); ok {
60 return res
61 }
62 return s.GetDefaultCPUSet()
63 }
64
65 func (s *stateMemory) GetCPUAssignments() ContainerCPUAssignments {
66 s.RLock()
67 defer s.RUnlock()
68 return s.assignments.Clone()
69 }
70
71 func (s *stateMemory) SetCPUSet(podUID string, containerName string, cset cpuset.CPUSet) {
72 s.Lock()
73 defer s.Unlock()
74
75 if _, ok := s.assignments[podUID]; !ok {
76 s.assignments[podUID] = make(map[string]cpuset.CPUSet)
77 }
78
79 s.assignments[podUID][containerName] = cset
80 klog.InfoS("Updated desired CPUSet", "podUID", podUID, "containerName", containerName, "cpuSet", cset)
81 }
82
83 func (s *stateMemory) SetDefaultCPUSet(cset cpuset.CPUSet) {
84 s.Lock()
85 defer s.Unlock()
86
87 s.defaultCPUSet = cset
88 klog.InfoS("Updated default CPUSet", "cpuSet", cset)
89 }
90
91 func (s *stateMemory) SetCPUAssignments(a ContainerCPUAssignments) {
92 s.Lock()
93 defer s.Unlock()
94
95 s.assignments = a.Clone()
96 klog.InfoS("Updated CPUSet assignments", "assignments", a)
97 }
98
99 func (s *stateMemory) Delete(podUID string, containerName string) {
100 s.Lock()
101 defer s.Unlock()
102
103 delete(s.assignments[podUID], containerName)
104 if len(s.assignments[podUID]) == 0 {
105 delete(s.assignments, podUID)
106 }
107 klog.V(2).InfoS("Deleted CPUSet assignment", "podUID", podUID, "containerName", containerName)
108 }
109
110 func (s *stateMemory) ClearState() {
111 s.Lock()
112 defer s.Unlock()
113
114 s.defaultCPUSet = cpuset.CPUSet{}
115 s.assignments = make(ContainerCPUAssignments)
116 klog.V(2).InfoS("Cleared state")
117 }
118
View as plain text