Cleanup cancelled disruption

Signed-off-by: Pablo Chacin <[email protected]>

pkg/disruptors/cmd_builders.go

@@ -113,3 +113,7 @@ func buildHTTPFaultCmd(
 
 	return cmd
 }
+
+func buildCleanupCmd() []string {
+	return []string{"xk6-disruptor-agent", "cleanup"}
+}

pkg/disruptors/controller.go

@@ -13,6 +13,14 @@ import (
 	corev1 "k8s.io/api/core/v1"
 )
 
+// VisitOpts define the options for visiting a Pod
+type VisitOpts struct {
+	// Command defines the command to be executed
+	Command []string
+	// Cleanup defines the command to execute for cleaning up if command execution fails
+	Cleanup []string
+}
+
 // AgentController defines the interface for controlling agents in a set of targets
 type AgentController interface {
 	// InjectDisruptorAgent injects the Disruptor agent in the target pods
@@ -22,7 +30,7 @@ type AgentController interface {
 	// Targets retrieves the names of the target of the controller
 	Targets(ctx context.Context) ([]string, error)
 	// Visit allows executing a different command on each target returned by a visiting function
-	Visit(ctx context.Context, visitor func(target corev1.Pod) ([]string, error)) error
+	Visit(ctx context.Context, visitor func(target corev1.Pod) (VisitOpts, error)) error
 }
 
 // AgentController controls de agents in a set of target pods
@@ -96,13 +104,13 @@ func (c *agentController) InjectDisruptorAgent(ctx context.Context) error {
 // ExecCommand executes a command in the targets of the AgentController and reports any error
 func (c *agentController) ExecCommand(ctx context.Context, cmd []string) error {
 	// visit each target with the same command
-	return c.Visit(ctx, func(corev1.Pod) ([]string, error) {
-		return cmd, nil
+	return c.Visit(ctx, func(corev1.Pod) (VisitOpts, error) {
+		return VisitOpts{Command: cmd}, nil
 	})
 }
 
 // Visit allows executing a different command on each target returned by a visiting function
-func (c *agentController) Visit(ctx context.Context, visitor func(corev1.Pod) ([]string, error)) error {
+func (c *agentController) Visit(ctx context.Context, visitor func(corev1.Pod) (VisitOpts, error)) error {
 	// if there are no targets, nothing to do
 	if len(c.targets) == 0 {
 		return nil
@@ -111,50 +119,57 @@ func (c *agentController) Visit(ctx context.Context, visitor func(corev1.Pod) ([
 	execContext, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	// ensure errors channel has enough space to avoid blocking gorutines
-	errors := make(chan error, len(c.targets))
+	// ensure errCh channel has enough space to avoid blocking gorutines
+	errCh := make(chan error, len(c.targets))
 	for _, pod := range c.targets {
 		// attach each container asynchronously
 		go func(pod corev1.Pod) {
 			// get the command to execute in the target
-			cmd, err := visitor(pod)
+			visitOpts, err := visitor(pod)
 			if err != nil {
-				errors <- fmt.Errorf("error building command for pod %s: %w", pod.Name, err)
+				errCh <- fmt.Errorf("error building command for pod %s: %w", pod.Name, err)
 				return
 			}
 
-			_, stderr, err := c.helper.Exec(execContext, pod.Name, "xk6-agent", cmd, []byte{})
+			_, stderr, err := c.helper.Exec(execContext, pod.Name, "xk6-agent", visitOpts.Command, []byte{})
+
+			// if command failed, ensure the agent execution is terminated
+			if err != nil && visitOpts.Cleanup != nil {
+				// we ignore errors because k6 was cancelled, so there's no point in reporting
+				// use a fresh context because the exec context may have been cancelled or expired
+				//nolint:contextcheck
+				_, _, _ = c.helper.Exec(context.TODO(), pod.Name, "xk6-agent", visitOpts.Cleanup, []byte{})
+			}
+
 			if err != nil {
-				errors <- fmt.Errorf("error invoking agent: %w \n%s", err, string(stderr))
-				return
+				errCh <- fmt.Errorf("error invoking agent: %w \n%s", err, string(stderr))
 			}
 
-			errors <- nil
+			errCh <- nil
 		}(pod)
 	}
 
+	var err error
 	pending := len(c.targets)
 	for {
 		select {
-		case err := <-errors:
-			if err != nil {
-				return err
-			}
+		case e := <-errCh:
 			pending--
+			if e != nil {
+				// cancel ongoing commands
+				cancel()
+				// Save first received error as reason for ending execution
+				err = e
+			}
+
 			if pending == 0 {
-				return nil
+				return err
 			}
 		case <-ctx.Done():
-			// FIXME: the command is hardcoded and breaks separation of concerns
-			cancelCmd := []string{"xk6-disruptor-agent", "cleanup"}
-			// TODO: consider sending these commands in parallel to all targets
-			for _, pod := range c.targets {
-				// we ignore errors because k6 was cancelled, so there's no point in reporting
-				// TODO: define which context to use, depending if this is done in "background"
-				//nolint:contextcheck
-				_, _, _ = c.helper.Exec(context.TODO(), pod.Name, "xk6-agent", cancelCmd, []byte{})
-			}
-			return ctx.Err()
+			// cancel ongoing commands
+			cancel()
+			// save the reason for ending execution
+			err = ctx.Err()
 		}
 	}
 }

pkg/disruptors/pod.go

@@ -119,23 +119,26 @@ func (d *podDisruptor) InjectHTTPFaults(
 		fault.Port = DefaultTargetPort
 	}
 
-	return d.controller.Visit(ctx, func(pod corev1.Pod) ([]string, error) {
+	return d.controller.Visit(ctx, func(pod corev1.Pod) (VisitOpts, error) {
 		if !utils.HasPort(pod, fault.Port) {
-			return nil, fmt.Errorf("pod %q does not expose port %d", pod.Name, fault.Port)
+			return VisitOpts{}, fmt.Errorf("pod %q does not expose port %d", pod.Name, fault.Port)
 		}
 
 		if utils.HasHostNetwork(pod) {
-			return nil, fmt.Errorf("pod %q cannot be safely injected as it has hostNetwork set to true", pod.Name)
+			return VisitOpts{}, fmt.Errorf("pod %q cannot be safely injected as it has hostNetwork set to true", pod.Name)
 		}
 
 		targetAddress, err := utils.PodIP(pod)
 		if err != nil {
-			return nil, err
+			return VisitOpts{}, err
 		}
 
-		cmd := buildHTTPFaultCmd(targetAddress, fault, duration, options)
+		visitOps := VisitOpts{
+			Command: buildHTTPFaultCmd(targetAddress, fault, duration, options),
+			Cleanup: buildCleanupCmd(),
+		}
 
-		return cmd, nil
+		return visitOps, nil
 	})
 }
 
@@ -146,17 +149,21 @@ func (d *podDisruptor) InjectGrpcFaults(
 	duration time.Duration,
 	options GrpcDisruptionOptions,
 ) error {
-	return d.controller.Visit(ctx, func(pod corev1.Pod) ([]string, error) {
+	return d.controller.Visit(ctx, func(pod corev1.Pod) (VisitOpts, error) {
 		if !utils.HasPort(pod, fault.Port) {
-			return nil, fmt.Errorf("pod %q does not expose port %d", pod.Name, fault.Port)
+			return VisitOpts{}, fmt.Errorf("pod %q does not expose port %d", pod.Name, fault.Port)
 		}
 
 		targetAddress, err := utils.PodIP(pod)
 		if err != nil {
-			return nil, err
+			return VisitOpts{}, err
+		}
+
+		visitOps := VisitOpts{
+			Command: buildGrpcFaultCmd(targetAddress, fault, duration, options),
+			Cleanup: buildCleanupCmd(),
 		}
 
-		cmd := buildGrpcFaultCmd(targetAddress, fault, duration, options)
-		return cmd, nil
+		return visitOps, nil
 	})
 }

pkg/disruptors/service.go

@@ -90,14 +90,14 @@ func (d *serviceDisruptor) InjectHTTPFaults(
 ) error {
 	// for each target, the port to inject the fault can be different
 	// we use the Visit function and generate a command for each pod
-	return d.controller.Visit(ctx, func(pod corev1.Pod) ([]string, error) {
+	return d.controller.Visit(ctx, func(pod corev1.Pod) (VisitOpts, error) {
 		port, err := utils.MapPort(d.service, fault.Port, pod)
 		if err != nil {
-			return nil, err
+			return VisitOpts{}, err
 		}
 
 		if utils.HasHostNetwork(pod) {
-			return nil, fmt.Errorf("pod %q cannot be safely injected as it has hostNetwork set to true", pod.Name)
+			return VisitOpts{}, fmt.Errorf("pod %q cannot be safely injected as it has hostNetwork set to true", pod.Name)
 		}
 
 		// copy fault to change target port for the pod
@@ -106,11 +106,15 @@ func (d *serviceDisruptor) InjectHTTPFaults(
 
 		targetAddress, err := utils.PodIP(pod)
 		if err != nil {
-			return nil, err
+			return VisitOpts{}, err
 		}
 
-		cmd := buildHTTPFaultCmd(targetAddress, podFault, duration, options)
-		return cmd, nil
+		visitOps := VisitOpts{
+			Command: buildHTTPFaultCmd(targetAddress, podFault, duration, options),
+			Cleanup: buildCleanupCmd(),
+		}
+
+		return visitOps, nil
 	})
 }
 
@@ -122,22 +126,26 @@ func (d *serviceDisruptor) InjectGrpcFaults(
 ) error {
 	// for each target, the port to inject the fault can be different
 	// we use the Visit function and generate a command for each pod
-	return d.controller.Visit(ctx, func(pod corev1.Pod) ([]string, error) {
+	return d.controller.Visit(ctx, func(pod corev1.Pod) (VisitOpts, error) {
 		port, err := utils.MapPort(d.service, fault.Port, pod)
 		if err != nil {
-			return nil, err
+			return VisitOpts{}, err
 		}
 
 		podFault := fault
 		podFault.Port = port
 
 		targetAddress, err := utils.PodIP(pod)
 		if err != nil {
-			return nil, err
+			return VisitOpts{}, err
+		}
+
+		visitOps := VisitOpts{
+			Command: buildGrpcFaultCmd(targetAddress, podFault, duration, options),
+			Cleanup: buildCleanupCmd(),
 		}
 
-		cmd := buildGrpcFaultCmd(targetAddress, podFault, duration, options)
-		return cmd, nil
+		return visitOps, nil
 	})
 }