adding flow to restore deployment env

resilience-bench/operator/src/main/java/io/resiliencebench/execution/steps/AbstractEnvironmentStep.java

@@ -0,0 +1,62 @@
+package io.resiliencebench.execution.steps;
+
+import io.fabric8.kubernetes.api.model.EnvVar;
+import io.fabric8.kubernetes.api.model.apps.Deployment;
+import io.fabric8.kubernetes.client.KubernetesClient;
+import io.resiliencebench.resources.scenario.Scenario;
+import io.resiliencebench.resources.service.ResilientService;
+import io.resiliencebench.support.CustomResourceRepository;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+abstract class AbstractEnvironmentStep extends ExecutorStep<Deployment> {
+
+    private final static Logger logger = LoggerFactory.getLogger(AbstractEnvironmentStep.class);
+    protected final CustomResourceRepository<ResilientService> resilientServiceRepository;
+
+    public AbstractEnvironmentStep(KubernetesClient kubernetesClient,
+            CustomResourceRepository<ResilientService> resilientServiceRepository) {
+        super(kubernetesClient);
+        this.resilientServiceRepository = resilientServiceRepository;
+    }
+
+    protected Deployment getDeployment(Scenario scenario, ResilientService resilientService) {
+        return kubernetesClient()
+                .apps()
+                .deployments()
+                .inNamespace(scenario.getMetadata().getNamespace())
+                .withLabelSelector(resilientService.getSpec().getSelector())
+                .list()
+                .getItems()
+                .stream()
+                .findFirst()
+                .orElse(null);
+    }
+
+    protected void updateDeployment(Deployment targetDeployment) {
+        kubernetesClient().apps().deployments()
+                .inNamespace(targetDeployment.getMetadata().getNamespace())
+                .resource(targetDeployment)
+                .update();
+    }
+
+    protected void restartPods(Deployment targetDeployment) {
+        logger.info("Waiting for the pods to restart");
+        kubernetesClient().pods()
+                .inNamespace(targetDeployment.getMetadata().getNamespace())
+                .withLabel("app", targetDeployment.getMetadata().getName())
+                .waitUntilReady(60, TimeUnit.SECONDS);
+        logger.info("Pods restarted successfully");
+    }
+
+    protected List<EnvVar> getActualEnv(Deployment targetDeployment, String containerName) {
+        var container = targetDeployment.getSpec().getTemplate().getSpec().getContainers().stream()
+                .filter(c -> c.getName().equals(containerName))
+                .findFirst()
+                .orElseThrow(() -> new RuntimeException("Container not found: " + containerName));
+        return container.getEnv();
+    }
+}

resilience-bench/operator/src/main/java/io/resiliencebench/execution/steps/EnvironmentPostStep.java

@@ -0,0 +1,76 @@
+package io.resiliencebench.execution.steps;
+
+import io.fabric8.kubernetes.api.model.apps.Deployment;
+import io.fabric8.kubernetes.client.KubernetesClient;
+import io.resiliencebench.resources.queue.ExecutionQueue;
+import io.resiliencebench.resources.scenario.Scenario;
+import io.resiliencebench.resources.service.ResilientService;
+import io.resiliencebench.support.CustomResourceRepository;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Service;
+
+import static io.resiliencebench.support.Annotations.CONTAINER;
+
+@Service
+public class EnvironmentPostStep extends AbstractEnvironmentStep {
+
+  private final static Logger logger = LoggerFactory.getLogger(EnvironmentPostStep.class);
+
+  public EnvironmentPostStep(KubernetesClient kubernetesClient,
+      CustomResourceRepository<ResilientService> resilientServiceRepository) {
+    super(kubernetesClient, resilientServiceRepository);
+  }
+
+  @Override
+  protected boolean isApplicable(Scenario scenario) {
+    return scenario
+        .getSpec()
+        .getConnectors()
+        .stream()
+        .anyMatch(connector -> connector.getDestination().getEnvs() != null ||
+            connector.getSource().getEnvs() != null);
+  }
+
+  @Override
+  protected Deployment internalExecute(Scenario scenario, ExecutionQueue queue) {
+    for (var connector : scenario.getSpec().getConnectors()) {
+      restoreEnvironment(scenario, connector.getSource());
+      restoreEnvironment(scenario, connector.getDestination());
+    }
+    return null;
+  }
+
+  public void restoreEnvironment(Scenario scenario, io.resiliencebench.resources.scenario.Service service) {
+    var resilientService = resilientServiceRepository.get(scenario.getMetadata().getNamespace(), service.getName());
+    var containerName = resilientService.getMetadata().getAnnotations().get(CONTAINER);
+
+    var env = resilientService.getSpec().getEnvs();
+    if (env == null) {
+      return;
+    }
+
+    var deployment = kubernetesClient()
+        .apps()
+        .deployments()
+        .inNamespace(scenario.getMetadata().getNamespace())
+        .withLabelSelector(resilientService.getSpec().getSelector())
+        .list()
+        .getItems()
+        .stream()
+        .findFirst();
+
+    if (deployment.isPresent()) {
+      var container = deployment.get().getSpec().getTemplate().getSpec().getContainers().stream()
+          .filter(c -> c.getName().equals(containerName))
+          .findFirst()
+          .orElseThrow(() -> new RuntimeException("Container not found: " + containerName));
+
+      container.setEnv(env);
+      updateDeployment(deployment.get());
+      restartPods(deployment.get());
+    } else {
+      logger.warn("Deployment not found for ResilientService {}", service.getName());
+    }
+  }
+}

