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  }