Drain cli, api, http

api/nodes.go

@@ -3,7 +3,6 @@ package api
 import (
 	"fmt"
 	"sort"
-	"strconv"
 	"time"
 )
 
@@ -42,10 +41,24 @@ func (n *Nodes) Info(nodeID string, q *QueryOptions) (*Node, *QueryMeta, error)
 	return &resp, qm, nil
 }
 
-// ToggleDrain is used to toggle drain mode on/off for a given node.
-func (n *Nodes) ToggleDrain(nodeID string, drain bool, q *WriteOptions) (*WriteMeta, error) {
-	drainArg := strconv.FormatBool(drain)
-	wm, err := n.client.write("/v1/node/"+nodeID+"/drain?enable="+drainArg, nil, nil, q)
+// NodeUpdateDrainRequest is used to update the drain specification for a node.
+type NodeUpdateDrainRequest struct {
+	// NodeID is the node to update the drain specification for.
+	NodeID string
+
+	// DrainSpec is the drain specification to set for the node. A nil DrainSpec
+	// will disable draining.
+	DrainSpec *DrainSpec
+}
+
+// UpdateDrain is used to update the drain strategy for a given node.
+func (n *Nodes) UpdateDrain(nodeID string, spec *DrainSpec, q *WriteOptions) (*WriteMeta, error) {
+	req := &NodeUpdateDrainRequest{
+		NodeID:    nodeID,
+		DrainSpec: spec,
+	}
+
+	wm, err := n.client.write("/v1/node/"+nodeID+"/drain", req, nil, q)
 	if err != nil {
 		return nil, err
 	}
@@ -108,25 +121,44 @@ type DriverInfo struct {
 
 // Node is used to deserialize a node entry.
 type Node struct {
-	ID                string
-	Datacenter        string
-	Name              string
-	HTTPAddr          string
-	TLSEnabled        bool
-	Attributes        map[string]string
-	Resources         *Resources
-	Reserved          *Resources
-	Links             map[string]string
-	Meta              map[string]string
-	NodeClass         string
-	Drain             bool
-	Status            string
-	StatusDescription string
-	StatusUpdatedAt   int64
-	Events            []*NodeEvent
-	Drivers           map[string]*DriverInfo
-	CreateIndex       uint64
-	ModifyIndex       uint64
+	ID                    string
+	Datacenter            string
+	Name                  string
+	HTTPAddr              string
+	TLSEnabled            bool
+	Attributes            map[string]string
+	Resources             *Resources
+	Reserved              *Resources
+	Links                 map[string]string
+	Meta                  map[string]string
+	NodeClass             string
+	Drain                 bool
+	DrainStrategy         *DrainStrategy
+	SchedulingEligibility string
+	Status                string
+	StatusDescription     string
+	StatusUpdatedAt       int64
+	Events                []*NodeEvent
+	Drivers               map[string]*DriverInfo
+	CreateIndex           uint64
+	ModifyIndex           uint64
+}
+
+// DrainStrategy describes a Node's drain behavior.
+type DrainStrategy struct {
+	// DrainSpec is the user declared drain specification
+	DrainSpec
+}
+
+// DrainSpec describes a Node's drain behavior.
+type DrainSpec struct {
+	// Deadline is the duration after StartTime when the remaining
+	// allocations on a draining Node should be told to stop.
+	Deadline time.Duration
+
+	// IgnoreSystemJobs allows systems jobs to remain on the node even though it
+	// has been marked for draining.
+	IgnoreSystemJobs bool
 }
 
 const (

api/nodes_test.go

@@ -174,7 +174,10 @@ func TestNodes_ToggleDrain(t *testing.T) {
 	}
 
 	// Toggle it on
-	wm, err := nodes.ToggleDrain(nodeID, true, nil)
+	spec := &DrainSpec{
+		Deadline: 10 * time.Second,
+	}
+	wm, err := nodes.UpdateDrain(nodeID, spec, nil)
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
@@ -185,12 +188,12 @@ func TestNodes_ToggleDrain(t *testing.T) {
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
-	if !out.Drain {
-		t.Fatalf("drain mode should be on")
+	if out.SchedulingEligibility != structs.NodeSchedulingIneligible {
+		t.Fatalf("bad eligibility: %v vs %v", out.SchedulingEligibility, structs.NodeSchedulingIneligible)
 	}
 
 	// Toggle off again
-	wm, err = nodes.ToggleDrain(nodeID, false, nil)
+	wm, err = nodes.UpdateDrain(nodeID, nil, nil)
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
@@ -204,6 +207,9 @@ func TestNodes_ToggleDrain(t *testing.T) {
 	if out.Drain {
 		t.Fatalf("drain mode should be off")
 	}
+	if out.DrainStrategy != nil {
+		t.Fatalf("drain strategy should be unset")
+	}
 }
 
 func TestNodes_Allocations(t *testing.T) {

command/agent/node_endpoint.go

@@ -2,9 +2,9 @@ package agent
 
 import (
 	"net/http"
-	"strconv"
 	"strings"
 
+	"github.com/hashicorp/nomad/api"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -101,19 +101,21 @@ func (s *HTTPServer) nodeToggleDrain(resp http.ResponseWriter, req *http.Request
 		return nil, CodedError(405, ErrInvalidMethod)
 	}
 
-	// Get the enable value
-	enableRaw := req.URL.Query().Get("enable")
-	if enableRaw == "" {
-		return nil, CodedError(400, "missing enable value")
-	}
-	enable, err := strconv.ParseBool(enableRaw)
-	if err != nil {
-		return nil, CodedError(400, "invalid enable value")
+	var drainRequest api.NodeUpdateDrainRequest
+	if err := decodeBody(req, &drainRequest); err != nil {
+		return nil, CodedError(400, err.Error())
 	}
 
 	args := structs.NodeUpdateDrainRequest{
 		NodeID: nodeID,
-		Drain:  enable,
+	}
+	if drainRequest.DrainSpec != nil {
+		args.DrainStrategy = &structs.DrainStrategy{
+			DrainSpec: structs.DrainSpec{
+				Deadline:         drainRequest.DrainSpec.Deadline,
+				IgnoreSystemJobs: drainRequest.DrainSpec.IgnoreSystemJobs,
+			},
+		}
 	}
 	s.parseWriteRequest(req, &args.WriteRequest)
 

command/agent/node_endpoint_test.go

@@ -4,10 +4,13 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
+	"github.com/hashicorp/nomad/api"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestHTTP_NodesList(t *testing.T) {
@@ -238,6 +241,7 @@ func TestHTTP_NodeAllocations(t *testing.T) {
 
 func TestHTTP_NodeDrain(t *testing.T) {
 	t.Parallel()
+	require := require.New(t)
 	httpTest(t, nil, func(s *TestAgent) {
 		// Create the node
 		node := mock.Node()
@@ -246,45 +250,55 @@ func TestHTTP_NodeDrain(t *testing.T) {
 			WriteRequest: structs.WriteRequest{Region: "global"},
 		}
 		var resp structs.NodeUpdateResponse
-		if err := s.Agent.RPC("Node.Register", &args, &resp); err != nil {
-			t.Fatalf("err: %v", err)
-		}
+		require.Nil(s.Agent.RPC("Node.Register", &args, &resp))
 
-		// Directly manipulate the state
-		state := s.Agent.server.State()
-		alloc1 := mock.Alloc()
-		alloc1.NodeID = node.ID
-		if err := state.UpsertJobSummary(999, mock.JobSummary(alloc1.JobID)); err != nil {
-			t.Fatal(err)
-		}
-		err := state.UpsertAllocs(1000, []*structs.Allocation{alloc1})
-		if err != nil {
-			t.Fatalf("err: %v", err)
+		drainReq := api.NodeUpdateDrainRequest{
+			NodeID: node.ID,
+			DrainSpec: &api.DrainSpec{
+				Deadline: 10 * time.Second,
+			},
 		}
 
 		// Make the HTTP request
-		req, err := http.NewRequest("POST", "/v1/node/"+node.ID+"/drain?enable=1", nil)
-		if err != nil {
-			t.Fatalf("err: %v", err)
-		}
+		buf := encodeReq(drainReq)
+		req, err := http.NewRequest("POST", "/v1/node/"+node.ID+"/drain", buf)
+		require.Nil(err)
 		respW := httptest.NewRecorder()
 
 		// Make the request
 		obj, err := s.Server.NodeSpecificRequest(respW, req)
-		if err != nil {
-			t.Fatalf("err: %v", err)
-		}
+		require.Nil(err)
 
 		// Check for the index
-		if respW.HeaderMap.Get("X-Nomad-Index") == "" {
-			t.Fatalf("missing index")
-		}
+		require.NotZero(respW.HeaderMap.Get("X-Nomad-Index"))
 
 		// Check the response
-		upd := obj.(structs.NodeDrainUpdateResponse)
-		if len(upd.EvalIDs) == 0 {
-			t.Fatalf("bad: %v", upd)
-		}
+		_, ok := obj.(structs.NodeDrainUpdateResponse)
+		require.True(ok)
+
+		// Check that the node has been updated
+		state := s.Agent.server.State()
+		out, err := state.NodeByID(nil, node.ID)
+		require.Nil(err)
+		require.True(out.Drain)
+		require.NotNil(out.DrainStrategy)
+		require.Equal(10*time.Second, out.DrainStrategy.Deadline)
+
+		// Make the HTTP request to unset drain
+		drainReq.DrainSpec = nil
+		buf = encodeReq(drainReq)
+		req, err = http.NewRequest("POST", "/v1/node/"+node.ID+"/drain", buf)
+		require.Nil(err)
+		respW = httptest.NewRecorder()
+
+		// Make the request
+		obj, err = s.Server.NodeSpecificRequest(respW, req)
+		require.Nil(err)
+
+		out, err = state.NodeByID(nil, node.ID)
+		require.Nil(err)
+		require.False(out.Drain)
+		require.Nil(out.DrainStrategy)
 	})
 }
 

command/node.go

@@ -0,0 +1,19 @@
+package command
+
+import "github.com/mitchellh/cli"
+
+type NodeCommand struct {
+	Meta
+}
+
+func (f *NodeCommand) Help() string {
+	return "This command is accessed by using one of the subcommands below."
+}
+
+func (f *NodeCommand) Synopsis() string {
+	return "Interact with nodes"
+}
+
+func (f *NodeCommand) Run(args []string) int {
+	return cli.RunResultHelp
+}

command/node_drain.go

@@ -3,18 +3,26 @@ package command
 import (
 	"fmt"
 	"strings"
+	"time"
 
+	"github.com/hashicorp/nomad/api"
 	"github.com/hashicorp/nomad/api/contexts"
 	"github.com/posener/complete"
 )
 
+var (
+	// defaultDrainDuration is the default drain duration if it is not specified
+	// explicitly
+	defaultDrainDuration = 1 * time.Hour
+)
+
 type NodeDrainCommand struct {
 	Meta
 }
 
 func (c *NodeDrainCommand) Help() string {
 	helpText := `
-Usage: nomad node-drain [options] <node>
+Usage: nomad node drain [options] <node>
 
   Toggles node draining on a specified node. It is required
   that either -enable or -disable is specified, but not both.
@@ -32,8 +40,24 @@ Node Drain Options:
   -enable
     Enable draining for the specified node.
 
+  -deadline <duration>
+    Set the deadline by which all allocations must be moved off the node.
+    Remaining allocations after the deadline are forced removed from the node.
+    If unspecified, a default deadline of one hour is applied.
+
+  -force
+    Force remove allocations off the node immediately.
+
+  -no-deadline
+    No deadline allows the allocations to drain off the node without being force
+    stopped after a certain deadline.
+
+  -ignore-system
+    Ignore system allows the drain to complete without stopping system job
+    allocations.
+
   -self
-    Query the status of the local node.
+    Set the drain status of the local node.
 
   -yes
     Automatic yes to prompts.
@@ -48,10 +72,14 @@ func (c *NodeDrainCommand) Synopsis() string {
 func (c *NodeDrainCommand) AutocompleteFlags() complete.Flags {
 	return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient),
 		complete.Flags{
-			"-disable": complete.PredictNothing,
-			"-enable":  complete.PredictNothing,
-			"-self":    complete.PredictNothing,
-			"-yes":     complete.PredictNothing,
+			"-disable":       complete.PredictNothing,
+			"-enable":        complete.PredictNothing,
+			"-deadline":      complete.PredictAnything,
+			"-force":         complete.PredictNothing,
+			"-no-deadline":   complete.PredictNothing,
+			"-ignore-system": complete.PredictNothing,
+			"-self":          complete.PredictNothing,
+			"-yes":           complete.PredictNothing,
 		})
 }
 
@@ -71,12 +99,18 @@ func (c *NodeDrainCommand) AutocompleteArgs() complete.Predictor {
 }
 
 func (c *NodeDrainCommand) Run(args []string) int {
-	var enable, disable, self, autoYes bool
+	var enable, disable, force,
+		noDeadline, ignoreSystem, self, autoYes bool
+	var deadline string
 
 	flags := c.Meta.FlagSet("node-drain", FlagSetClient)
 	flags.Usage = func() { c.Ui.Output(c.Help()) }
 	flags.BoolVar(&enable, "enable", false, "Enable drain mode")
 	flags.BoolVar(&disable, "disable", false, "Disable drain mode")
+	flags.StringVar(&deadline, "deadline", "", "Deadline after which allocations are force stopped")
+	flags.BoolVar(&force, "force", false, "Force immediate drain")
+	flags.BoolVar(&noDeadline, "no-deadline", false, "Drain node with no deadline")
+	flags.BoolVar(&ignoreSystem, "ignore-system", false, "Do not drain system job allocations from the node")
 	flags.BoolVar(&self, "self", false, "")
 	flags.BoolVar(&autoYes, "yes", false, "Automatic yes to prompts.")
 
@@ -93,10 +127,46 @@ func (c *NodeDrainCommand) Run(args []string) int {
 	// Check that we got a node ID
 	args = flags.Args()
 	if l := len(args); self && l != 0 || !self && l != 1 {
-		c.Ui.Error(c.Help())
+		c.Ui.Error("Node ID must be specified if -self isn't being used")
 		return 1
 	}
 
+	// Validate a compatible set of flags were set
+	if disable && (deadline != "" || force || noDeadline || ignoreSystem) {
+		c.Ui.Error("-disable can't be combined with flags configuring drain strategy")
+		return 1
+	}
+	if deadline != "" && (force || noDeadline) {
+		c.Ui.Error("-deadline can't be combined with -force or -no-deadline")
+		return 1
+	}
+	if force && noDeadline {
+		c.Ui.Error("-force and -no-deadline are mutually exclusive")
+		return 1
+	}
+
+	// Parse the duration
+	var d time.Duration
+	if force {
+		d = -1 * time.Second
+	} else if noDeadline {
+		d = 0
+	} else if deadline != "" {
+		dur, err := time.ParseDuration(deadline)
+		if err != nil {
+			c.Ui.Error(fmt.Sprintf("Failed to parse deadline %q: %v", deadline, err))
+			return 1
+		}
+		if dur <= 0 {
+			c.Ui.Error("A positive drain duration must be given")
+			return 1
+		}
+
+		d = dur
+	} else {
+		d = defaultDrainDuration
+	}
+
 	// Get the HTTP client
 	client, err := c.Meta.Client()
 	if err != nil {
@@ -186,9 +256,17 @@ func (c *NodeDrainCommand) Run(args []string) int {
 		}
 	}
 
+	var spec *api.DrainSpec
+	if enable {
+		spec = &api.DrainSpec{
+			Deadline:         d,
+			IgnoreSystemJobs: ignoreSystem,
+		}
+	}
+
 	// Toggle node draining
-	if _, err := client.Nodes().ToggleDrain(node.ID, enable, nil); err != nil {
-		c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
+	if _, err := client.Nodes().UpdateDrain(node.ID, spec, nil); err != nil {
+		c.Ui.Error(fmt.Sprintf("Error updating drain specification: %s", err))
 		return 1
 	}
 	return 0

command/node_drain_test.go

@@ -85,6 +85,49 @@ func TestNodeDrainCommand_Fails(t *testing.T) {
 	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()
+
+	// Fail on disable being used with drain strategy flags
+	for _, flag := range []string{"-force", "-no-deadline", "-ignore-system"} {
+		if code := cmd.Run([]string{"-address=" + url, "-disable", flag, "12345678-abcd-efab-cdef-123456789abc"}); code != 1 {
+			t.Fatalf("expected exit 1, got: %d", code)
+		}
+		if out := ui.ErrorWriter.String(); !strings.Contains(out, "combined with flags configuring drain strategy") {
+			t.Fatalf("got: %s", out)
+		}
+		ui.ErrorWriter.Reset()
+	}
+
+	// Fail on setting a deadline plus deadline modifying flags
+	for _, flag := range []string{"-force", "-no-deadline"} {
+		if code := cmd.Run([]string{"-address=" + url, "-enable", "-deadline=10s", flag, "12345678-abcd-efab-cdef-123456789abc"}); code != 1 {
+			t.Fatalf("expected exit 1, got: %d", code)
+		}
+		if out := ui.ErrorWriter.String(); !strings.Contains(out, "deadline can't be combined with") {
+			t.Fatalf("got: %s", out)
+		}
+		ui.ErrorWriter.Reset()
+	}
+
+	// Fail on setting a force and no deadline
+	if code := cmd.Run([]string{"-address=" + url, "-enable", "-force", "-no-deadline", "12345678-abcd-efab-cdef-123456789abc"}); code != 1 {
+		t.Fatalf("expected exit 1, got: %d", code)
+	}
+	if out := ui.ErrorWriter.String(); !strings.Contains(out, "mutually exclusive") {
+		t.Fatalf("got: %s", out)
+	}
+	ui.ErrorWriter.Reset()
+
+	// Fail on setting a bad deadline
+	for _, flag := range []string{"-deadline=0s", "-deadline=-1s"} {
+		if code := cmd.Run([]string{"-address=" + url, "-enable", flag, "12345678-abcd-efab-cdef-123456789abc"}); code != 1 {
+			t.Fatalf("expected exit 1, got: %d", code)
+		}
+		if out := ui.ErrorWriter.String(); !strings.Contains(out, "positive") {
+			t.Fatalf("got: %s", out)
+		}
+		ui.ErrorWriter.Reset()
+	}
 }
 
 func TestNodeDrainCommand_AutocompleteArgs(t *testing.T) {

command/node_status.go

@@ -37,7 +37,7 @@ type NodeStatusCommand struct {
 
 func (c *NodeStatusCommand) Help() string {
 	helpText := `
-Usage: nomad node-status [options] <node>
+Usage: nomad node status [options] <node>
 
   Display status information about a given node. The list of nodes
   returned includes only nodes which jobs may be scheduled to, and

commands.go

@@ -258,17 +258,31 @@ func Commands(metaPtr *command.Meta) map[string]cli.CommandFactory {
 				Meta: meta,
 			}, nil
 		},
+		"node": func() (cli.Command, error) {
+			return &command.NodeCommand{
+				Meta: meta,
+			}, nil
+		},
 		"node-drain": func() (cli.Command, error) {
 			return &command.NodeDrainCommand{
 				Meta: meta,
 			}, nil
 		},
+		"node drain": func() (cli.Command, error) {
+			return &command.NodeDrainCommand{
+				Meta: meta,
+			}, nil
+		},
 		"node-status": func() (cli.Command, error) {
 			return &command.NodeStatusCommand{
 				Meta: meta,
 			}, nil
 		},
-
+		"node status": func() (cli.Command, error) {
+			return &command.NodeStatusCommand{
+				Meta: meta,
+			}, nil
+		},
 		"operator": func() (cli.Command, error) {
 			return &command.OperatorCommand{
 				Meta: meta,

main.go

@@ -37,6 +37,7 @@ func RunCustom(args []string, commands map[string]cli.CommandFactory) int {
 		case "quota list", "quota delete", "quota apply", "quota status", "quota inspect", "quota init":
 		case "operator raft", "operator raft list-peers", "operator raft remove-peer":
 		case "acl policy", "acl policy apply", "acl token", "acl token create":
+		case "node-drain", "node-status":
 		default:
 			commandsInclude = append(commandsInclude, k)
 		}

nomad/structs/structs.go

@@ -1191,12 +1191,8 @@ const (
 	NodeSchedulingIneligible = "ineligible"
 )
 
-// DrainStrategy describes a Node's drain behavior.
-type DrainStrategy struct {
-	// StartTime as nanoseconds since Unix epoch indicating when a drain
-	// began for deadline calcuations.
-	StartTime int64
-
+// DrainSpec describes a Node's desired drain behavior.
+type DrainSpec struct {
 	// Deadline is the duration after StartTime when the remaining
 	// allocations on a draining Node should be told to stop.
 	Deadline time.Duration
@@ -1206,6 +1202,16 @@ type DrainStrategy struct {
 	IgnoreSystemJobs bool
 }
 
+// DrainStrategy describes a Node's drain behavior.
+type DrainStrategy struct {
+	// DrainSpec is the user declared drain specification
+	DrainSpec
+
+	// StartTime as nanoseconds since Unix epoch indicating when a drain
+	// began for deadline calcuations.
+	StartTime int64
+}
+
 func (d *DrainStrategy) Copy() *DrainStrategy {
 	if d == nil {
 		return nil