github.com/kubewharf/katalyst-core@v0.5.3/pkg/util/native/pod_finder.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 native 18 19 import ( 20 "fmt" 21 22 core "k8s.io/api/core/v1" 23 "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" 24 "k8s.io/apimachinery/pkg/labels" 25 "k8s.io/apimachinery/pkg/runtime" 26 "k8s.io/apimachinery/pkg/types" 27 v1 "k8s.io/client-go/informers/core/v1" 28 corelisters "k8s.io/client-go/listers/core/v1" 29 "k8s.io/client-go/tools/cache" 30 "k8s.io/klog/v2" 31 ) 32 33 type PodLabelIndexer string 34 35 const nodeNameKeyIndex = "spec.nodeName" 36 37 // IndexFunc is used to construct informer index for labels in pod 38 func (p PodLabelIndexer) IndexFunc(obj interface{}) ([]string, error) { 39 pod, ok := obj.(*core.Pod) 40 if !ok { 41 return nil, fmt.Errorf("failed to reflect a obj to pod") 42 } 43 44 key := string(p) 45 value := pod.Labels[key] 46 47 return []string{value}, nil 48 } 49 50 // getPodListForWorkloadWithIndex is used to get pod list that belongs to the given workload 51 func getPodListForWorkloadWithIndex(selector labels.Selector, podIndexer cache.Indexer, podIndexerKey []string) ([]*core.Pod, error) { 52 podMap := make(map[types.UID]*core.Pod) 53 for _, key := range podIndexerKey { 54 value, ok := selector.RequiresExactMatch(key) 55 if !ok { 56 klog.Warningf("key %v without value for selector %v", key, selector.String()) 57 continue 58 } 59 60 objs, err := podIndexer.ByIndex(key, value) 61 if err != nil { 62 klog.Warningf("pods for index %s/%s err: %v", key, value, err) 63 continue 64 } 65 66 for _, obj := range objs { 67 if pod, exist := obj.(*core.Pod); !exist { 68 continue 69 } else if _, exist = podMap[pod.UID]; !exist && selector.Matches(labels.Set(pod.Labels)) { 70 podMap[pod.UID] = pod 71 } 72 } 73 } 74 75 if len(podMap) == 0 { 76 return nil, fmt.Errorf("no pod matched with selector %v", selector.String()) 77 } 78 79 var pods []*core.Pod 80 for _, pod := range podMap { 81 pods = append(pods, pod) 82 } 83 return pods, nil 84 } 85 86 // GetPodListForWorkload returns pod list that belong to the given workload 87 // we will use label selector to find pods, and this may require that the given 88 // workload is limited to several selected objects. 89 func GetPodListForWorkload(workloadObj runtime.Object, podIndexer cache.Indexer, labelKeyList []string, podLister corelisters.PodLister) ([]*core.Pod, error) { 90 workload := workloadObj.(*unstructured.Unstructured) 91 92 selector, err := GetUnstructuredSelector(workload) 93 if err != nil || selector == nil { 94 klog.Errorf("failed to get workload selector %v: %v", workloadObj, err) 95 return nil, err 96 } 97 98 if podIndexer != nil { 99 if podList, err := getPodListForWorkloadWithIndex(selector, podIndexer, labelKeyList); err == nil { 100 klog.V(3).Infof("get pods with index successfully, selector: %v", selector) 101 return podList, nil 102 } 103 } 104 105 pods, err := podLister.List(selector) 106 if err != nil { 107 return nil, err 108 } 109 return pods, nil 110 } 111 112 // GetPodsAssignedToNode returns pods that belong to this node by indexer 113 func GetPodsAssignedToNode(nodeName string, podIndexer cache.Indexer) ([]*core.Pod, error) { 114 objs, err := podIndexer.ByIndex(nodeNameKeyIndex, nodeName) 115 if err != nil { 116 return nil, err 117 } 118 119 pods := make([]*core.Pod, 0, len(objs)) 120 for _, obj := range objs { 121 pod, ok := obj.(*core.Pod) 122 if !ok { 123 continue 124 } 125 pods = append(pods, pod) 126 } 127 return pods, nil 128 } 129 130 // AddNodeNameIndexerForPod add node name index for pod informer 131 func AddNodeNameIndexerForPod(podInformer v1.PodInformer) error { 132 if _, ok := podInformer.Informer().GetIndexer().GetIndexers()[nodeNameKeyIndex]; !ok { 133 return podInformer.Informer().GetIndexer().AddIndexers(cache.Indexers{ 134 nodeNameKeyIndex: func(obj interface{}) ([]string, error) { 135 pod, ok := obj.(*core.Pod) 136 if !ok { 137 return []string{}, nil 138 } 139 if len(pod.Spec.NodeName) == 0 { 140 return []string{}, nil 141 } 142 return []string{pod.Spec.NodeName}, nil 143 }, 144 }) 145 } 146 return nil 147 }