1package command
2
3import (
4	"context"
5	"fmt"
6	"strings"
7	"time"
8
9	"github.com/hashicorp/nomad/api"
10	"github.com/hashicorp/nomad/api/contexts"
11	"github.com/posener/complete"
12)
13
14var (
15	// defaultDrainDuration is the default drain duration if it is not specified
16	// explicitly
17	defaultDrainDuration = 1 * time.Hour
18)
19
20type NodeDrainCommand struct {
21	Meta
22}
23
24func (c *NodeDrainCommand) Help() string {
25	helpText := `
26Usage: nomad node drain [options] <node>
27
28  Toggles node draining on a specified node. It is required
29  that either -enable or -disable is specified, but not both.
30  The -self flag is useful to drain the local node.
31
32General Options:
33
34  ` + generalOptionsUsage() + `
35
36Node Drain Options:
37
38  -disable
39    Disable draining for the specified node.
40
41  -enable
42    Enable draining for the specified node.
43
44  -deadline <duration>
45    Set the deadline by which all allocations must be moved off the node.
46    Remaining allocations after the deadline are forced removed from the node.
47    If unspecified, a default deadline of one hour is applied.
48
49  -detach
50    Return immediately instead of entering monitor mode.
51
52  -monitor
53    Enter monitor mode directly without modifying the drain status.
54
55  -force
56    Force remove allocations off the node immediately.
57
58  -no-deadline
59    No deadline allows the allocations to drain off the node without being force
60    stopped after a certain deadline.
61
62  -ignore-system
63    Ignore system allows the drain to complete without stopping system job
64    allocations. By default system jobs are stopped last.
65
66  -keep-ineligible
67    Keep ineligible will maintain the node's scheduling ineligibility even if
68    the drain is being disabled. This is useful when an existing drain is being
69    cancelled but additional scheduling on the node is not desired.
70
71  -self
72    Set the drain status of the local node.
73
74  -yes
75    Automatic yes to prompts.
76`
77	return strings.TrimSpace(helpText)
78}
79
80func (c *NodeDrainCommand) Synopsis() string {
81	return "Toggle drain mode on a given node"
82}
83
84func (c *NodeDrainCommand) AutocompleteFlags() complete.Flags {
85	return mergeAutocompleteFlags(c.Meta.AutocompleteFlags(FlagSetClient),
86		complete.Flags{
87			"-disable":         complete.PredictNothing,
88			"-enable":          complete.PredictNothing,
89			"-deadline":        complete.PredictAnything,
90			"-detach":          complete.PredictNothing,
91			"-force":           complete.PredictNothing,
92			"-no-deadline":     complete.PredictNothing,
93			"-ignore-system":   complete.PredictNothing,
94			"-keep-ineligible": complete.PredictNothing,
95			"-self":            complete.PredictNothing,
96			"-yes":             complete.PredictNothing,
97		})
98}
99
100func (c *NodeDrainCommand) AutocompleteArgs() complete.Predictor {
101	return complete.PredictFunc(func(a complete.Args) []string {
102		client, err := c.Meta.Client()
103		if err != nil {
104			return nil
105		}
106
107		resp, _, err := client.Search().PrefixSearch(a.Last, contexts.Nodes, nil)
108		if err != nil {
109			return []string{}
110		}
111		return resp.Matches[contexts.Nodes]
112	})
113}
114
115func (c *NodeDrainCommand) Name() string { return "node-drain" }
116
117func (c *NodeDrainCommand) Run(args []string) int {
118	var enable, disable, detach, force,
119		noDeadline, ignoreSystem, keepIneligible,
120		self, autoYes, monitor bool
121	var deadline string
122
123	flags := c.Meta.FlagSet(c.Name(), FlagSetClient)
124	flags.Usage = func() { c.Ui.Output(c.Help()) }
125	flags.BoolVar(&enable, "enable", false, "Enable drain mode")
126	flags.BoolVar(&disable, "disable", false, "Disable drain mode")
127	flags.StringVar(&deadline, "deadline", "", "Deadline after which allocations are force stopped")
128	flags.BoolVar(&detach, "detach", false, "")
129	flags.BoolVar(&force, "force", false, "Force immediate drain")
130	flags.BoolVar(&noDeadline, "no-deadline", false, "Drain node with no deadline")
131	flags.BoolVar(&ignoreSystem, "ignore-system", false, "Do not drain system job allocations from the node")
132	flags.BoolVar(&keepIneligible, "keep-ineligible", false, "Do not update the nodes scheduling eligibility")
133	flags.BoolVar(&self, "self", false, "")
134	flags.BoolVar(&autoYes, "yes", false, "Automatic yes to prompts.")
135	flags.BoolVar(&monitor, "monitor", false, "Monitor drain status.")
136
137	if err := flags.Parse(args); err != nil {
138		return 1
139	}
140
141	// Check that enable or disable is not set with monitor
142	if monitor && (enable || disable) {
143		c.Ui.Error("The -monitor flag cannot be used with the '-enable' or '-disable' flags")
144		c.Ui.Error(commandErrorText(c))
145		return 1
146	}
147
148	// Check that we got either enable or disable, but not both.
149	if (enable && disable) || (!monitor && !enable && !disable) {
150		c.Ui.Error("Ethier the '-enable' or '-disable' flag must be set, unless using '-monitor'")
151		c.Ui.Error(commandErrorText(c))
152		return 1
153	}
154
155	// Check that we got a node ID
156	args = flags.Args()
157	if l := len(args); self && l != 0 || !self && l != 1 {
158		c.Ui.Error("Node ID must be specified if -self isn't being used")
159		c.Ui.Error(commandErrorText(c))
160		return 1
161	}
162
163	// Validate a compatible set of flags were set
164	if disable && (deadline != "" || force || noDeadline || ignoreSystem) {
165		c.Ui.Error("-disable can't be combined with flags configuring drain strategy")
166		c.Ui.Error(commandErrorText(c))
167		return 1
168	}
169	if deadline != "" && (force || noDeadline) {
170		c.Ui.Error("-deadline can't be combined with -force or -no-deadline")
171		c.Ui.Error(commandErrorText(c))
172		return 1
173	}
174	if force && noDeadline {
175		c.Ui.Error("-force and -no-deadline are mutually exclusive")
176		c.Ui.Error(commandErrorText(c))
177		return 1
178	}
179
180	// Parse the duration
181	var d time.Duration
182	if force {
183		d = -1 * time.Second
184	} else if noDeadline {
185		d = 0
186	} else if deadline != "" {
187		dur, err := time.ParseDuration(deadline)
188		if err != nil {
189			c.Ui.Error(fmt.Sprintf("Failed to parse deadline %q: %v", deadline, err))
190			return 1
191		}
192		if dur <= 0 {
193			c.Ui.Error("A positive drain duration must be given")
194			return 1
195		}
196
197		d = dur
198	} else {
199		d = defaultDrainDuration
200	}
201
202	// Get the HTTP client
203	client, err := c.Meta.Client()
204	if err != nil {
205		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
206		return 1
207	}
208
209	// If -self flag is set then determine the current node.
210	var nodeID string
211	if !self {
212		nodeID = args[0]
213	} else {
214		var err error
215		if nodeID, err = getLocalNodeID(client); err != nil {
216			c.Ui.Error(err.Error())
217			return 1
218		}
219	}
220
221	// Check if node exists
222	if len(nodeID) == 1 {
223		c.Ui.Error(fmt.Sprintf("Identifier must contain at least two characters."))
224		return 1
225	}
226
227	nodeID = sanitizeUUIDPrefix(nodeID)
228	nodes, _, err := client.Nodes().PrefixList(nodeID)
229	if err != nil {
230		c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
231		return 1
232	}
233	// Return error if no nodes are found
234	if len(nodes) == 0 {
235		c.Ui.Error(fmt.Sprintf("No node(s) with prefix or id %q found", nodeID))
236		return 1
237	}
238	if len(nodes) > 1 {
239		c.Ui.Error(fmt.Sprintf("Prefix matched multiple nodes\n\n%s",
240			formatNodeStubList(nodes, true)))
241		return 1
242	}
243
244	// Prefix lookup matched a single node
245	node, meta, err := client.Nodes().Info(nodes[0].ID, nil)
246	if err != nil {
247		c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
248		return 1
249	}
250
251	// If monitoring the drain start the montior and return when done
252	if monitor {
253		if node.DrainStrategy == nil {
254			c.Ui.Warn("No drain strategy set")
255			return 0
256		}
257		c.Ui.Info(fmt.Sprintf("%s: Monitoring node %q: Ctrl-C to detach monitoring", formatTime(time.Now()), node.ID))
258		c.monitorDrain(client, context.Background(), node, meta.LastIndex, ignoreSystem)
259		return 0
260	}
261
262	// Confirm drain if the node was a prefix match.
263	if nodeID != node.ID && !autoYes {
264		verb := "enable"
265		if disable {
266			verb = "disable"
267		}
268		question := fmt.Sprintf("Are you sure you want to %s drain mode for node %q? [y/N]", verb, node.ID)
269		answer, err := c.Ui.Ask(question)
270		if err != nil {
271			c.Ui.Error(fmt.Sprintf("Failed to parse answer: %v", err))
272			return 1
273		}
274
275		if answer == "" || strings.ToLower(answer)[0] == 'n' {
276			// No case
277			c.Ui.Output("Canceling drain toggle")
278			return 0
279		} else if strings.ToLower(answer)[0] == 'y' && len(answer) > 1 {
280			// Non exact match yes
281			c.Ui.Output("For confirmation, an exact ‘y’ is required.")
282			return 0
283		} else if answer != "y" {
284			c.Ui.Output("No confirmation detected. For confirmation, an exact 'y' is required.")
285			return 1
286		}
287	}
288
289	var spec *api.DrainSpec
290	if enable {
291		spec = &api.DrainSpec{
292			Deadline:         d,
293			IgnoreSystemJobs: ignoreSystem,
294		}
295	}
296
297	// Toggle node draining
298	updateMeta, err := client.Nodes().UpdateDrain(node.ID, spec, !keepIneligible, nil)
299	if err != nil {
300		c.Ui.Error(fmt.Sprintf("Error updating drain specification: %s", err))
301		return 1
302	}
303
304	if !enable || detach {
305		if enable {
306			c.Ui.Output(fmt.Sprintf("Node %q drain strategy set", node.ID))
307		} else {
308			c.Ui.Output(fmt.Sprintf("Node %q drain strategy unset", node.ID))
309		}
310	}
311
312	if enable && !detach {
313		now := time.Now()
314		c.Ui.Info(fmt.Sprintf("%s: Ctrl-C to stop monitoring: will not cancel the node drain", formatTime(now)))
315		c.Ui.Output(fmt.Sprintf("%s: Node %q drain strategy set", formatTime(now), node.ID))
316		c.monitorDrain(client, context.Background(), node, updateMeta.LastIndex, ignoreSystem)
317	}
318	return 0
319}
320
321func (c *NodeDrainCommand) monitorDrain(client *api.Client, ctx context.Context, node *api.Node, index uint64, ignoreSystem bool) {
322	outCh := client.Nodes().MonitorDrain(ctx, node.ID, index, ignoreSystem)
323	for msg := range outCh {
324		switch msg.Level {
325		case api.MonitorMsgLevelInfo:
326			c.Ui.Info(fmt.Sprintf("%s: %s", formatTime(time.Now()), msg))
327		case api.MonitorMsgLevelWarn:
328			c.Ui.Warn(fmt.Sprintf("%s: %s", formatTime(time.Now()), msg))
329		case api.MonitorMsgLevelError:
330			c.Ui.Error(fmt.Sprintf("%s: %s", formatTime(time.Now()), msg))
331		default:
332			c.Ui.Output(fmt.Sprintf("%s: %s", formatTime(time.Now()), msg))
333		}
334	}
335}
336