forked from hashicorp/nomad
-
Notifications
You must be signed in to change notification settings - Fork 2
/
node_drain.go
127 lines (108 loc) · 3.04 KB
/
node_drain.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
package command
import (
"fmt"
"strings"
)
type NodeDrainCommand struct {
Meta
}
func (c *NodeDrainCommand) Help() string {
helpText := `
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.
General Options:
` + generalOptionsUsage() + `
Node Drain Options:
-disable
Disable draining for the specified node.
-enable
Enable draining for the specified node.
`
return strings.TrimSpace(helpText)
}
func (c *NodeDrainCommand) Synopsis() string {
return "Toggle drain mode on a given node"
}
func (c *NodeDrainCommand) Run(args []string) int {
var enable, disable bool
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")
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 len(args) != 1 {
c.Ui.Error(c.Help())
return 1
}
nodeID := args[0]
// Get the HTTP client
client, err := c.Meta.Client()
if err != nil {
c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
return 1
}
// Check if node exists
node, _, err := client.Nodes().Info(nodeID, nil)
if err != nil {
if len(nodeID) == 1 {
c.Ui.Error(fmt.Sprintf("Identifier must contain at least two characters."))
return 1
}
if len(nodeID)%2 == 1 {
// Identifiers must be of even length, so we strip off the last byte
// to provide a consistent user experience.
nodeID = nodeID[:len(nodeID)-1]
}
// Exact lookup failed, try with prefix based search
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.Output(fmt.Sprintf("Prefix matched multiple nodes\n\n%s", formatList(out)))
return 0
}
// 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 draining
if _, err := client.Nodes().ToggleDrain(node.ID, enable, nil); err != nil {
c.Ui.Error(fmt.Sprintf("Error toggling drain mode: %s", err))
return 1
}
return 0
}