Drain cli, api, http

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