diff --git a/client/alloc_endpoint.go b/client/alloc_endpoint.go index c22179606ce9..9ce4b50b30f9 100644 --- a/client/alloc_endpoint.go +++ b/client/alloc_endpoint.go @@ -48,6 +48,7 @@ func (a *Allocations) GarbageCollect(args *nstructs.AllocSpecificRequest, reply return nil } +// Signal is used to send a signal to an allocation's tasks on a client. func (a *Allocations) Signal(args *nstructs.AllocSignalRequest, reply *nstructs.GenericResponse) error { defer metrics.MeasureSince([]string{"client", "allocations", "signal"}, time.Now()) diff --git a/client/alloc_endpoint_test.go b/client/alloc_endpoint_test.go index b1bf686ba8d4..a29d60eeeccc 100644 --- a/client/alloc_endpoint_test.go +++ b/client/alloc_endpoint_test.go @@ -356,14 +356,6 @@ func TestAllocations_Signal_ACL(t *testing.T) { } } -func TestAllocations_Signal_Subtask(t *testing.T) { - t.Parallel() -} - -func TestAllocations_Signal_EntireAlloc(t *testing.T) { - t.Parallel() -} - func TestAllocations_Stats(t *testing.T) { t.Parallel() require := require.New(t) diff --git a/command/alloc_signal.go b/command/alloc_signal.go index 0aa40bb3d3d4..e879dcc1e93b 100644 --- a/command/alloc_signal.go +++ b/command/alloc_signal.go @@ -3,6 +3,10 @@ package command import ( "fmt" "strings" + + "github.com/hashicorp/nomad/api" + "github.com/hashicorp/nomad/api/contexts" + "github.com/posener/complete" ) type AllocSignalCommand struct { @@ -71,11 +75,6 @@ func (c *AllocSignalCommand) Run(args []string) int { allocID = sanitizeUUIDPrefix(allocID) - var taskName string - if len(args) == 2 { - taskName = args[1] - } - // Get the HTTP client client, err := c.Meta.Client() if err != nil { @@ -108,6 +107,17 @@ func (c *AllocSignalCommand) Run(args []string) int { return 1 } + var taskName string + if len(args) == 2 { + // Validate Task + taskName = args[1] + err := validateTaskExistsInAllocation(taskName, alloc) + if err != nil { + c.Ui.Error(err.Error()) + return 1 + } + } + err = client.Allocations().Signal(alloc, nil, taskName, signal) if err != nil { c.Ui.Error(fmt.Sprintf("Error signalling allocation: %s", err)) @@ -120,3 +130,74 @@ func (c *AllocSignalCommand) Run(args []string) int { func (a *AllocSignalCommand) Synopsis() string { return "Signal a running allocation" } + +func (c *AllocSignalCommand) AutocompleteFlags() complete.Flags { + return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient), + complete.Flags{ + "-s": complete.PredictNothing, + "-verbose": complete.PredictNothing, + }) +} +func (c *AllocSignalCommand) AutocompleteArgs() complete.Predictor { + return complete.PredictFunc(func(a complete.Args) []string { + client, err := c.Meta.Client() + if err != nil { + return nil + } + + // The complete.All array will return [signal, argN...], including partial + // args. + // Here we autocomplete the alloc id if there is either no partial arg or if + // there is a prefix to search by. + // We then autocomplete task names from the allocation for the subsequent + // argument and return nothing if there are more args than [signal, alloc, task] + + count := len(a.All) + + if count > 3 { + return []string{} + } + + if count <= 2 { + // Searching for an alloc + resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Allocs, nil) + if err != nil { + return []string{} + } + return resp.Matches[contexts.Allocs] + } + + // Searching for a task name in the specified alloc + + allocID := sanitizeUUIDPrefix(a.All[1]) + allocs, _, err := client.Allocations().PrefixList(allocID) + + if len(allocs) != 1 { + return []string{} + } + + // Prefix lookup matched a single allocation + alloc, _, err := client.Allocations().Info(allocs[0].ID, nil) + if err != nil { + return []string{} + } + + return filterAllocationTasks(alloc, a.Last) + }) +} + +func filterAllocationTasks(a *api.Allocation, prefix string) []string { + tg := a.Job.LookupTaskGroup(a.TaskGroup) + if tg == nil { + return []string{} + } + + var matching []string + for _, task := range tg.Tasks { + if strings.HasPrefix(task.Name, prefix) { + matching = append(matching, task.Name) + } + } + + return matching +} diff --git a/command/alloc_signal_test.go b/command/alloc_signal_test.go index 02fea72419d3..92e0b3b83fa7 100644 --- a/command/alloc_signal_test.go +++ b/command/alloc_signal_test.go @@ -1,9 +1,16 @@ package command import ( + "fmt" "testing" + "github.com/hashicorp/nomad/api" + "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" ) @@ -48,3 +55,99 @@ func TestAllocSignalCommand_Fails(t *testing.T) { require.Contains(ui.ErrorWriter.String(), "must contain at least two characters.") ui.ErrorWriter.Reset() } + +func TestAllocSignalCommand_AutocompleteArgs(t *testing.T) { + assert := assert.New(t) + + srv, _, url := testServer(t, true, nil) + defer srv.Shutdown() + + ui := new(cli.MockUi) + cmd := &AllocSignalCommand{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{All: []string{"signal", prefix}, Last: prefix} + predictor := cmd.AutocompleteArgs() + + // Match Allocs + res := predictor.Predict(args) + assert.Equal(1, len(res)) + assert.Equal(a.ID, res[0]) + + // Match Tasknames + args = complete.Args{All: []string{"signal", a.ID, "w"}, Last: "w"} + res = predictor.Predict(args) + assert.Equal(1, len(res)) + assert.Equal("web", res[0]) + + // Ignore extra args + args = complete.Args{All: []string{"signal", a.ID, "web", "we"}, Last: "we"} + res = predictor.Predict(args) + assert.Empty(res) +} + +func TestAllocSignalCommand_Run(t *testing.T) { + srv, client, url := testServer(t, true, nil) + defer srv.Shutdown() + + require := require.New(t) + + // 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) { + t.Fatalf("err: %v", err) + }) + + ui := new(cli.MockUi) + cmd := &AllocSignalCommand{Meta: Meta{Ui: ui}} + + jobID := "job1_sfx" + job1 := testJob(jobID) + resp, _, err := client.Jobs().Register(job1, nil) + require.NoError(err) + if code := waitForSuccess(ui, client, fullId, t, resp.EvalID); code != 0 { + t.Fatalf("status code non zero saw %d", code) + } + // 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(allocId1, "unable to find allocation") + + // Wait for alloc to be running + testutil.WaitForResult(func() (bool, error) { + alloc, _, err := client.Allocations().Info(allocId1, nil) + if err != nil { + return false, err + } + if alloc.ClientStatus == api.AllocClientStatusRunning { + return true, nil + } + return false, fmt.Errorf("alloc is not running, is: %s", alloc.ClientStatus) + }, func(err error) { + t.Fatalf("err: %v", err) + }) + + require.Equal(cmd.Run([]string{"-address=" + url, allocId1}), 0, "expected successful exit code") + + ui.OutputWriter.Reset() +}