From bad77f1ece5938b112f37acf5c5fce9b72e636e2 Mon Sep 17 00:00:00 2001 From: Mahmood Ali Date: Sun, 28 Apr 2019 17:35:17 -0400 Subject: [PATCH] add CLI commands for nomad exec --- command/alloc_exec.go | 315 ++++++++++++++++++++++++++++++++++ command/alloc_exec_test.go | 287 +++++++++++++++++++++++++++++++ command/alloc_exec_unix.go | 14 ++ command/alloc_exec_windows.go | 9 + command/alloc_logs.go | 41 +++-- command/alloc_restart.go | 10 +- command/commands.go | 10 ++ 7 files changed, 660 insertions(+), 26 deletions(-) create mode 100644 command/alloc_exec.go create mode 100644 command/alloc_exec_test.go create mode 100644 command/alloc_exec_unix.go create mode 100644 command/alloc_exec_windows.go diff --git a/command/alloc_exec.go b/command/alloc_exec.go new file mode 100644 index 000000000..7417e61ec --- /dev/null +++ b/command/alloc_exec.go @@ -0,0 +1,315 @@ +package command + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "os" + "os/signal" + "strings" + "syscall" + + "github.com/docker/docker/pkg/term" + "github.com/hashicorp/nomad/api" + "github.com/hashicorp/nomad/api/contexts" + "github.com/posener/complete" +) + +type AllocExecCommand struct { + Meta + + Stdin io.Reader + Stdout io.WriteCloser + Stderr io.WriteCloser +} + +func (l *AllocExecCommand) Help() string { + helpText := ` +Usage: nomad alloc exec [options] + + Run command inside the environment of the given allocation and task. + +General Options: + + ` + generalOptionsUsage() + ` + +Exec Specific Options: + + -task + Sets the task to exec command in + + -job + Use a random allocation from the specified job ID. + + -i + Pass stdin to the container, defaults to true + + -t + Allocate a pseudo-tty, defaults to true if stdin is detected to be a tty session + ` + return strings.TrimSpace(helpText) +} + +func (l *AllocExecCommand) Synopsis() string { + return "Execute commands in task" +} + +func (c *AllocExecCommand) AutocompleteFlags() complete.Flags { + return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient), + complete.Flags{ + "--task": complete.PredictAnything, + "-job": complete.PredictAnything, + "-i": complete.PredictNothing, + "-t": complete.PredictNothing, + }) +} + +func (l *AllocExecCommand) AutocompleteArgs() complete.Predictor { + return complete.PredictFunc(func(a complete.Args) []string { + client, err := l.Meta.Client() + if err != nil { + return nil + } + + resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Allocs, nil) + if err != nil { + return []string{} + } + return resp.Matches[contexts.Allocs] + }) +} + +func (l *AllocExecCommand) Name() string { return "alloc exec" } + +func (l *AllocExecCommand) Run(args []string) int { + var job, stdinOpt, ttyOpt bool + var task string + + flags := l.Meta.FlagSet(l.Name(), FlagSetClient) + flags.Usage = func() { l.Ui.Output(l.Help()) } + flags.BoolVar(&job, "job", false, "") + flags.StringVar(&task, "task", "", "") + + flags.BoolVar(&stdinOpt, "i", true, "") + + stdinTty := isStdinTty() + flags.BoolVar(&ttyOpt, "t", stdinTty, "") + + if err := flags.Parse(args); err != nil { + return 1 + } + args = flags.Args() + + if ttyOpt && !stdinOpt { + l.Ui.Error("-i must be enabled if running with tty") + return -1 + } + + if numArgs := len(args); numArgs < 1 { + if job { + l.Ui.Error("A job ID is required") + } else { + l.Ui.Error("An allocation ID is required") + } + + l.Ui.Error(commandErrorText(l)) + return 1 + } else if numArgs < 2 { + l.Ui.Error("A command is required") + l.Ui.Error(commandErrorText(l)) + return 1 + } + + command := args[1:] + + client, err := l.Meta.Client() + if err != nil { + l.Ui.Error(fmt.Sprintf("Error initializing client: %v", err)) + return 1 + } + + // If -job is specified, use random allocation, otherwise use provided allocation + allocID := args[0] + if job { + allocID, err = getRandomJobAlloc(client, args[0]) + if err != nil { + l.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err)) + return 1 + } + } + + length := shortId + + // Query the allocation info + if len(allocID) == 1 { + l.Ui.Error(fmt.Sprintf("Alloc ID must contain at least two characters.")) + return 1 + } + + allocID = sanitizeUUIDPrefix(allocID) + allocs, _, err := client.Allocations().PrefixList(allocID) + if err != nil { + l.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err)) + return 1 + } + if len(allocs) == 0 { + l.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID)) + return 1 + } + if len(allocs) > 1 { + // Format the allocs + out := formatAllocListStubs(allocs, false, length) + l.Ui.Error(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", out)) + return 1 + } + // Prefix lookup matched a single allocation + alloc, _, err := client.Allocations().Info(allocs[0].ID, nil) + if err != nil { + l.Ui.Error(fmt.Sprintf("Error querying allocation: %s", err)) + return 1 + } + + if task == "" { + task, err = lookupAllocTask(alloc) + + if err != nil { + l.Ui.Error(err.Error()) + l.Ui.Error("\nPlease specify the task.") + return 1 + } + } + + if err := validateTaskExistsInAllocation(task, alloc); err != nil { + l.Ui.Error(err.Error()) + return 1 + } + + if l.Stdin == nil { + l.Stdin = os.Stdin + } + if l.Stdout == nil { + l.Stdout = os.Stdout + } + if l.Stderr == nil { + l.Stderr = os.Stderr + } + + var stdin io.Reader = l.Stdin + if !stdinOpt { + stdin = bytes.NewReader(nil) + } + + code, err := l.execImpl(client, alloc, task, ttyOpt, command, stdin, l.Stdout, l.Stderr) + if err != nil { + l.Ui.Error(fmt.Sprintf("failed to exec into task: %v", err)) + return 1 + } + + return code +} + +func isStdinTty() bool { + _, isTerminal := term.GetFdInfo(os.Stdin) + return isTerminal +} + +func setRawTerminal(stream interface{}) (cleanup func(), err error) { + fd, isTerminal := term.GetFdInfo(stream) + if !isTerminal { + return nil, errors.New("not a terminal") + } + + state, err := term.SetRawTerminal(fd) + if err != nil { + return nil, err + } + + return func() { term.RestoreTerminal(fd, state) }, nil +} + +func watchTerminalSize(out io.Writer, resize chan<- api.TerminalSize) (func(), error) { + fd, isTerminal := term.GetFdInfo(out) + if !isTerminal { + return nil, errors.New("not a terminal") + } + + ctx, cancel := context.WithCancel(context.Background()) + + signalCh := make(chan os.Signal, 1) + setupWindowNotification(signalCh) + + sendTerminalSize := func() { + s, err := term.GetWinsize(fd) + if err != nil { + return + } + + resize <- api.TerminalSize{ + Height: int(s.Height), + Width: int(s.Width), + } + } + go func() { + for { + select { + case <-ctx.Done(): + return + case <-signalCh: + sendTerminalSize() + } + } + }() + + go func() { + // send initial size + sendTerminalSize() + }() + + return cancel, nil +} + +func (l *AllocExecCommand) execImpl(client *api.Client, alloc *api.Allocation, task string, tty bool, + command []string, stdin io.Reader, stdout, stderr io.WriteCloser) (int, error) { + + sizeCh := make(chan api.TerminalSize, 1) + + if tty { + if stdin == nil { + return -1, fmt.Errorf("stdin is null") + } + + inCleanup, err := setRawTerminal(stdin) + if err != nil { + return -1, err + } + defer inCleanup() + + outCleanup, err := setRawTerminal(stdout) + if err != nil { + return -1, err + } + defer outCleanup() + + sizeCleanup, err := watchTerminalSize(stdout, sizeCh) + if err != nil { + return -1, err + } + defer sizeCleanup() + } + + ctx, cancelFn := context.WithCancel(context.Background()) + defer cancelFn() + + signalCh := make(chan os.Signal, 1) + signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM) + go func() { + for range signalCh { + cancelFn() + } + }() + + return client.Allocations().Exec(ctx, + alloc, task, tty, command, stdin, stdout, stderr, sizeCh, nil) +} diff --git a/command/alloc_exec_test.go b/command/alloc_exec_test.go new file mode 100644 index 000000000..ead498e53 --- /dev/null +++ b/command/alloc_exec_test.go @@ -0,0 +1,287 @@ +package command + +import ( + "bytes" + "fmt" + "strings" + "testing" + + "github.com/hashicorp/nomad/helper/uuid" + "github.com/hashicorp/nomad/nomad/mock" + "github.com/hashicorp/nomad/nomad/structs" + "github.com/hashicorp/nomad/testutil" + "github.com/mitchellh/cli" + "github.com/posener/complete" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// static check +var _ cli.Command = &AllocExecCommand{} + +func TestAllocExecCommand_Fails(t *testing.T) { + t.Parallel() + srv, client, url := testServer(t, true, nil) + defer srv.Shutdown() + + cases := []struct { + name string + args []string + expectedError string + }{ + { + "misuse", + []string{"bad"}, + commandErrorText(&AllocExecCommand{}), + }, + { + "connection failure", + []string{"-address=nope", "26470238-5CF2-438F-8772-DC67CFB0705C", "/bin/bash"}, + "Error querying allocation", + }, + { + "not found alloc", + []string{"-address=" + url, "26470238-5CF2-438F-8772-DC67CFB0705C", "/bin/bash"}, + "No allocation(s) with prefix or id", + }, + { + "not found job", + []string{"-address=" + url, "-job", "example", "/bin/bash"}, + `job "example" doesn't exist`, + }, + { + "too short allocis", + []string{"-address=" + url, "2", "/bin/bash"}, + "Alloc ID must contain at least two characters", + }, + { + "missing command", + []string{"-address=" + url, "26470238-5CF2-438F-8772-DC67CFB0705C"}, + "A command is required", + }, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + ui := new(cli.MockUi) + cmd := &AllocExecCommand{Meta: Meta{Ui: ui}} + + code := cmd.Run(c.args) + require.Equal(t, 1, code) + + require.Contains(t, ui.ErrorWriter.String(), c.expectedError) + + ui.ErrorWriter.Reset() + ui.OutputWriter.Reset() + + }) + } + + // Wait for a node to be ready + testutil.WaitForResult(func() (bool, error) { + nodes, _, err := client.Nodes().List(nil) + if err != nil { + return false, err + } + for _, node := range nodes { + if _, ok := node.Drivers["mock_driver"]; ok && + node.Status == structs.NodeStatusReady { + return true, nil + } + } + return false, fmt.Errorf("no ready nodes") + }, func(err error) { + require.NoError(t, err) + }) + + t.Run("non existent task", func(t *testing.T) { + ui := new(cli.MockUi) + cmd := &AllocExecCommand{Meta: Meta{Ui: ui}} + + jobID := "job1_sfx" + job1 := testJob(jobID) + resp, _, err := client.Jobs().Register(job1, nil) + require.NoError(t, err) + code := waitForSuccess(ui, client, fullId, t, resp.EvalID) + require.Zero(t, code, "status code not zero") + + // get an alloc id + allocId1 := "" + if allocs, _, err := client.Jobs().Allocations(jobID, false, nil); err == nil { + if len(allocs) > 0 { + allocId1 = allocs[0].ID + } + } + require.NotEmpty(t, allocId1, "unable to find allocation") + + // by alloc + require.Equal(t, 1, cmd.Run([]string{"-address=" + url, "-task=nonexistenttask1", allocId1, "/bin/bash"})) + require.Contains(t, ui.ErrorWriter.String(), "Could not find task named: nonexistenttask1") + ui.ErrorWriter.Reset() + + // by jobID + require.Equal(t, 1, cmd.Run([]string{"-address=" + url, "-task=nonexistenttask2", "-job", jobID, "/bin/bash"})) + require.Contains(t, ui.ErrorWriter.String(), "Could not find task named: nonexistenttask2") + ui.ErrorWriter.Reset() + }) + +} + +func TestAllocExecCommand_AutocompleteArgs(t *testing.T) { + assert := assert.New(t) + t.Parallel() + + srv, _, url := testServer(t, true, nil) + defer srv.Shutdown() + + ui := new(cli.MockUi) + cmd := &AllocExecCommand{Meta: Meta{Ui: ui, flagAddress: url}} + + // Create a fake alloc + state := srv.Agent.Server().State() + a := mock.Alloc() + assert.Nil(state.UpsertAllocs(1000, []*structs.Allocation{a})) + + prefix := a.ID[:5] + args := complete.Args{Last: prefix} + predictor := cmd.AutocompleteArgs() + + res := predictor.Predict(args) + assert.Equal(1, len(res)) + assert.Equal(a.ID, res[0]) +} + +func TestAllocExecCommand_Run(t *testing.T) { + t.Parallel() + srv, client, url := testServer(t, true, nil) + defer srv.Shutdown() + + // Wait for a node to be ready + testutil.WaitForResult(func() (bool, error) { + nodes, _, err := client.Nodes().List(nil) + if err != nil { + return false, err + } + + for _, node := range nodes { + if _, ok := node.Drivers["mock_driver"]; ok && + node.Status == structs.NodeStatusReady { + return true, nil + } + } + return false, fmt.Errorf("no ready nodes") + }, func(err error) { + require.NoError(t, err) + }) + + jobID := uuid.Generate() + job := testJob(jobID) + job.TaskGroups[0].Tasks[0].Config = map[string]interface{}{ + "run_for": "10s", + "exec_command": map[string]interface{}{ + "run_for": "1ms", + "exit_code": 21, + "stdout_string": "sample stdout output\n", + "stderr_string": "sample stderr output\n", + }, + } + resp, _, err := client.Jobs().Register(job, nil) + require.NoError(t, err) + + evalUi := new(cli.MockUi) + code := waitForSuccess(evalUi, client, fullId, t, resp.EvalID) + require.Equal(t, 0, code, "failed to get status - output: %v", evalUi.ErrorWriter.String()) + + allocId := "" + + testutil.WaitForResult(func() (bool, error) { + allocs, _, err := client.Jobs().Allocations(jobID, false, nil) + if err != nil { + return false, fmt.Errorf("failed to get allocations: %v", err) + } + + if len(allocs) < 0 { + return false, fmt.Errorf("no allocations yet") + } + + alloc := allocs[0] + if alloc.ClientStatus != "running" { + return false, fmt.Errorf("alloc is not running yet: %v", alloc.ClientStatus) + } + + allocId = alloc.ID + return true, nil + }, func(err error) { + require.NoError(t, err) + + }) + + cases := []struct { + name string + command string + stdin string + + stdout string + stderr string + exitCode int + }{ + { + name: "basic stdout/err", + command: "simplecommand", + stdin: "", + stdout: "sample stdout output", + stderr: "sample stderr output", + exitCode: 21, + }, + { + name: "notty: streamining input", + command: "showinput", + stdin: "hello from stdin", + stdout: "TTY: false\nStdin:\nhello from stdin", + exitCode: 0, + }, + } + + for _, c := range cases { + t.Run("by id: "+c.name, func(t *testing.T) { + ui := new(cli.MockUi) + var stdout, stderr bufferCloser + + cmd := &AllocExecCommand{ + Meta: Meta{Ui: ui}, + Stdin: strings.NewReader(c.stdin), + Stdout: &stdout, + Stderr: &stderr, + } + + code = cmd.Run([]string{"-address=" + url, allocId, c.command}) + assert.Equal(t, c.exitCode, code) + assert.Equal(t, c.stdout, strings.TrimSpace(stdout.String())) + assert.Equal(t, c.stderr, strings.TrimSpace(stderr.String())) + }) + t.Run("by job: "+c.name, func(t *testing.T) { + ui := new(cli.MockUi) + var stdout, stderr bufferCloser + + cmd := &AllocExecCommand{ + Meta: Meta{Ui: ui}, + Stdin: strings.NewReader(c.stdin), + Stdout: &stdout, + Stderr: &stderr, + } + + code = cmd.Run([]string{"-address=" + url, "-job", jobID, c.command}) + assert.Equal(t, c.exitCode, code) + assert.Equal(t, c.stdout, strings.TrimSpace(stdout.String())) + assert.Equal(t, c.stderr, strings.TrimSpace(stderr.String())) + }) + } +} + +type bufferCloser struct { + bytes.Buffer +} + +func (b *bufferCloser) Close() error { + return nil +} diff --git a/command/alloc_exec_unix.go b/command/alloc_exec_unix.go new file mode 100644 index 000000000..36c3901df --- /dev/null +++ b/command/alloc_exec_unix.go @@ -0,0 +1,14 @@ +// +build darwin dragonfly freebsd linux netbsd openbsd solaris + +package command + +import ( + "os" + "os/signal" + + "golang.org/x/sys/unix" +) + +func setupWindowNotification(ch chan<- os.Signal) { + signal.Notify(ch, unix.SIGWINCH) +} diff --git a/command/alloc_exec_windows.go b/command/alloc_exec_windows.go new file mode 100644 index 000000000..882c757c0 --- /dev/null +++ b/command/alloc_exec_windows.go @@ -0,0 +1,9 @@ +package command + +import ( + "os" +) + +func setupWindowNotification(ch chan<- os.Signal) { + // do nothing +} diff --git a/command/alloc_logs.go b/command/alloc_logs.go index 3691c52bf..2717a01ab 100644 --- a/command/alloc_logs.go +++ b/command/alloc_logs.go @@ -1,6 +1,7 @@ package command import ( + "errors" "fmt" "io" "os" @@ -185,25 +186,10 @@ func (l *AllocLogsCommand) Run(args []string) int { } } else { - // Try to determine the tasks name from the allocation - var tasks []*api.Task - for _, tg := range alloc.Job.TaskGroups { - if *tg.Name == alloc.TaskGroup { - if len(tg.Tasks) == 1 { - task = tg.Tasks[0].Name - break - } + task, err = lookupAllocTask(alloc) - tasks = tg.Tasks - break - } - } - - if task == "" { - l.Ui.Error(fmt.Sprintf("Allocation %q is running the following tasks:", limit(alloc.ID, length))) - for _, t := range tasks { - l.Ui.Error(fmt.Sprintf(" * %s", t.Name)) - } + if err != nil { + l.Ui.Error(err.Error()) l.Ui.Error("\nPlease specify the task.") return 1 } @@ -294,3 +280,22 @@ func (l *AllocLogsCommand) followFile(client *api.Client, alloc *api.Allocation, return r, nil } + +func lookupAllocTask(alloc *api.Allocation) (string, error) { + tg := alloc.Job.LookupTaskGroup(alloc.TaskGroup) + if tg == nil { + return "", fmt.Errorf("Could not find allocation task group: %s", alloc.TaskGroup) + } + + if len(tg.Tasks) == 1 { + return tg.Tasks[0].Name, nil + } + + var errStr strings.Builder + fmt.Fprintf(&errStr, "Allocation %q is running the following tasks:\n", limit(alloc.ID, shortId)) + for _, t := range tg.Tasks { + fmt.Fprintf(&errStr, " * %s\n", t.Name) + } + fmt.Fprintf(&errStr, "\nPlease specify the task.") + return "", errors.New(errStr.String()) +} diff --git a/command/alloc_restart.go b/command/alloc_restart.go index 68b2f5637..b27b0f363 100644 --- a/command/alloc_restart.go +++ b/command/alloc_restart.go @@ -128,21 +128,15 @@ func validateTaskExistsInAllocation(taskName string, alloc *api.Allocation) erro return fmt.Errorf("Could not find allocation task group: %s", alloc.TaskGroup) } - taskExists := false foundTaskNames := make([]string, len(tg.Tasks)) for i, task := range tg.Tasks { foundTaskNames[i] = task.Name if task.Name == taskName { - taskExists = true - break + return nil } } - if !taskExists { - return fmt.Errorf("Could not find task named: %s, found:\n%s", taskName, formatList(foundTaskNames)) - } - - return nil + return fmt.Errorf("Could not find task named: %s, found:\n%s", taskName, formatList(foundTaskNames)) } func (a *AllocRestartCommand) Synopsis() string { diff --git a/command/commands.go b/command/commands.go index 9619a44f3..1c1666a49 100644 --- a/command/commands.go +++ b/command/commands.go @@ -145,6 +145,11 @@ func Commands(metaPtr *Meta, agentUi cli.Ui) map[string]cli.CommandFactory { Meta: meta, }, nil }, + "alloc exec": func() (cli.Command, error) { + return &AllocExecCommand{ + Meta: meta, + }, nil + }, "alloc signal": func() (cli.Command, error) { return &AllocSignalCommand{ Meta: meta, @@ -252,6 +257,11 @@ func Commands(metaPtr *Meta, agentUi cli.Ui) map[string]cli.CommandFactory { Meta: meta, }, nil }, + "exec": func() (cli.Command, error) { + return &AllocExecCommand{ + Meta: meta, + }, nil + }, "executor": func() (cli.Command, error) { return &ExecutorPluginCommand{ Meta: meta,