Adding node_label_selector for pod scenarios

Signed-off-by: Sahil Shah <[email protected]>

Author: shahsahil264

krkn/scenario_plugins/pod_disruption/models/models.py

@@ -11,11 +11,13 @@ def __init__(self, config: dict[str,any] = None):
             self.label_selector = config["label_selector"] if "label_selector" in config else ""
             self.namespace_pattern = config["namespace_pattern"] if "namespace_pattern" in config else ""
             self.name_pattern = config["name_pattern"] if "name_pattern" in config else ""
+            self.node_label_selector = config["node_label_selector"] if "node_label_selector" in config else ""
 
     namespace_pattern: str
     krkn_pod_recovery_time: int
     timeout: int
     duration: int
     kill: int
     label_selector: str
-    name_pattern: str
+    name_pattern: str
+    node_label_selector: str

krkn/scenario_plugins/pod_disruption/pod_disruption_scenario_plugin.py

@@ -106,7 +106,7 @@ def start_monitoring(self, kill_scenario: InputParams, pool: PodsMonitorPool):
             )
 
 
-    def get_pods(self, name_pattern, label_selector,namespace, kubecli: KrknKubernetes, field_selector: str =None): 
+    def get_pods(self, name_pattern, label_selector,namespace, kubecli: KrknKubernetes, field_selector: str =None, node_label_selector: str = None): 
         if label_selector and name_pattern: 
             logging.error('Only, one of name pattern or label pattern can be specified')
         elif label_selector:
@@ -115,6 +115,22 @@ def get_pods(self, name_pattern, label_selector,namespace, kubecli: KrknKubernet
             pods = kubecli.select_pods_by_name_pattern_and_namespace_pattern(pod_name_pattern=name_pattern, namespace_pattern=namespace, field_selector=field_selector)
         else: 
             logging.error('Name pattern or label pattern must be specified ')
+        
+        # Filter pods by node label selector if specified
+        if node_label_selector and pods:
+            filtered_pods = []
+            nodes_with_label = kubecli.list_nodes(label_selector=node_label_selector)
+            if not nodes_with_label:
+                return []
+            
+            for pod_name, pod_namespace in pods:
+                pod_info = kubecli.read_pod(pod_name, pod_namespace)
+                pod_node_name = pod_info.spec.node_name
+                if pod_node_name and pod_node_name in nodes_with_label:
+                    filtered_pods.append((pod_name, pod_namespace))
+            
+            return filtered_pods
+        
         return pods
 
     def killing_pods(self, config: InputParams, kubecli: KrknKubernetes):
@@ -124,7 +140,7 @@ def killing_pods(self, config: InputParams, kubecli: KrknKubernetes):
         if not namespace: 
             logging.error('Namespace pattern must be specified')
 
-        pods = self.get_pods(config.name_pattern,config.label_selector,config.namespace_pattern, kubecli, field_selector="status.phase=Running")
+        pods = self.get_pods(config.name_pattern,config.label_selector,config.namespace_pattern, kubecli, field_selector="status.phase=Running", node_label_selector=config.node_label_selector)
         pods_count = len(pods)
         if len(pods) < config.kill:
             logging.error("Not enough pods match the criteria, expected {} but found only {} pods".format(
@@ -133,23 +149,22 @@ def killing_pods(self, config: InputParams, kubecli: KrknKubernetes):
 
         random.shuffle(pods)
         for i in range(config.kill):
-
             pod = pods[i]
             logging.info(pod)
             logging.info(f'Deleting pod {pod[0]}')
             kubecli.delete_pod(pod[0], pod[1])
 
-        self.wait_for_pods(config.label_selector,config.name_pattern,config.namespace_pattern, pods_count, config.duration, config.timeout, kubecli)
+        self.wait_for_pods(config.label_selector,config.name_pattern,config.namespace_pattern, pods_count, config.duration, config.timeout, kubecli, config.node_label_selector)
         return 0
 
     def wait_for_pods(
-        self, label_selector, pod_name, namespace, pod_count, duration, wait_timeout, kubecli: KrknKubernetes
+        self, label_selector, pod_name, namespace, pod_count, duration, wait_timeout, kubecli: KrknKubernetes, node_label_selector
     ):
         timeout = False
         start_time = datetime.now()
 
         while not timeout:
-            pods = self.get_pods(name_pattern=pod_name, label_selector=label_selector,namespace=namespace, field_selector="status.phase=Running", kubecli=kubecli)
+            pods = self.get_pods(name_pattern=pod_name, label_selector=label_selector,namespace=namespace, field_selector="status.phase=Running", kubecli=kubecli, node_label_selector=node_label_selector)
             if pod_count == len(pods):
                 return
 

scenarios/kind/pod_etcd.yml

@@ -2,4 +2,5 @@
   config:
     namespace_pattern: "kube-system"
     label_selector: "component=etcd"
+    node_label_selector: node-role.kubernetes.io/control-plane=  # Target control-plane nodes (kind clusters)
     krkn_pod_recovery_time: 120

scenarios/kube/pod.yml

@@ -5,3 +5,4 @@
     namespace_pattern: ^default$
     kill: 1
     krkn_pod_recovery_time: 120
+    node_label_selector: node-role.kubernetes.io/worker=  # Target worker nodes only

scenarios/openshift/customapp_pod.yaml

@@ -3,4 +3,5 @@
   config:
     namespace_pattern: ^acme-air$
     name_pattern: .*
-    krkn_pod_recovery_time: 120
+    krkn_pod_recovery_time: 120
+    node_label_selector: node-role.kubernetes.io/master=              # Target master nodes only

scenarios/openshift/openshift-apiserver.yml

@@ -3,5 +3,6 @@
   config:
     namespace_pattern: ^openshift-apiserver$
     label_selector: app=openshift-apiserver-a
+    node_label_selector: node-role.kubernetes.io/master=  # Target master nodes only
     krkn_pod_recovery_time: 120
 

scenarios/openshift/openshift-kube-apiserver.yml

@@ -3,5 +3,6 @@
   config:
     namespace_pattern: ^openshift-kube-apiserver$
     label_selector: app=openshift-kube-apiserver
+    node_label_selector: node-role.kubernetes.io/master=  # Target master nodes only
     krkn_pod_recovery_time: 120
 

scenarios/openshift/prom_kill.yml

@@ -2,4 +2,5 @@
   config:
     namespace_pattern: ^openshift-monitoring$
     label_selector: statefulset.kubernetes.io/pod-name=prometheus-k8s-0
-    krkn_pod_recovery_time: 120
+    krkn_pod_recovery_time: 120
+    node_label_selector: node-role.kubernetes.io/worker=              # Target worker nodes only

scenarios/openshift/regex_openshift_pod_kill.yml

@@ -5,3 +5,4 @@
     name_pattern: .*
     kill: 3
     krkn_pod_recovery_time: 120
+    node_label_selector: node-role.kubernetes.io/worker=