resilience-bench/operator/src/main/java/io/resiliencebench/execution/steps/EnvironmentStep.java

@@ -1,6 +1,6 @@
 package io.resiliencebench.execution.steps;
 
-import com.fasterxml.jackson.databind.JsonNode;
+import io.fabric8.kubernetes.api.model.EnvVar;
 import io.fabric8.kubernetes.api.model.apps.Deployment;
 import io.fabric8.kubernetes.client.KubernetesClient;
 import io.resiliencebench.resources.queue.ExecutionQueue;
@@ -11,32 +11,28 @@
 import org.slf4j.LoggerFactory;
 import org.springframework.stereotype.Service;
 
-import java.util.Map;
-import java.util.concurrent.TimeUnit;
+import java.util.List;
 
 import static io.resiliencebench.support.Annotations.CONTAINER;
 
 @Service
-public class EnvironmentStep extends ExecutorStep<Deployment> {
+public class EnvironmentStep extends AbstractEnvironmentStep {
 
   private final static Logger logger = LoggerFactory.getLogger(EnvironmentStep.class);
 
-  private final CustomResourceRepository<ResilientService> resilientServiceRepository;
-
   public EnvironmentStep(KubernetesClient kubernetesClient,
-                        CustomResourceRepository<ResilientService> resilientServiceRepository) {
-    super(kubernetesClient);
-    this.resilientServiceRepository = resilientServiceRepository;
+      CustomResourceRepository<ResilientService> resilientServiceRepository) {
+    super(kubernetesClient, resilientServiceRepository);
   }
 
   @Override
   protected boolean isApplicable(Scenario scenario) {
     return scenario
-            .getSpec()
-            .getConnectors()
-            .stream()
-            .anyMatch(connector -> connector.getDestination().getEnvs() != null ||
-                    connector.getSource().getEnvs() != null);
+        .getSpec()
+        .getConnectors()
+        .stream()
+        .anyMatch(connector -> connector.getDestination().getEnvs() != null ||
+            connector.getSource().getEnvs() != null);
   }
 
   public void applyServiceEnvironment(Scenario scenario, io.resiliencebench.resources.scenario.Service service) {
@@ -48,17 +44,28 @@ public void applyServiceEnvironment(Scenario scenario, io.resiliencebench.resour
     var containerName = resilientService.getMetadata().getAnnotations().get(CONTAINER);
 
     var deployment = kubernetesClient()
-            .apps()
-            .deployments()
-            .inNamespace(scenario.getMetadata().getNamespace())
-            .withLabelSelector(resilientService.getSpec().getSelector())
-            .list()
-            .getItems()
-            .stream()
-            .findFirst();
+        .apps()
+        .deployments()
+        .inNamespace(scenario.getMetadata().getNamespace())
+        .withLabelSelector(resilientService.getSpec().getSelector())
+        .list()
+        .getItems()
+        .stream()
+        .findFirst();
 
     if (deployment.isPresent()) {
-      applyNewVariables(deployment.get(), containerName, env);
+      var deploymentVars = getActualEnv(deployment.get(), CONTAINER);
+
+      for (var variable : deploymentVars) {
+        var newValue = env.get(variable.getName());
+        if (newValue != null) {
+          logger.info("Container {}. EnvVar {}={}", containerName, variable.getName(), newValue);
+          variable.setValue(newValue.toString());
+        }
+      }
+      updateDeployment(deployment.get());
+      restartPods(deployment.get());
+      saveActualEnv(deploymentVars, resilientService);
     } else {
       logger.warn("Deployment not found for ResilientService {}", service.getName());
     }
@@ -73,41 +80,10 @@ protected Deployment internalExecute(Scenario scenario, ExecutionQueue queue) {
     return null;
   }
 
-  public void applyNewVariables(Deployment targetDeployment, String containerName, Map<String, JsonNode> env) {
-    var container = targetDeployment.getSpec().getTemplate().getSpec().getContainers().stream()
-            .filter(c -> c.getName().equals(containerName))
-            .findFirst()
-            .orElseThrow(() -> new RuntimeException("Container not found: " + containerName));
-    var deploymentVars = container
-            .getEnv()
-            .stream()
-            .toList();
-
-    for (var variable : deploymentVars) {
-      var newValue = env.get(variable.getName());
-      if (newValue != null) {
-        logger.info("Container {}. EnvVar {}={}", containerName, variable.getName(), newValue);
-        variable.setValue(newValue.toString());
-      }
-    }
-
-    updateDeployment(targetDeployment);
-    restartPods(targetDeployment);
-  }
-
-  private void updateDeployment(Deployment targetDeployment) {
-    kubernetesClient().apps().deployments()
-            .inNamespace(targetDeployment.getMetadata().getNamespace())
-            .resource(targetDeployment)
-            .update();
+  private void saveActualEnv(List<EnvVar> deploymentVars, ResilientService resilientService) {
+    resilientService.getSpec().getEnvs().clear();
+    resilientService.getSpec().getEnvs().addAll(deploymentVars);
+    resilientServiceRepository.update(resilientService);
   }
 
-  private void restartPods(Deployment targetDeployment) {
-    logger.info("Waiting for the pods to restart");
-    kubernetesClient().pods()
-            .inNamespace(targetDeployment.getMetadata().getNamespace())
-            .withLabel("app", targetDeployment.getMetadata().getName())
-            .waitUntilReady(60, TimeUnit.SECONDS);
-    logger.info("Pods restarted successfully");
-  }
 }

resilience-bench/operator/src/main/java/io/resiliencebench/execution/steps/StepRegister.java

@@ -16,25 +16,25 @@ public class StepRegister {
   private final List<ExecutorStep<?>> postExecutionSteps;
 
   public StepRegister(UpdateStatusQueueStep updateStatusQueueStep,
-                      ResultLocalFileStep resultLocalFileStep,
-                      IstioCircuitBreakerStep istioCircuitBreakerStep,
-                      IstioRetryStep istioRetryStep,
-                      IstioTimeoutStep istioTimeoutStep,
-                      IstioFaultStep istioFaultStep,
-                      EnvironmentStep environmentStep) {
+      ResultLocalFileStep resultLocalFileStep,
+      IstioCircuitBreakerStep istioCircuitBreakerStep,
+      IstioRetryStep istioRetryStep,
+      IstioTimeoutStep istioTimeoutStep,
+      IstioFaultStep istioFaultStep,
+      EnvironmentStep environmentStep,
+      EnvironmentPostStep environmentPostStep) {
 
     preparationSteps = List.of(
-            updateStatusQueueStep,
-            istioRetryStep,
-            istioCircuitBreakerStep,
-            istioTimeoutStep,
-            istioFaultStep,
-            environmentStep
-    );
+        updateStatusQueueStep,
+        istioRetryStep,
+        istioCircuitBreakerStep,
+        istioTimeoutStep,
+        istioFaultStep,
+        environmentStep);
     postExecutionSteps = List.of(
-            updateStatusQueueStep,
-            resultLocalFileStep
-    );
+        updateStatusQueueStep,
+        resultLocalFileStep,
+        environmentPostStep);
   }
 
   public List<ExecutorStep<?>> getPostExecutionSteps() {

resilience-bench/operator/src/main/java/io/resiliencebench/resources/service/ResilientServiceSpec.java

@@ -1,10 +1,13 @@
 package io.resiliencebench.resources.service;
 
+import java.util.List;
+import io.fabric8.kubernetes.api.model.EnvVar;
 import io.fabric8.kubernetes.api.model.LabelSelector;
 
 public class ResilientServiceSpec {
 
   private LabelSelector selector;
+  private List<EnvVar> envs;
 
   public LabelSelector getSelector() {
     return selector;
@@ -13,4 +16,12 @@ public LabelSelector getSelector() {
   public void setSelector(LabelSelector selector) {
     this.selector = selector;
   }
+
+  public List<EnvVar> getEnvs() {
+    return envs;
+  }
+
+  public void setEnvs(List<EnvVar> envs) {
+    this.envs = envs;
+  }
 }