selector.go

Documentation: github.com/chaos-mesh/chaos-mesh/pkg/selector

     1  // Copyright 2019 Chaos Mesh Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package selector
    15  
    16  import (
    17  	"context"
    18  	"errors"
    19  	"fmt"
    20  	"math"
    21  	"math/rand"
    22  	"strconv"
    23  	"strings"
    24  
    25  	"github.com/chaos-mesh/chaos-mesh/api/v1alpha1"
    26  	"github.com/chaos-mesh/chaos-mesh/pkg/label"
    27  	"github.com/chaos-mesh/chaos-mesh/pkg/mock"
    28  
    29  	ctrl "sigs.k8s.io/controller-runtime"
    30  	"sigs.k8s.io/controller-runtime/pkg/client"
    31  
    32  	v1 "k8s.io/api/core/v1"
    33  	apierrors "k8s.io/apimachinery/pkg/api/errors"
    34  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    35  	"k8s.io/apimachinery/pkg/fields"
    36  	"k8s.io/apimachinery/pkg/labels"
    37  	"k8s.io/apimachinery/pkg/selection"
    38  	"k8s.io/apimachinery/pkg/types"
    39  )
    40  
    41  var log = ctrl.Log.WithName("selector")
    42  
    43  const injectAnnotationKey = "chaos-mesh.org/inject"
    44  
    45  type SelectSpec interface {
    46  	GetSelector() v1alpha1.SelectorSpec
    47  	GetMode() v1alpha1.PodMode
    48  	GetValue() string
    49  }
    50  
    51  // SelectAndFilterPods returns the list of pods that filtered by selector and PodMode
    52  func SelectAndFilterPods(ctx context.Context, c client.Client, r client.Reader, spec SelectSpec, clusterScoped bool, targetNamespace string, enableFilterNamespace bool) ([]v1.Pod, error) {
    53  	if pods := mock.On("MockSelectAndFilterPods"); pods != nil {
    54  		return pods.(func() []v1.Pod)(), nil
    55  	}
    56  	if err := mock.On("MockSelectedAndFilterPodsError"); err != nil {
    57  		return nil, err.(error)
    58  	}
    59  
    60  	selector := spec.GetSelector()
    61  	mode := spec.GetMode()
    62  	value := spec.GetValue()
    63  
    64  	pods, err := SelectPods(ctx, c, r, selector, clusterScoped, targetNamespace, enableFilterNamespace)
    65  	if err != nil {
    66  		return nil, err
    67  	}
    68  
    69  	if len(pods) == 0 {
    70  		err = errors.New("no pod is selected")
    71  		return nil, err
    72  	}
    73  
    74  	filteredPod, err := filterPodsByMode(pods, mode, value)
    75  	if err != nil {
    76  		return nil, err
    77  	}
    78  
    79  	return filteredPod, nil
    80  }
    81  
    82  //revive:disable:flag-parameter
    83  
    84  // SelectPods returns the list of pods that are available for pod chaos action.
    85  // It returns all pods that match the configured label, annotation and namespace selectors.
    86  // If pods are specifically specified by `selector.Pods`, it just returns the selector.Pods.
    87  func SelectPods(ctx context.Context, c client.Client, r client.Reader, selector v1alpha1.SelectorSpec, clusterScoped bool, targetNamespace string, enableFilterNamespace bool) ([]v1.Pod, error) {
    88  	// TODO: refactor: make different selectors to replace if-else logics
    89  	var pods []v1.Pod
    90  
    91  	// pods are specifically specified
    92  	if len(selector.Pods) > 0 {
    93  		for ns, names := range selector.Pods {
    94  			if !clusterScoped {
    95  				if targetNamespace != ns {
    96  					log.Info("skip namespace because ns is out of scope within namespace scoped mode", "namespace", ns)
    97  					continue
    98  				}
    99  			}
   100  			for _, name := range names {
   101  				var pod v1.Pod
   102  				err := c.Get(ctx, types.NamespacedName{
   103  					Namespace: ns,
   104  					Name:      name,
   105  				}, &pod)
   106  				if err == nil {
   107  					pods = append(pods, pod)
   108  					continue
   109  				}
   110  
   111  				if apierrors.IsNotFound(err) {
   112  					log.Error(err, "Pod is not found", "namespace", ns, "pod name", name)
   113  					continue
   114  				}
   115  
   116  				return nil, err
   117  			}
   118  		}
   119  
   120  		return pods, nil
   121  	}
   122  
   123  	if !clusterScoped {
   124  		if len(selector.Namespaces) > 1 {
   125  			return nil, fmt.Errorf("could NOT use more than 1 namespace selector within namespace scoped mode")
   126  		} else if len(selector.Namespaces) == 1 {
   127  			if selector.Namespaces[0] != targetNamespace {
   128  				return nil, fmt.Errorf("could NOT list pods from out of scoped namespace: %s", selector.Namespaces[0])
   129  			}
   130  		}
   131  	}
   132  
   133  	var listOptions = client.ListOptions{}
   134  	if !clusterScoped {
   135  		listOptions.Namespace = targetNamespace
   136  	}
   137  	if len(selector.LabelSelectors) > 0 || len(selector.ExpressionSelectors) > 0 {
   138  		metav1Ls := &metav1.LabelSelector{
   139  			MatchLabels:      selector.LabelSelectors,
   140  			MatchExpressions: selector.ExpressionSelectors,
   141  		}
   142  		ls, err := metav1.LabelSelectorAsSelector(metav1Ls)
   143  		if err != nil {
   144  			return nil, err
   145  		}
   146  		listOptions.LabelSelector = ls
   147  	}
   148  
   149  	listFunc := c.List
   150  
   151  	if len(selector.FieldSelectors) > 0 {
   152  		listOptions.FieldSelector = fields.SelectorFromSet(selector.FieldSelectors)
   153  
   154  		// Since FieldSelectors need to implement index creation, Reader.List is used to get the pod list.
   155  		// Otherwise, just call Client.List directly, which can be obtained through cache.
   156  		if r != nil {
   157  			listFunc = r.List
   158  		}
   159  	}
   160  
   161  	var podList v1.PodList
   162  	if len(selector.Namespaces) > 0 {
   163  		for _, namespace := range selector.Namespaces {
   164  			listOptions.Namespace = namespace
   165  
   166  			if err := listFunc(ctx, &podList, &listOptions); err != nil {
   167  				return nil, err
   168  			}
   169  
   170  			pods = append(pods, podList.Items...)
   171  		}
   172  	} else {
   173  		if err := listFunc(ctx, &podList, &listOptions); err != nil {
   174  			return nil, err
   175  		}
   176  
   177  		pods = append(pods, podList.Items...)
   178  	}
   179  
   180  	var (
   181  		nodes           []v1.Node
   182  		nodeList        v1.NodeList
   183  		nodeListOptions = client.ListOptions{}
   184  	)
   185  	// if both setting Nodes and NodeSelectors, the node list will be combined.
   186  	if len(selector.Nodes) > 0 || len(selector.NodeSelectors) > 0 {
   187  		if len(selector.Nodes) > 0 {
   188  			for _, nodename := range selector.Nodes {
   189  				var node v1.Node
   190  				if err := c.Get(ctx, types.NamespacedName{Name: nodename}, &node); err != nil {
   191  					return nil, err
   192  				}
   193  				nodes = append(nodes, node)
   194  			}
   195  		}
   196  		if len(selector.NodeSelectors) > 0 {
   197  			nodeListOptions.LabelSelector = labels.SelectorFromSet(selector.NodeSelectors)
   198  			if err := c.List(ctx, &nodeList, &nodeListOptions); err != nil {
   199  				return nil, err
   200  			}
   201  			nodes = append(nodes, nodeList.Items...)
   202  		}
   203  		pods = filterPodByNode(pods, nodes)
   204  	}
   205  	if enableFilterNamespace {
   206  		pods = filterByNamespaces(ctx, c, pods)
   207  	}
   208  
   209  	namespaceSelector, err := parseSelector(strings.Join(selector.Namespaces, ","))
   210  	if err != nil {
   211  		return nil, err
   212  	}
   213  	pods, err = filterByNamespaceSelector(pods, namespaceSelector)
   214  	if err != nil {
   215  		return nil, err
   216  	}
   217  
   218  	annotationsSelector, err := parseSelector(label.Label(selector.AnnotationSelectors).String())
   219  	if err != nil {
   220  		return nil, err
   221  	}
   222  	pods = filterByAnnotations(pods, annotationsSelector)
   223  
   224  	phaseSelector, err := parseSelector(strings.Join(selector.PodPhaseSelectors, ","))
   225  	if err != nil {
   226  		return nil, err
   227  	}
   228  	pods, err = filterByPhaseSelector(pods, phaseSelector)
   229  	if err != nil {
   230  		return nil, err
   231  	}
   232  
   233  	return pods, nil
   234  }
   235  
   236  //revive:enable:flag-parameter
   237  
   238  // GetService get k8s service by service name
   239  func GetService(ctx context.Context, c client.Client, namespace, controllerNamespace string, serviceName string) (*v1.Service, error) {
   240  	// use the environment value if namespace is empty
   241  	if len(namespace) == 0 {
   242  		namespace = controllerNamespace
   243  	}
   244  
   245  	service := &v1.Service{}
   246  	err := c.Get(ctx, client.ObjectKey{
   247  		Namespace: namespace,
   248  		Name:      serviceName,
   249  	}, service)
   250  	if err != nil {
   251  		return nil, err
   252  	}
   253  
   254  	return service, nil
   255  }
   256  
   257  // CheckPodMeetSelector checks if this pod meets the selection criteria.
   258  // TODO: support to check fieldsSelector
   259  func CheckPodMeetSelector(pod v1.Pod, selector v1alpha1.SelectorSpec) (bool, error) {
   260  	if len(selector.Pods) > 0 {
   261  		meet := false
   262  		for ns, names := range selector.Pods {
   263  			if pod.Namespace != ns {
   264  				continue
   265  			}
   266  
   267  			for _, name := range names {
   268  				if pod.Name == name {
   269  					meet = true
   270  				}
   271  			}
   272  
   273  			if !meet {
   274  				return false, nil
   275  			}
   276  		}
   277  	}
   278  
   279  	// check pod labels.
   280  	if pod.Labels == nil {
   281  		pod.Labels = make(map[string]string)
   282  	}
   283  
   284  	if selector.LabelSelectors == nil {
   285  		selector.LabelSelectors = make(map[string]string)
   286  	}
   287  
   288  	if len(selector.LabelSelectors) > 0 || len(selector.ExpressionSelectors) > 0 {
   289  		metav1Ls := &metav1.LabelSelector{
   290  			MatchLabels:      selector.LabelSelectors,
   291  			MatchExpressions: selector.ExpressionSelectors,
   292  		}
   293  		ls, err := metav1.LabelSelectorAsSelector(metav1Ls)
   294  		if err != nil {
   295  			return false, err
   296  		}
   297  		podLabels := labels.Set(pod.Labels)
   298  		if len(pod.Labels) == 0 || !ls.Matches(podLabels) {
   299  			return false, nil
   300  		}
   301  	}
   302  
   303  	pods := []v1.Pod{pod}
   304  
   305  	namespaceSelector, err := parseSelector(strings.Join(selector.Namespaces, ","))
   306  	if err != nil {
   307  		return false, err
   308  	}
   309  
   310  	pods, err = filterByNamespaceSelector(pods, namespaceSelector)
   311  	if err != nil {
   312  		return false, err
   313  	}
   314  
   315  	annotationsSelector, err := parseSelector(label.Label(selector.AnnotationSelectors).String())
   316  	if err != nil {
   317  		return false, err
   318  	}
   319  
   320  	pods = filterByAnnotations(pods, annotationsSelector)
   321  
   322  	phaseSelector, err := parseSelector(strings.Join(selector.PodPhaseSelectors, ","))
   323  	if err != nil {
   324  		return false, err
   325  	}
   326  	pods, err = filterByPhaseSelector(pods, phaseSelector)
   327  	if err != nil {
   328  		return false, err
   329  	}
   330  
   331  	if len(pods) > 0 {
   332  		return true, nil
   333  	}
   334  
   335  	return false, nil
   336  }
   337  
   338  func filterPodByNode(pods []v1.Pod, nodes []v1.Node) []v1.Pod {
   339  	if len(nodes) == 0 {
   340  		return nil
   341  	}
   342  	var filteredList []v1.Pod
   343  	for _, pod := range pods {
   344  		for _, node := range nodes {
   345  			if pod.Spec.NodeName == node.Name {
   346  				filteredList = append(filteredList, pod)
   347  			}
   348  		}
   349  	}
   350  	return filteredList
   351  }
   352  
   353  // filterPodsByMode filters pods by mode from pod list
   354  func filterPodsByMode(pods []v1.Pod, mode v1alpha1.PodMode, value string) ([]v1.Pod, error) {
   355  	if len(pods) == 0 {
   356  		return nil, errors.New("cannot generate pods from empty list")
   357  	}
   358  
   359  	switch mode {
   360  	case v1alpha1.OnePodMode:
   361  		index := rand.Intn(len(pods))
   362  		pod := pods[index]
   363  
   364  		return []v1.Pod{pod}, nil
   365  	case v1alpha1.AllPodMode:
   366  		return pods, nil
   367  	case v1alpha1.FixedPodMode:
   368  		num, err := strconv.Atoi(value)
   369  		if err != nil {
   370  			return nil, err
   371  		}
   372  
   373  		if len(pods) < num {
   374  			num = len(pods)
   375  		}
   376  
   377  		if num <= 0 {
   378  			return nil, errors.New("cannot select any pod as value below or equal 0")
   379  		}
   380  
   381  		return getFixedSubListFromPodList(pods, num), nil
   382  	case v1alpha1.FixedPercentPodMode:
   383  		percentage, err := strconv.Atoi(value)
   384  		if err != nil {
   385  			return nil, err
   386  		}
   387  
   388  		if percentage == 0 {
   389  			return nil, errors.New("cannot select any pod as value below or equal 0")
   390  		}
   391  
   392  		if percentage < 0 || percentage > 100 {
   393  			return nil, fmt.Errorf("fixed percentage value of %d is invalid, Must be (0,100]", percentage)
   394  		}
   395  
   396  		num := int(math.Floor(float64(len(pods)) * float64(percentage) / 100))
   397  
   398  		return getFixedSubListFromPodList(pods, num), nil
   399  	case v1alpha1.RandomMaxPercentPodMode:
   400  		maxPercentage, err := strconv.Atoi(value)
   401  		if err != nil {
   402  			return nil, err
   403  		}
   404  
   405  		if maxPercentage == 0 {
   406  			return nil, errors.New("cannot select any pod as value below or equal 0")
   407  		}
   408  
   409  		if maxPercentage < 0 || maxPercentage > 100 {
   410  			return nil, fmt.Errorf("fixed percentage value of %d is invalid, Must be [0-100]", maxPercentage)
   411  		}
   412  
   413  		percentage := rand.Intn(maxPercentage + 1) // + 1 because Intn works with half open interval [0,n) and we want [0,n]
   414  		num := int(math.Floor(float64(len(pods)) * float64(percentage) / 100))
   415  
   416  		return getFixedSubListFromPodList(pods, num), nil
   417  	default:
   418  		return nil, fmt.Errorf("mode %s not supported", mode)
   419  	}
   420  }
   421  
   422  // filterByAnnotations filters a list of pods by a given annotation selector.
   423  func filterByAnnotations(pods []v1.Pod, annotations labels.Selector) []v1.Pod {
   424  	// empty filter returns original list
   425  	if annotations.Empty() {
   426  		return pods
   427  	}
   428  
   429  	var filteredList []v1.Pod
   430  
   431  	for _, pod := range pods {
   432  		// convert the pod's annotations to an equivalent label selector
   433  		selector := labels.Set(pod.Annotations)
   434  
   435  		// include pod if its annotations match the selector
   436  		if annotations.Matches(selector) {
   437  			filteredList = append(filteredList, pod)
   438  		}
   439  	}
   440  
   441  	return filteredList
   442  }
   443  
   444  // filterByPhaseSet filters a list of pods by a given PodPhase selector.
   445  func filterByPhaseSelector(pods []v1.Pod, phases labels.Selector) ([]v1.Pod, error) {
   446  	if phases.Empty() {
   447  		return pods, nil
   448  	}
   449  
   450  	reqs, _ := phases.Requirements()
   451  	var (
   452  		reqIncl []labels.Requirement
   453  		reqExcl []labels.Requirement
   454  
   455  		filteredList []v1.Pod
   456  	)
   457  
   458  	for _, req := range reqs {
   459  		switch req.Operator() {
   460  		case selection.Exists:
   461  			reqIncl = append(reqIncl, req)
   462  		case selection.DoesNotExist:
   463  			reqExcl = append(reqExcl, req)
   464  		default:
   465  			return nil, fmt.Errorf("unsupported operator: %s", req.Operator())
   466  		}
   467  	}
   468  
   469  	for _, pod := range pods {
   470  		included := len(reqIncl) == 0
   471  		selector := labels.Set{string(pod.Status.Phase): ""}
   472  
   473  		// include pod if one including requirement matches
   474  		for _, req := range reqIncl {
   475  			if req.Matches(selector) {
   476  				included = true
   477  				break
   478  			}
   479  		}
   480  
   481  		// exclude pod if it is filtered out by at least one excluding requirement
   482  		for _, req := range reqExcl {
   483  			if !req.Matches(selector) {
   484  				included = false
   485  				break
   486  			}
   487  		}
   488  
   489  		if included {
   490  			filteredList = append(filteredList, pod)
   491  		}
   492  	}
   493  
   494  	return filteredList, nil
   495  }
   496  
   497  func filterByNamespaces(ctx context.Context, c client.Client, pods []v1.Pod) []v1.Pod {
   498  	var filteredList []v1.Pod
   499  
   500  	for _, pod := range pods {
   501  		ok, err := IsAllowedNamespaces(ctx, c, pod.Namespace)
   502  		if err != nil {
   503  			log.Error(err, "fail to check whether this namespace is allowed", "namespace", pod.Namespace)
   504  			continue
   505  		}
   506  
   507  		if ok {
   508  			filteredList = append(filteredList, pod)
   509  		} else {
   510  			log.Info("namespace is not enabled for chaos-mesh", "namespace", pod.Namespace)
   511  		}
   512  	}
   513  	return filteredList
   514  }
   515  
   516  func IsAllowedNamespaces(ctx context.Context, c client.Client, namespace string) (bool, error) {
   517  	ns := &v1.Namespace{}
   518  
   519  	err := c.Get(ctx, types.NamespacedName{Name: namespace}, ns)
   520  	if err != nil {
   521  		return false, err
   522  	}
   523  
   524  	if ns.Annotations[injectAnnotationKey] == "enabled" {
   525  		return true, nil
   526  	}
   527  
   528  	return false, nil
   529  }
   530  
   531  // filterByNamespaceSelector filters a list of pods by a given namespace selector.
   532  func filterByNamespaceSelector(pods []v1.Pod, namespaces labels.Selector) ([]v1.Pod, error) {
   533  	// empty filter returns original list
   534  	if namespaces.Empty() {
   535  		return pods, nil
   536  	}
   537  
   538  	// split requirements into including and excluding groups
   539  	reqs, _ := namespaces.Requirements()
   540  
   541  	var (
   542  		reqIncl []labels.Requirement
   543  		reqExcl []labels.Requirement
   544  
   545  		filteredList []v1.Pod
   546  	)
   547  
   548  	for _, req := range reqs {
   549  		switch req.Operator() {
   550  		case selection.Exists:
   551  			reqIncl = append(reqIncl, req)
   552  		case selection.DoesNotExist:
   553  			reqExcl = append(reqExcl, req)
   554  		default:
   555  			return nil, fmt.Errorf("unsupported operator: %s", req.Operator())
   556  		}
   557  	}
   558  
   559  	for _, pod := range pods {
   560  		// if there aren't any including requirements, we're in by default
   561  		included := len(reqIncl) == 0
   562  
   563  		// convert the pod's namespace to an equivalent label selector
   564  		selector := labels.Set{pod.Namespace: ""}
   565  
   566  		// include pod if one including requirement matches
   567  		for _, req := range reqIncl {
   568  			if req.Matches(selector) {
   569  				included = true
   570  				break
   571  			}
   572  		}
   573  
   574  		// exclude pod if it is filtered out by at least one excluding requirement
   575  		for _, req := range reqExcl {
   576  			if !req.Matches(selector) {
   577  				included = false
   578  				break
   579  			}
   580  		}
   581  
   582  		if included {
   583  			filteredList = append(filteredList, pod)
   584  		}
   585  	}
   586  
   587  	return filteredList, nil
   588  }
   589  
   590  func parseSelector(str string) (labels.Selector, error) {
   591  	selector, err := labels.Parse(str)
   592  	if err != nil {
   593  		return nil, err
   594  	}
   595  	return selector, nil
   596  }
   597  
   598  func getFixedSubListFromPodList(pods []v1.Pod, num int) []v1.Pod {
   599  	indexes := RandomFixedIndexes(0, uint(len(pods)), uint(num))
   600  
   601  	var filteredPods []v1.Pod
   602  
   603  	for _, index := range indexes {
   604  		index := index
   605  		filteredPods = append(filteredPods, pods[index])
   606  	}
   607  
   608  	return filteredPods
   609  }
   610  
   611  // RandomFixedIndexes returns the `count` random indexes between `start` and `end`.
   612  // [start, end)
   613  func RandomFixedIndexes(start, end, count uint) []uint {
   614  	var indexes []uint
   615  	m := make(map[uint]uint, count)
   616  
   617  	if end < start {
   618  		return indexes
   619  	}
   620  
   621  	if count > end-start {
   622  		for i := start; i < end; i++ {
   623  			indexes = append(indexes, i)
   624  		}
   625  
   626  		return indexes
   627  	}
   628  
   629  	for i := 0; i < int(count); {
   630  		index := uint(rand.Intn(int(end-start))) + start
   631  
   632  		_, exist := m[index]
   633  		if exist {
   634  			continue
   635  		}
   636  
   637  		m[index] = index
   638  		indexes = append(indexes, index)
   639  		i++
   640  	}
   641  
   642  	return indexes
   643  }
   644