WIP: Smart Node Drain

api/allocations.go

@@ -81,6 +81,7 @@ type Allocation struct {
 	Metrics            *AllocationMetric
 	DesiredStatus      string
 	DesiredDescription string
+	DesiredTransition  DesiredTransition
 	ClientStatus       string
 	ClientDescription  string
 	TaskStates         map[string]*TaskState
@@ -205,3 +206,17 @@ type RescheduleEvent struct {
 	// PrevNodeID is the node ID of the previous allocation
 	PrevNodeID string
 }
+
+// DesiredTransition is used to mark an allocation as having a desired state
+// transition. This information can be used by the scheduler to make the
+// correct decision.
+type DesiredTransition struct {
+	// Migrate is used to indicate that this allocation should be stopped and
+	// migrated to another node.
+	Migrate *bool
+}
+
+// ShouldMigrate returns whether the transition object dictates a migration.
+func (d DesiredTransition) ShouldMigrate() bool {
+	return d.Migrate != nil && *d.Migrate
+}

api/allocations_test.go

@@ -239,3 +239,10 @@ func TestAllocations_RescheduleInfo(t *testing.T) {
 	}
 
 }
+
+func TestAllocations_ShouldMigrate(t *testing.T) {
+	t.Parallel()
+	require.True(t, DesiredTransition{Migrate: helper.BoolToPtr(true)}.ShouldMigrate())
+	require.False(t, DesiredTransition{}.ShouldMigrate())
+	require.False(t, DesiredTransition{Migrate: helper.BoolToPtr(false)}.ShouldMigrate())
+}

api/jobs.go

@@ -559,6 +559,7 @@ type Job struct {
 	ParameterizedJob  *ParameterizedJobConfig
 	Payload           []byte
 	Reschedule        *ReschedulePolicy
+	Migrate           *MigrateStrategy
 	Meta              map[string]string
 	VaultToken        *string `mapstructure:"vault_token"`
 	Status            *string

api/nodes.go

@@ -3,7 +3,9 @@ package api
 import (
 	"fmt"
 	"sort"
-	"strconv"
+	"time"
+
+	"github.com/hashicorp/nomad/nomad/structs"
 )
 
 // Nodes is used to query node-related API endpoints
@@ -41,10 +43,56 @@ 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
+
+	// MarkEligible marks the node as eligible if removing the drain strategy.
+	MarkEligible bool
+}
+
+// UpdateDrain is used to update the drain strategy for a given node. If
+// markEligible is true and the drain is being removed, the node will be marked
+// as having its scheduling being elibile
+func (n *Nodes) UpdateDrain(nodeID string, spec *DrainSpec, markEligible bool, q *WriteOptions) (*WriteMeta, error) {
+	req := &NodeUpdateDrainRequest{
+		NodeID:       nodeID,
+		DrainSpec:    spec,
+		MarkEligible: markEligible,
+	}
+
+	wm, err := n.client.write("/v1/node/"+nodeID+"/drain", req, nil, q)
+	if err != nil {
+		return nil, err
+	}
+	return wm, nil
+}
+
+// NodeUpdateEligibilityRequest is used to update the drain specification for a node.
+type NodeUpdateEligibilityRequest struct {
+	// NodeID is the node to update the drain specification for.
+	NodeID      string
+	Eligibility string
+}
+
+// ToggleEligibility is used to update the scheduling eligibility of the node
+func (n *Nodes) ToggleEligibility(nodeID string, eligible bool, q *WriteOptions) (*WriteMeta, error) {
+	e := structs.NodeSchedulingEligible
+	if !eligible {
+		e = structs.NodeSchedulingIneligible
+	}
+
+	req := &NodeUpdateEligibilityRequest{
+		NodeID:      nodeID,
+		Eligibility: e,
+	}
+
+	wm, err := n.client.write("/v1/node/"+nodeID+"/eligibility", req, nil, q)
 	if err != nil {
 		return nil, err
 	}
@@ -98,24 +146,47 @@ func (n *Nodes) GcAlloc(allocID string, q *QueryOptions) error {
 
 // 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
-	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
+	CreateIndex           uint64
+	ModifyIndex           uint64
+}
+
+// DrainStrategy describes a Node's drain behavior.
+type DrainStrategy struct {
+	// DrainSpec is the user declared drain specification
+	DrainSpec
+
+	// ForceDeadline is the deadline time for the drain after which drains will
+	// be forced
+	ForceDeadline time.Time
+}
+
+// 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 (
@@ -170,17 +241,18 @@ type HostDiskStats struct {
 // NodeListStub is a subset of information returned during
 // node list operations.
 type NodeListStub struct {
-	Address           string
-	ID                string
-	Datacenter        string
-	Name              string
-	NodeClass         string
-	Version           string
-	Drain             bool
-	Status            string
-	StatusDescription string
-	CreateIndex       uint64
-	ModifyIndex       uint64
+	Address               string
+	ID                    string
+	Datacenter            string
+	Name                  string
+	NodeClass             string
+	Version               string
+	Drain                 bool
+	SchedulingEligibility string
+	Status                string
+	StatusDescription     string
+	CreateIndex           uint64
+	ModifyIndex           uint64
 }
 
 // NodeIndexSort reverse sorts nodes by CreateIndex

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, false, 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, true, nil)
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
@@ -204,6 +207,78 @@ 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")
+	}
+	if out.SchedulingEligibility != structs.NodeSchedulingEligible {
+		t.Fatalf("should be eligible")
+	}
+}
+
+func TestNodes_ToggleEligibility(t *testing.T) {
+	t.Parallel()
+	c, s := makeClient(t, nil, func(c *testutil.TestServerConfig) {
+		c.DevMode = true
+	})
+	defer s.Stop()
+	nodes := c.Nodes()
+
+	// Wait for node registration and get the ID
+	var nodeID string
+	testutil.WaitForResult(func() (bool, error) {
+		out, _, err := nodes.List(nil)
+		if err != nil {
+			return false, err
+		}
+		if n := len(out); n != 1 {
+			return false, fmt.Errorf("expected 1 node, got: %d", n)
+		}
+		nodeID = out[0].ID
+		return true, nil
+	}, func(err error) {
+		t.Fatalf("err: %s", err)
+	})
+
+	// Check for eligibility
+	out, _, err := nodes.Info(nodeID, nil)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if out.SchedulingEligibility != structs.NodeSchedulingEligible {
+		t.Fatalf("node should be eligible")
+	}
+
+	// Toggle it off
+	wm, err := nodes.ToggleEligibility(nodeID, false, nil)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	assertWriteMeta(t, wm)
+
+	// Check again
+	out, _, err = nodes.Info(nodeID, nil)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if out.SchedulingEligibility != structs.NodeSchedulingIneligible {
+		t.Fatalf("bad eligibility: %v vs %v", out.SchedulingEligibility, structs.NodeSchedulingIneligible)
+	}
+
+	// Toggle on
+	wm, err = nodes.ToggleEligibility(nodeID, true, nil)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	assertWriteMeta(t, wm)
+
+	// Check again
+	out, _, err = nodes.Info(nodeID, nil)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if out.SchedulingEligibility != structs.NodeSchedulingEligible {
+		t.Fatalf("bad eligibility: %v vs %v", out.SchedulingEligibility, structs.NodeSchedulingEligible)
+	}
 }
 
 func TestNodes_Allocations(t *testing.T) {