diff --git a/command/node_drain.go b/command/node_drain.go index 0b92a0a8e..18bd695c4 100644 --- a/command/node_drain.go +++ b/command/node_drain.go @@ -269,5 +269,7 @@ func (c *NodeDrainCommand) Run(args []string) int { c.Ui.Error(fmt.Sprintf("Error updating drain specification: %s", err)) return 1 } + + c.Ui.Output(fmt.Sprintf("Node %q drain strategy set", node.ID)) return 0 } diff --git a/command/node_eligibility.go b/command/node_eligibility.go new file mode 100644 index 000000000..2db14ddc2 --- /dev/null +++ b/command/node_eligibility.go @@ -0,0 +1,168 @@ +package command + +import ( + "fmt" + "strings" + + "github.com/hashicorp/nomad/api/contexts" + "github.com/posener/complete" +) + +type NodeEligibilityCommand struct { + Meta +} + +func (c *NodeEligibilityCommand) Help() string { + helpText := ` +Usage: nomad node eligibility [options] + + Toggles the nodes scheduling eligibility. When a node is marked as ineligible, + no new allocations will be placed on it but existing allocations will remain. + To remove existing allocations, use the node drain command. + + It is required that either -enable or -disable is specified, but not both. + The -self flag is useful to drain the local node. + +General Options: + + ` + generalOptionsUsage() + ` + +Node Eligibility Options: + + -disable + Mark the specified node as ineligible for new allocations. + + -enable + Mark the specified node as eligible for new allocations. + + -self + Set the eligibility of the local node. +` + return strings.TrimSpace(helpText) +} + +func (c *NodeEligibilityCommand) Synopsis() string { + return "Toggle scheduling eligibility for a given node" +} + +func (c *NodeEligibilityCommand) AutocompleteFlags() complete.Flags { + return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient), + complete.Flags{ + "-disable": complete.PredictNothing, + "-enable": complete.PredictNothing, + "-self": complete.PredictNothing, + }) +} + +func (c *NodeEligibilityCommand) AutocompleteArgs() complete.Predictor { + return complete.PredictFunc(func(a complete.Args) []string { + client, err := c.Meta.Client() + if err != nil { + return nil + } + + resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Nodes, nil) + if err != nil { + return []string{} + } + return resp.Matches[contexts.Nodes] + }) +} + +func (c *NodeEligibilityCommand) Run(args []string) int { + var enable, disable, self bool + + flags := c.Meta.FlagSet("node-eligibility", FlagSetClient) + flags.Usage = func() { c.Ui.Output(c.Help()) } + flags.BoolVar(&enable, "enable", false, "Mark node as eligibile for scheduling") + flags.BoolVar(&disable, "disable", false, "Mark node as ineligibile for scheduling") + flags.BoolVar(&self, "self", false, "") + + if err := flags.Parse(args); err != nil { + return 1 + } + + // Check that we got either enable or disable, but not both. + if (enable && disable) || (!enable && !disable) { + c.Ui.Error(c.Help()) + return 1 + } + + // Check that we got a node ID + args = flags.Args() + if l := len(args); self && l != 0 || !self && l != 1 { + c.Ui.Error("Node ID must be specified if -self isn't being used") + return 1 + } + + // Get the HTTP client + client, err := c.Meta.Client() + if err != nil { + c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err)) + return 1 + } + + // If -self flag is set then determine the current node. + var nodeID string + if !self { + nodeID = args[0] + } else { + var err error + if nodeID, err = getLocalNodeID(client); err != nil { + c.Ui.Error(err.Error()) + return 1 + } + } + + // Check if node exists + if len(nodeID) == 1 { + c.Ui.Error(fmt.Sprintf("Identifier must contain at least two characters.")) + return 1 + } + + nodeID = sanatizeUUIDPrefix(nodeID) + nodes, _, err := client.Nodes().PrefixList(nodeID) + if err != nil { + c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err)) + return 1 + } + // Return error if no nodes are found + if len(nodes) == 0 { + c.Ui.Error(fmt.Sprintf("No node(s) with prefix or id %q found", nodeID)) + return 1 + } + if len(nodes) > 1 { + // Format the nodes list that matches the prefix so that the user + // can create a more specific request + out := make([]string, len(nodes)+1) + out[0] = "ID|Datacenter|Name|Class|Drain|Status" + for i, node := range nodes { + out[i+1] = fmt.Sprintf("%s|%s|%s|%s|%v|%s", + node.ID, + node.Datacenter, + node.Name, + node.NodeClass, + node.Drain, + node.Status) + } + // Dump the output + c.Ui.Error(fmt.Sprintf("Prefix matched multiple nodes\n\n%s", formatList(out))) + return 1 + } + + // Prefix lookup matched a single node + node, _, err := client.Nodes().Info(nodes[0].ID, nil) + if err != nil { + c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err)) + return 1 + } + + // Toggle node eligibility + if _, err := client.Nodes().ToggleEligibility(node.ID, enable, nil); err != nil { + c.Ui.Error(fmt.Sprintf("Error updating scheduling eligibility: %s", err)) + return 1 + } + + c.Ui.Output(fmt.Sprintf("Node %q scheduling eligibility set", node.ID)) + return 0 +} diff --git a/command/node_eligibility_test.go b/command/node_eligibility_test.go new file mode 100644 index 000000000..3129fe86a --- /dev/null +++ b/command/node_eligibility_test.go @@ -0,0 +1,124 @@ +package command + +import ( + "fmt" + "strings" + "testing" + + "github.com/hashicorp/nomad/testutil" + "github.com/mitchellh/cli" + "github.com/posener/complete" + "github.com/stretchr/testify/assert" +) + +func TestNodeEligibilityCommand_Implements(t *testing.T) { + t.Parallel() + var _ cli.Command = &NodeEligibilityCommand{} +} + +func TestNodeEligibilityCommand_Fails(t *testing.T) { + t.Parallel() + srv, _, url := testServer(t, false, nil) + defer srv.Shutdown() + + ui := new(cli.MockUi) + cmd := &NodeEligibilityCommand{Meta: Meta{Ui: ui}} + + // Fails on misuse + if code := cmd.Run([]string{"some", "bad", "args"}); code != 1 { + t.Fatalf("expected exit code 1, got: %d", code) + } + if out := ui.ErrorWriter.String(); !strings.Contains(out, cmd.Help()) { + t.Fatalf("expected help output, got: %s", out) + } + ui.ErrorWriter.Reset() + + // Fails on connection failure + if code := cmd.Run([]string{"-address=nope", "-enable", "12345678-abcd-efab-cdef-123456789abc"}); code != 1 { + t.Fatalf("expected exit code 1, got: %d", code) + } + if out := ui.ErrorWriter.String(); !strings.Contains(out, "Error toggling") { + t.Fatalf("expected failed toggle error, got: %s", out) + } + ui.ErrorWriter.Reset() + + // Fails on non-existent node + if code := cmd.Run([]string{"-address=" + url, "-enable", "12345678-abcd-efab-cdef-123456789abc"}); code != 1 { + t.Fatalf("expected exit 1, got: %d", code) + } + if out := ui.ErrorWriter.String(); !strings.Contains(out, "No node(s) with prefix or id") { + t.Fatalf("expected not exist error, got: %s", out) + } + ui.ErrorWriter.Reset() + + // Fails if both enable and disable specified + if code := cmd.Run([]string{"-enable", "-disable", "12345678-abcd-efab-cdef-123456789abc"}); code != 1 { + t.Fatalf("expected exit 1, got: %d", code) + } + if out := ui.ErrorWriter.String(); !strings.Contains(out, cmd.Help()) { + t.Fatalf("expected help output, got: %s", out) + } + ui.ErrorWriter.Reset() + + // Fails if neither enable or disable specified + if code := cmd.Run([]string{"12345678-abcd-efab-cdef-123456789abc"}); code != 1 { + t.Fatalf("expected exit 1, got: %d", code) + } + if out := ui.ErrorWriter.String(); !strings.Contains(out, cmd.Help()) { + t.Fatalf("expected help output, got: %s", out) + } + ui.ErrorWriter.Reset() + + // Fail on identifier with too few characters + if code := cmd.Run([]string{"-address=" + url, "-enable", "1"}); code != 1 { + t.Fatalf("expected exit 1, got: %d", code) + } + if out := ui.ErrorWriter.String(); !strings.Contains(out, "must contain at least two characters.") { + t.Fatalf("expected too few characters error, got: %s", out) + } + ui.ErrorWriter.Reset() + + // Identifiers with uneven length should produce a query result + if code := cmd.Run([]string{"-address=" + url, "-enable", "123"}); code != 1 { + t.Fatalf("expected exit 1, got: %d", code) + } + if out := ui.ErrorWriter.String(); !strings.Contains(out, "No node(s) with prefix or id") { + t.Fatalf("expected not exist error, got: %s", out) + } + ui.ErrorWriter.Reset() +} + +func TestNodeEligibilityCommand_AutocompleteArgs(t *testing.T) { + assert := assert.New(t) + t.Parallel() + + srv, client, url := testServer(t, true, nil) + defer srv.Shutdown() + + // Wait for a node to appear + var nodeID string + testutil.WaitForResult(func() (bool, error) { + nodes, _, err := client.Nodes().List(nil) + if err != nil { + return false, err + } + if len(nodes) == 0 { + return false, fmt.Errorf("missing node") + } + nodeID = nodes[0].ID + return true, nil + }, func(err error) { + t.Fatalf("err: %s", err) + }) + + ui := new(cli.MockUi) + cmd := &NodeEligibilityCommand{Meta: Meta{Ui: ui, flagAddress: url}} + + prefix := nodeID[:len(nodeID)-5] + args := complete.Args{Last: prefix} + predictor := cmd.AutocompleteArgs() + + res := predictor.Predict(args) + assert.Equal(1, len(res)) + assert.Equal(nodeID, res[0]) +} diff --git a/commands.go b/commands.go index 9e27664af..0b3a422f0 100644 --- a/commands.go +++ b/commands.go @@ -273,6 +273,11 @@ func Commands(metaPtr *command.Meta) map[string]cli.CommandFactory { Meta: meta, }, nil }, + "node eligibility": func() (cli.Command, error) { + return &command.NodeEligibilityCommand{ + Meta: meta, + }, nil + }, "node-status": func() (cli.Command, error) { return &command.NodeStatusCommand{ Meta: meta, diff --git a/main.go b/main.go index 3c178e145..f482ca283 100644 --- a/main.go +++ b/main.go @@ -47,7 +47,7 @@ func RunCustom(args []string, commands map[string]cli.CommandFactory) int { // users should not be running should be placed here, versus hiding // subcommands from the main help, which should be filtered out of the // commands above. - hidden := []string{"check", "executor", "syslog"} + hidden := []string{"check", "executor", "syslog", "node-drain", "node-status"} cli := &cli.CLI{ Name: "nomad",