add support for on_exit deferring it for all the commands

README.md

@@ -463,6 +463,27 @@ example installing curl package if not installed already:
     cond: "! command -v curl"
 ```
 
+### Deferred actions (`on_exit`)
+
+Each command may have `on_exit` parameter defined. It allows executing a command on the remote host after the task with all commands is completed. The command is called regardless of the task's exit code.
+
+This is useful in several scenarios:
+
+- a temporary script copied to the remote host and executed and should be removed after execution with `on_exit: "rm -fv /tmp/script.sh"`
+- a service should be restarted after the new version is deployed with `on_exit: "systemctl restart myservice"`
+
+
+```yaml
+  - name: "copy script"
+    copy: {src: "testdata/script.sh", "dst": "/tmp/script.sh", "chmod+x": true}
+    on_exit: "rm -fv /tmp/script.sh" # register deferred action to remove script.sh after execution
+  - name: "run script"
+    script: "/tmp/script.sh"
+```
+
+In the example above, the `script.sh` is copied to the remote host, executed, and removed after completion of the task.
+
+
 ### Script Execution
 
 Spot allows executing scripts on remote hosts, or locally if `options.local` is set to true. Scripts can be executed in two different ways, depending on whether they are single-line or multi-line scripts.

pkg/config/command.go

@@ -30,6 +30,7 @@ type Cmd struct {
 	Options     CmdOptions        `yaml:"options" toml:"options,omitempty"`
 	Condition   string            `yaml:"cond" toml:"cond,omitempty"`
 	Register    []string          `yaml:"register" toml:"register"` // register variables from command
+	OnExit      string            `yaml:"on_exit" toml:"on_exit"`   // script to run on exit
 
 	Secrets  map[string]string `yaml:"-" toml:"-"` // loaded secrets, filled by playbook
 	SSHShell string            `yaml:"-" toml:"-"` // shell to use for ssh commands, filled by playbook

pkg/runner/commands.go

@@ -29,12 +29,14 @@ type execCmd struct {
 	exec     executor.Interface
 	verbose  bool
 	sshShell string
+	onExit   string
 }
 
 type execCmdResp struct {
 	details string
 	verbose string
 	vars    map[string]string
+	onExit  execCmd
 }
 
 type execCmdErr struct {

pkg/runner/runner.go

@@ -182,6 +182,19 @@ func (p *Process) runTaskOnHost(ctx context.Context, tsk *config.Task, hostAddr,
 	// copy task to prevent one task on hostA modifying task on hostB as it does updateVars
 	activeTask := deepcopy.Copy(*tsk).(config.Task)
 
+	onExitCmds := []execCmd{}
+	defer func() {
+		// run on-exit commands if any. it is executed after all commands of the task are done or on error
+		if len(onExitCmds) > 0 {
+			log.Printf("[DEBUG] run %d on-exit commands on %s", len(onExitCmds), hostAddr)
+			for _, ec := range onExitCmds {
+				if _, err := ec.Script(ctx); err != nil {
+					report(ec.hostAddr, ec.hostName, "failed on-exit command %q (%v)", ec.cmd.Name, err)
+				}
+			}
+		}
+	}()
+
 	for _, cmd := range activeTask.Commands {
 		if !p.shouldRunCmd(cmd, hostName, hostAddr) {
 			continue
@@ -191,7 +204,7 @@ func (p *Process) runTaskOnHost(ctx context.Context, tsk *config.Task, hostAddr,
 		stCmd := time.Now()
 
 		ec := execCmd{cmd: cmd, hostAddr: hostAddr, hostName: hostName, tsk: &activeTask, exec: remote,
-			verbose: p.Verbose, sshShell: p.SSHShell}
+			verbose: p.Verbose, sshShell: p.SSHShell, onExit: cmd.OnExit}
 		ec = p.pickCmdExecutor(cmd, ec, hostAddr, hostName) // pick executor on dry run or local command
 
 		repHostAddr, repHostName := ec.hostAddr, ec.hostName
@@ -205,6 +218,10 @@ func (p *Process) runTaskOnHost(ctx context.Context, tsk *config.Task, hostAddr,
 		}
 
 		exResp, err := p.execCommand(ctx, ec)
+		if exResp.onExit.cmd.Name != "" { // we have on-exit command, save it for later execution
+			// this is intentionally before error check, we want to run on-exit command even if the main command failed
+			onExitCmds = append(onExitCmds, exResp.onExit)
+		}
 		if err != nil {
 			if !cmd.Options.IgnoreErrors {
 				return count, nil, fmt.Errorf("failed command %q on host %s (%s): %w", cmd.Name, ec.hostAddr, ec.hostName, err)
@@ -238,13 +255,28 @@ func (p *Process) runTaskOnHost(ctx context.Context, tsk *config.Task, hostAddr,
 	} else {
 		report("localhost", "", "completed task %q, commands: %d (%v)\n", activeTask.Name, count, since(stTask))
 	}
+
 	return count, tskVars, nil
 }
 
 // execCommand executes a single command on a target host.
 // It detects the command type based on the fields what are set.
 // Even if multiple fields for multiple commands are set, only one will be executed.
 func (p *Process) execCommand(ctx context.Context, ec execCmd) (resp execCmdResp, err error) {
+
+	if ec.cmd.OnExit != "" {
+		// register on-exit command if any set
+		defer func() {
+			// we need to defer it because it changes the command name and script
+			log.Printf("[DEBUG] defer execution on_exit script on %s for %s", ec.hostAddr, ec.cmd.Name)
+			// use the same executor as for the main command but with different script and name
+			ec.cmd.Name = "on exit for " + ec.cmd.Name
+			ec.cmd.Script = ec.cmd.OnExit
+			ec.cmd.OnExit = "" // prevent recursion
+			resp.onExit = ec
+		}()
+	}
+
 	switch {
 	case ec.cmd.Script != "":
 		log.Printf("[DEBUG] execute script %q on %s", ec.cmd.Name, ec.hostAddr)

pkg/runner/runner_test.go

@@ -721,6 +721,58 @@ func TestProcess_RunFailedErrIgnored(t *testing.T) {
 	require.NoError(t, err, "error ignored")
 }
 
+func TestProcess_RunWithOnExit(t *testing.T) {
+	ctx := context.Background()
+	testingHostAndPort, teardown := startTestContainer(t)
+	defer teardown()
+
+	logs := executor.MakeLogs(false, false, nil)
+	connector, err := executor.NewConnector("testdata/test_ssh_key", time.Second*10, logs)
+	require.NoError(t, err)
+	conf, err := config.New("testdata/conf.yml", nil, nil)
+	require.NoError(t, err)
+
+	p := Process{
+		Concurrency: 1,
+		Connector:   connector,
+		Playbook:    conf,
+		Logs:        logs,
+	}
+
+	t.Run("on_exit called on script completion", func(t *testing.T) {
+		var buf bytes.Buffer
+		log.SetOutput(&buf)
+
+		_, err = p.Run(ctx, "with_onexit", testingHostAndPort)
+		require.NoError(t, err)
+		t.Log(buf.String())
+		require.Contains(t, buf.String(), "> file content")
+		require.Contains(t, buf.String(), "> on exit called. task: with_onexit")
+		require.Contains(t, buf.String(), "> /bin/sh -c 'ls -la /tmp/file.txt'")
+	})
+
+	t.Run("on_exit called on script failed", func(t *testing.T) {
+		var buf bytes.Buffer
+		log.SetOutput(&buf)
+
+		_, err = p.Run(ctx, "with_onexit_failed", testingHostAndPort)
+		require.Error(t, err)
+		t.Log(buf.String())
+		require.Contains(t, buf.String(), "> on exit called on failed. task: with_onexit_failed")
+	})
+
+	t.Run("on_exit called on copy completion", func(t *testing.T) {
+		var buf bytes.Buffer
+		log.SetOutput(&buf)
+
+		_, err = p.Run(ctx, "with_onexit_copy", testingHostAndPort)
+		require.NoError(t, err)
+		t.Log(buf.String())
+		require.Contains(t, buf.String(), "> on exit called for copy. task: with_onexit_copy")
+		require.Contains(t, buf.String(), "> removed '/tmp/conf-blah.yml'")
+	})
+}
+
 func TestProcess_RunTaskWithWait(t *testing.T) {
 	ctx := context.Background()
 	testingHostAndPort, teardown := startTestContainer(t)

pkg/runner/testdata/conf.yml

@@ -165,3 +165,34 @@ tasks:
           echo good command 2
       - name: good command
         script: echo good command 3
+
+  - name: with_onexit
+    commands:
+      - name: some command
+        script: |
+          echo good command 1
+          echo "file content" > /tmp/file.txt
+          echo good command 2
+        on_exit: |
+          echo "on exit called. task: ${SPOT_TASK}, host: ${SPOT_REMOTE_HOST}, error: ${SPOT_ERROR}"
+          cat /tmp/file.txt
+      - name: "print file info"
+        script: ls -la /tmp/file.txt
+
+  - name: with_onexit_failed
+    commands:
+      - name: some command
+        script: |
+          echo good command 1
+          ls /no-such-dir
+        on_exit: |
+          echo "on exit called on failed. task: ${SPOT_TASK}, host: ${SPOT_REMOTE_HOST}, error: ${SPOT_ERROR}"
+
+  - name: with_onexit_copy
+    commands:
+      - name: some command
+        copy: {"src": "testdata/conf.yml", "dst": "/tmp/conf-blah.yml", "mkdir": true, "force": true}
+        on_exit: |
+          echo "on exit called for copy. task: ${SPOT_TASK}, host: ${SPOT_REMOTE_HOST}, error: ${SPOT_ERROR}"
+          rm -fv /tmp/conf-blah.yml
+