github.com/kubewharf/katalyst-core@v0.5.3/pkg/scheduler/cache/cache.go (about) 1 /* 2 Copyright 2022 The Katalyst Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package cache 18 19 import ( 20 "sync" 21 22 "github.com/pkg/errors" 23 v1 "k8s.io/api/core/v1" 24 "k8s.io/klog/v2" 25 "k8s.io/kubernetes/pkg/scheduler/framework" 26 27 apis "github.com/kubewharf/katalyst-api/pkg/apis/node/v1alpha1" 28 ) 29 30 type extendedCache struct { 31 // This mutex guards all fields within this extendedCache struct. 32 mu sync.RWMutex 33 nodes map[string]*NodeInfo 34 } 35 36 var cache *extendedCache 37 38 func init() { 39 cache = &extendedCache{ 40 nodes: make(map[string]*NodeInfo), 41 } 42 } 43 44 func GetCache() *extendedCache { 45 return cache 46 } 47 48 func (cache *extendedCache) AddPod(pod *v1.Pod) error { 49 key, err := framework.GetPodKey(pod) 50 if err != nil { 51 return err 52 } 53 54 cache.mu.Lock() 55 defer cache.mu.Unlock() 56 57 n, ok := cache.nodes[pod.Spec.NodeName] 58 if !ok { 59 n = NewNodeInfo() 60 cache.nodes[pod.Spec.NodeName] = n 61 } 62 n.AddPod(key, pod) 63 64 return nil 65 } 66 67 func (cache *extendedCache) RemovePod(pod *v1.Pod) error { 68 key, err := framework.GetPodKey(pod) 69 if err != nil { 70 return err 71 } 72 73 cache.mu.Lock() 74 defer cache.mu.Unlock() 75 76 n, ok := cache.nodes[pod.Spec.NodeName] 77 if !ok { 78 klog.ErrorS(nil, "Node not found when trying to remove pod", "node", klog.KRef("", pod.Spec.NodeName), "pod", klog.KObj(pod)) 79 } else { 80 n.RemovePod(key, pod) 81 } 82 83 return nil 84 } 85 86 func (cache *extendedCache) AddOrUpdateCNR(cnr *apis.CustomNodeResource) { 87 cache.mu.Lock() 88 defer cache.mu.Unlock() 89 90 n, ok := cache.nodes[cnr.Name] 91 if !ok { 92 n = NewNodeInfo() 93 cache.nodes[cnr.Name] = n 94 } 95 n.UpdateNodeInfo(cnr) 96 } 97 98 // RemoveCNR removes a Node from the cache. 99 func (cache *extendedCache) RemoveCNR(cnr *apis.CustomNodeResource) { 100 cache.mu.Lock() 101 defer cache.mu.Unlock() 102 103 delete(cache.nodes, cnr.Name) 104 } 105 106 // GetNodeInfo returns the NodeInfo. 107 func (cache *extendedCache) GetNodeInfo(name string) (*NodeInfo, error) { 108 cache.mu.RLock() 109 defer cache.mu.RUnlock() 110 111 nodeInfo, ok := cache.nodes[name] 112 if !ok { 113 return nil, errors.New("node not found in the extended cache") 114 } 115 116 return nodeInfo, nil 117 } 118 119 func (cache *extendedCache) ReserveNodeResource(nodeName string, pod *v1.Pod) { 120 cache.mu.Lock() 121 defer cache.mu.Unlock() 122 123 nodeInfo, ok := cache.nodes[nodeName] 124 if !ok { 125 nodeInfo = NewNodeInfo() 126 } 127 128 nodeInfo.AddAssumedPod(pod) 129 cache.nodes[nodeName] = nodeInfo 130 } 131 132 func (cache *extendedCache) UnreserveNodeResource(nodeName string, pod *v1.Pod) { 133 cache.mu.Lock() 134 defer cache.mu.Unlock() 135 136 nodeInfo, ok := cache.nodes[nodeName] 137 if !ok { 138 klog.Warningf("UnreserveNodeResource fail, node %v not exist in extendedCache", nodeName) 139 return 140 } 141 142 nodeInfo.DeleteAssumedPod(pod) 143 } 144 145 // GetNodeResourceTopology assumedPodResource will be added to nodeResourceTopology 146 func (cache *extendedCache) GetNodeResourceTopology(nodeName string, filterFn podFilter) *ResourceTopology { 147 cache.mu.RLock() 148 defer cache.mu.RUnlock() 149 150 nodeInfo, ok := cache.nodes[nodeName] 151 if !ok { 152 return nil 153 } 154 return nodeInfo.GetResourceTopologyCopy(filterFn) 155 }