agent: node maintenance mode works

command/agent/agent.go

@@ -30,8 +30,9 @@ const (
 		"be left behind. If the path looks correct, remove the file " +
 		"and try again."
 
-	// The ID of the faux health check for maintenance mode
+	// The ID of the faux health checks for maintenance mode
-	maintCheckID = "_maintenance_"
+	serviceMaintCheckID = "_service_maintenance"
+	nodeMaintCheckID    = "_node_maintenenace"
 )
 
 /*
@@ -1008,14 +1009,14 @@ func (a *Agent) EnableServiceMaintenance(serviceID string) error {
 	}
 
 	// Ensure maintenance mode is not already enabled
-	if _, ok := a.state.Checks()[maintCheckID]; ok {
+	if _, ok := a.state.Checks()[serviceMaintCheckID]; ok {
 		return nil
 	}
 
 	// Create and register the critical health check
 	check := &structs.HealthCheck{
 		Node:        a.config.NodeName,
-		CheckID:     maintCheckID,
+		CheckID:     serviceMaintCheckID,
 		Name:        "Service Maintenance Mode",
 		Notes:       "Maintenance mode is enabled for this service",
 		ServiceID:   service.ID,
@@ -1035,6 +1036,29 @@ func (a *Agent) DisableServiceMaintenance(serviceID string) error {
 	}
 
 	// Deregister the maintenance check
-	a.RemoveCheck(maintCheckID, true)
+	a.RemoveCheck(serviceMaintCheckID, true)
 	return nil
 }
+
+// EnableNodeMaintenance places a node into maintenance mode.
+func (a *Agent) EnableNodeMaintenance() {
+	// Ensure node maintenance is not already enabled
+	if _, ok := a.state.Checks()[nodeMaintCheckID]; ok {
+		return
+	}
+
+	// Create and register the node maintenance check
+	check := &structs.HealthCheck{
+		Node:    a.config.NodeName,
+		CheckID: nodeMaintCheckID,
+		Name:    "Node Maintenance Mode",
+		Notes:   "Maintenance mode is enabled for this node",
+		Status:  structs.HealthCritical,
+	}
+	a.AddCheck(check, nil, true)
+}
+
+// DisableNodeMaintenance removes a node from maintenance mode
+func (a *Agent) DisableNodeMaintenance() {
+	a.RemoveCheck(nodeMaintCheckID, true)
+}

command/agent/agent_endpoint.go

@@ -227,3 +227,39 @@ func (s *HTTPServer) AgentServiceMaintenance(resp http.ResponseWriter, req *http
 	}
 	return nil, err
 }
+
+func (s *HTTPServer) AgentNodeMaintenance(resp http.ResponseWriter, req *http.Request) (interface{}, error) {
+	// Only PUT supported
+	if req.Method != "PUT" {
+		resp.WriteHeader(405)
+		return nil, nil
+	}
+
+	// Ensure we have some action
+	params := req.URL.Query()
+	if _, ok := params["enable"]; !ok {
+		resp.WriteHeader(400)
+		resp.Write([]byte("Missing value for enable"))
+		return nil, nil
+	}
+
+	var enable bool
+	raw := params.Get("enable")
+	switch raw {
+	case "true":
+		enable = true
+	case "false":
+		enable = false
+	default:
+		resp.WriteHeader(400)
+		resp.Write([]byte(fmt.Sprintf("Invalid value for enable: %q", raw)))
+		return nil, nil
+	}
+
+	if enable {
+		s.agent.EnableNodeMaintenance()
+	} else {
+		s.agent.DisableNodeMaintenance()
+	}
+	return nil, nil
+}

command/agent/agent_endpoint_test.go

@@ -567,7 +567,7 @@ func TestHTTPAgent_EnableServiceMaintenance(t *testing.T) {
 	}
 
 	// Ensure the maintenance check was registered
-	if _, ok := srv.agent.state.Checks()[maintCheckID]; !ok {
+	if _, ok := srv.agent.state.Checks()[serviceMaintCheckID]; !ok {
 		t.Fatalf("should have registered maintenance check")
 	}
 }
@@ -603,7 +603,7 @@ func TestHTTPAgent_DisableServiceMaintenance(t *testing.T) {
 	}
 
 	// Ensure the maintenance check was removed
-	if _, ok := srv.agent.state.Checks()[maintCheckID]; ok {
+	if _, ok := srv.agent.state.Checks()[serviceMaintCheckID]; ok {
 		t.Fatalf("should have removed maintenance check")
 	}
 }

command/agent/agent_test.go

@@ -807,7 +807,7 @@ func TestAgent_MaintenanceMode(t *testing.T) {
 
 	// Make sure the critical health check was added
 	for _, check := range agent.state.Checks() {
-		if check.CheckID == maintCheckID {
+		if check.CheckID == serviceMaintCheckID {
 			return
 		}
 	}
@@ -822,7 +822,7 @@ func TestAgent_MaintenanceMode(t *testing.T) {
 
 	// Ensure the check was deregistered
 	for _, check := range agent.state.Checks() {
-		if check.CheckID == maintCheckID {
+		if check.CheckID == serviceMaintCheckID {
 			t.Fatalf("should have deregistered maintenance check")
 		}
 	}

command/agent/http.go

@@ -181,6 +181,7 @@ func (s *HTTPServer) registerHandlers(enableDebug bool) {
 	s.mux.HandleFunc("/v1/health/service/", s.wrap(s.HealthServiceNodes))
 
 	s.mux.HandleFunc("/v1/agent/self", s.wrap(s.AgentSelf))
+	s.mux.HandleFunc("/v1/agent/self/maintenance", s.wrap(s.AgentNodeMaintenance))
 	s.mux.HandleFunc("/v1/agent/services", s.wrap(s.AgentServices))
 	s.mux.HandleFunc("/v1/agent/checks", s.wrap(s.AgentChecks))
 	s.mux.HandleFunc("/v1/agent/members", s.wrap(s.AgentMembers))