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  }