-
Notifications
You must be signed in to change notification settings - Fork 117
/
priority_class_cache.go
129 lines (110 loc) · 3.85 KB
/
priority_class_cache.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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package admission
import (
schedulingv1 "k8s.io/api/scheduling/v1"
informersv1 "k8s.io/client-go/informers/scheduling/v1"
"k8s.io/client-go/tools/cache"
"github.com/apache/yunikorn-k8shim/pkg/common/constants"
"github.com/apache/yunikorn-k8shim/pkg/common/utils"
"github.com/apache/yunikorn-k8shim/pkg/locking"
"github.com/apache/yunikorn-k8shim/pkg/log"
)
type PriorityClassCache struct {
priorityClasses map[string]bool
locking.RWMutex
}
// NewPriorityClassCache creates a new cache and registers the handler for the cache with the Informer.
func NewPriorityClassCache(priorityClasses informersv1.PriorityClassInformer) *PriorityClassCache {
pcc := &PriorityClassCache{
priorityClasses: make(map[string]bool),
}
if priorityClasses != nil {
priorityClasses.Informer().AddEventHandler(&priorityClassUpdateHandler{cache: pcc})
}
return pcc
}
// isPreemptSelfAllowed returns the preemption value. Only returns false if configured.
func (pcc *PriorityClassCache) isPreemptSelfAllowed(priorityClassName string) bool {
pcc.RLock()
defer pcc.RUnlock()
value, ok := pcc.priorityClasses[priorityClassName]
if !ok {
return true
}
return value
}
// priorityClassExists for test only to see if the PriorityClass has been added to the cache or not.
func (pcc *PriorityClassCache) priorityClassExists(priorityClassName string) bool {
pcc.RLock()
defer pcc.RUnlock()
_, ok := pcc.priorityClasses[priorityClassName]
return ok
}
// priorityClassUpdateHandler implements the K8s ResourceEventHandler interface for PriorityClass.
type priorityClassUpdateHandler struct {
cache *PriorityClassCache
}
// OnAdd adds or replaces the priority class entry in the cache.
// The cached value is only the resulting value of the annotation, not the whole PriorityClass object.
// An empty string for the Name is technically possible but should not occur.
func (h *priorityClassUpdateHandler) OnAdd(obj interface{}, _ bool) {
pc := utils.Convert2PriorityClass(obj)
if pc == nil {
return
}
b := getAnnotationBoolean(pc.Annotations, constants.AnnotationAllowPreemption)
h.cache.Lock()
defer h.cache.Unlock()
h.cache.priorityClasses[pc.Name] = b
}
// OnUpdate calls OnAdd for processing the PriorityClass cache update.
func (h *priorityClassUpdateHandler) OnUpdate(_, newObj interface{}) {
h.OnAdd(newObj, false)
}
// OnDelete removes the PriorityClass from the cache.
func (h *priorityClassUpdateHandler) OnDelete(obj interface{}) {
var pc *schedulingv1.PriorityClass
switch t := obj.(type) {
case *schedulingv1.PriorityClass:
pc = t
case cache.DeletedFinalStateUnknown:
pc = utils.Convert2PriorityClass(obj)
default:
log.Log(log.Admission).Warn("unable to convert to PriorityClass")
return
}
if pc == nil {
return
}
h.cache.Lock()
defer h.cache.Unlock()
delete(h.cache.priorityClasses, pc.Name)
}
// getAnnotationBoolean retrieves the value from the map and returns it.
// Defaults to true if the name does not exist.
func getAnnotationBoolean(m map[string]string, name string) bool {
strVal, ok := m[name]
if !ok {
return true
}
switch strVal {
case constants.False:
return false
default:
return true
}
}