shell.go 3.56 KB
Newer Older
ale's avatar
ale committed
1 2 3 4
package tabacco

import (
	"context"
5
	"io/ioutil"
ale's avatar
ale committed
6 7 8
	"log"
	"os"
	"os/exec"
9
	"path/filepath"
ale's avatar
ale committed
10
	"strconv"
11
	"strings"
12

ale's avatar
ale committed
13
	"git.autistici.org/ai3/tools/tabacco/jobs"
ale's avatar
ale committed
14 15
)

16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41
// Environment variables starting with any of the following strings
// are removed from the command execution environment.
var filteredEnvVars = []string{
	"LANG=",
	"LC_",
	"RESTIC_",
}

// Return a copy of os.Environ(), with filteredEnvVars removed.
func safeEnv() []string {
	var env []string
	for _, s := range os.Environ() {
		skip := false
		for _, pfx := range filteredEnvVars {
			if strings.HasPrefix(s, pfx) {
				skip = true
				break
			}
		}
		if !skip {
			env = append(env, s)
		}
	}
	return env
}

ale's avatar
ale committed
42 43 44 45 46 47 48
// Shell runs commands, with some options (a global dry-run flag
// preventing all executions, nice level, i/o class). As one may guess
// by the name, commands are run using the shell, so variable
// substitutions and other shell features are available.
type Shell struct {
	dryRun      bool
	niceLevel   int
ale's avatar
ale committed
49
	ioniceClass int
50
	env         []string
ale's avatar
ale committed
51 52 53 54 55 56
}

// NewShell creates a new Shell.
func NewShell(dryRun bool) *Shell {
	return &Shell{
		dryRun:      dryRun,
ale's avatar
ale committed
57 58
		niceLevel:   10,
		ioniceClass: 2,
59
		env:         safeEnv(),
ale's avatar
ale committed
60 61 62 63 64 65 66 67 68
	}
}

// SetNiceLevel sets the nice(1) level.
func (s *Shell) SetNiceLevel(n int) {
	s.niceLevel = n
}

// SetIOClass sets the ionice(1) i/o class.
ale's avatar
ale committed
69 70
func (s *Shell) SetIOClass(n int) {
	s.ioniceClass = n
ale's avatar
ale committed
71 72
}

73 74 75
// command builds an exec.Cmd and gets some parameters from the
// context - notably it sets log output and working directory to be in
// the working dir if the job has been wrapped in WithWorkDir().
ale's avatar
ale committed
76 77 78 79 80 81 82 83 84 85
func (s *Shell) command(ctx context.Context, arg string) *exec.Cmd {
	var args []string
	if s.dryRun {
		args = []string{"/bin/echo", arg}
	} else {
		args = []string{"/bin/sh", "-c", arg}
	}

	if s.niceLevel != 0 {
		args = append(
86
			[]string{"nice", "-n", strconv.Itoa(s.niceLevel)},
ale's avatar
ale committed
87 88 89
			args...,
		)
	}
ale's avatar
ale committed
90
	if s.ioniceClass != 0 {
ale's avatar
ale committed
91
		args = append(
92
			[]string{"ionice", "-c", strconv.Itoa(s.ioniceClass)},
ale's avatar
ale committed
93 94 95 96 97
			args...,
		)
	}

	c := exec.CommandContext(ctx, args[0], args[1:]...) // #nosec
98 99 100
	var env []string
	env = append(env, s.env...)
	c.Env = env
ale's avatar
ale committed
101
	c.Stderr = os.Stderr
102
	c.Dir = getWorkDir(ctx)
103 104 105

	log.Printf("sh: %s", arg)

ale's avatar
ale committed
106 107 108 109 110 111 112 113 114 115 116 117 118 119
	return c
}

// Run a command.
func (s *Shell) Run(ctx context.Context, arg string) error {
	c := s.command(ctx, arg)
	c.Stdout = os.Stdout
	return c.Run()
}

// Output runs a command and returns the standard output.
func (s *Shell) Output(ctx context.Context, arg string) ([]byte, error) {
	return s.command(ctx, arg).Output()
}
120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172

// Manage a work directory (temporary scratch with lots of space).
type workdirManager struct {
	dir string
}

func newWorkdirManager(tmpdir string) (*workdirManager, error) {
	dir, err := ioutil.TempDir(tmpdir, "tabacco-")
	if err != nil {
		return nil, err
	}
	return &workdirManager{dir}, nil
}

func (w *workdirManager) Close() {
	os.RemoveAll(w.dir) // nolint
}

func (w *workdirManager) withWorkDir(j jobs.Job) jobs.Job {
	jobDir := filepath.Join(w.dir, "job-"+j.ID())
	return &workdirJob{
		Job: j,
		dir: jobDir,
	}
}

type workdirJob struct {
	jobs.Job
	dir string
}

type workdirKeyType int

var workdirKey workdirKeyType = 1

func (j *workdirJob) RunContext(ctx context.Context) error {
	if err := os.Mkdir(j.dir, 0700); err != nil {
		return err
	}

	ctx = context.WithValue(ctx, workdirKey, j.dir)
	return j.Job.RunContext(ctx)
}

// getWorkDir returns the current working directory, provided the
// WithWorkDir wrapper has been called.
func getWorkDir(ctx context.Context) string {
	dir, ok := ctx.Value(workdirKey).(string)
	if ok {
		return dir
	}
	return "/tmp"
}