node drain metadata wip

api/nodes.go

@@ -71,6 +71,9 @@ type NodeUpdateDrainRequest struct {
 	// MarkEligible marks the node as eligible for scheduling if removing
 	// the drain strategy.
 	MarkEligible bool
+
+	// Meta allows operators to specify metadata related to the drain operation
+	Meta map[string]string
 }
 
 // NodeDrainUpdateResponse is used to respond to a node drain update
@@ -81,14 +84,45 @@ type NodeDrainUpdateResponse struct {
 	WriteMeta
 }
 
+// DrainOptions is used to pass through node drain parameters
+type DrainOptions struct {
+	// DrainSpec contains the drain specification for the node. If non-nil,
+	// the node will be marked ineligible and begin/continue draining according
+	// to the provided drain spec.
+	// If nil, any existing drain operation will be canceled.
+	DrainSpec *DrainSpec
+
+	// MarkEligible indicates whether the node should be marked as eligible when
+	// canceling a drain operation.
+	MarkEligible bool
+
+	// Meta is metadata that is persisted in Node.LastDrain about this
+	// drain update.
+	Meta map[string]string
+}
+
 // 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 eligible
 func (n *Nodes) UpdateDrain(nodeID string, spec *DrainSpec, markEligible bool, q *WriteOptions) (*NodeDrainUpdateResponse, error) {
-	req := &NodeUpdateDrainRequest{
-		NodeID:       nodeID,
+	resp, err := n.UpdateDrainOpts(nodeID, DrainOptions{
 		DrainSpec:    spec,
 		MarkEligible: markEligible,
+		Meta:         nil,
+	}, q)
+	return resp, err
+}
+
+// UpdateDrainWithMeta 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 eligible
+func (n *Nodes) UpdateDrainOpts(nodeID string, opts DrainOptions, q *WriteOptions) (*NodeDrainUpdateResponse,
+	error) {
+	req := &NodeUpdateDrainRequest{
+		NodeID:       nodeID,
+		DrainSpec:    opts.DrainSpec,
+		MarkEligible: opts.MarkEligible,
+		Meta:         opts.Meta,
 	}
 
 	var resp NodeDrainUpdateResponse

api/nodes_test.go

@@ -248,7 +248,14 @@ func TestNodes_ToggleDrain(t *testing.T) {
 	spec := &DrainSpec{
 		Deadline: 10 * time.Second,
 	}
-	drainOut, err := nodes.UpdateDrain(nodeID, spec, false, nil)
+	drainMeta := map[string]string{
+		"reason": "this node needs to go",
+	}
+	drainOut, err := nodes.UpdateDrainOpts(nodeID, DrainOptions{
+		DrainSpec:    spec,
+		MarkEligible: false,
+		Meta:         drainMeta,
+	}, nil)
 	require.Nil(err)
 	assertWriteMeta(t, &drainOut.WriteMeta)
 
@@ -280,16 +287,14 @@ func TestNodes_ToggleDrain(t *testing.T) {
 					now := time.Now()
 					require.True(!node.LastDrain.StartedAt.Before(timeBeforeDrain) && !node.LastDrain.StartedAt.After(now),
 						"wanted %v <= %v <= %v", timeBeforeDrain, node.LastDrain.StartedAt, now)
-					// TODO: test meta vvvv
-					// require.Equal(node.LastDrain.Meta["reason"])
+					require.Equal(drainMeta, node.LastDrain.Meta)
 					sawDraining = node.ModifyIndex
 				} else if sawDraining != 0 && node.ModifyIndex > sawDraining &&
 					!node.Drain && node.SchedulingEligibility == NodeSchedulingIneligible {
 					require.NotNil(node.LastDrain)
 					require.Equal(DrainStatusCompleted, node.LastDrain.Status)
 					require.True(!node.LastDrain.UpdatedAt.Before(node.LastDrain.StartedAt))
-					// TODO: test meta vvvv
-					// require.Equal(node.LastDrain.Meta["reason"])
+					require.Equal(drainMeta, node.LastDrain.Meta)
 					sawDrainComplete = node.ModifyIndex
 				}
 			}

command/agent/node_endpoint.go

@@ -125,6 +125,7 @@ func (s *HTTPServer) nodeToggleDrain(resp http.ResponseWriter, req *http.Request
 	args := structs.NodeUpdateDrainRequest{
 		NodeID:       nodeID,
 		MarkEligible: drainRequest.MarkEligible,
+		Meta:         drainRequest.Meta,
 	}
 	if drainRequest.DrainSpec != nil {
 		args.DrainStrategy = &structs.DrainStrategy{

nomad/state/state_store.go

@@ -1015,27 +1015,38 @@ func (s *StateStore) updateNodeDrainImpl(txn *txn, index uint64, nodeID string,
 
 	// Update LastDrain
 	updateTime := time.Unix(updatedAt, 0)
+	// when done with this method, copyNode.LastDrain should be set
+	// this is either a new LastDrain struct or an update of the existing one
+	//
 	// if starting a new drain operation, create a new LastDrain. otherwise, update the existing one.
-	// if LastDrain doesn't exist, we'll need to create a new one. this might happen if we upgrade the
-	// server to 1.0.4 during a drain operation
-	if existingNode.DrainStrategy == nil && drain != nil || existingNode.LastDrain == nil {
-		// starting a new drain operation
+	// if already draining and LastDrain doesn't exist, we'll need to create a new one.
+	// this might happen if we upgrade/transition to 1.1 during a drain operation
+	if existingNode.DrainStrategy == nil && copyNode.DrainStrategy != nil ||
+		copyNode.LastDrain == nil {
+
 		copyNode.LastDrain = &structs.DrainMetadata{
 			StartedAt:  updateTime,
 			UpdatedAt:  updateTime,
-			Status:     structs.DrainStatusDraining,
 			AccessorID: accessorId,
 			Meta:       drainMeta,
 		}
+		switch {
+		case drain != nil:
+			copyNode.LastDrain.Status = structs.DrainStatusDraining
+		case drainCompleted:
+			copyNode.LastDrain.Status = structs.DrainStatusCompleted
+		default:
+			copyNode.LastDrain.Status = structs.DrainStatusCancelled
+		}
 	} else {
 		copyNode.LastDrain.UpdatedAt = updateTime
 		if accessorId != "" {
 			// we won't have an accessor ID for drain complete; don't overwrite the existing one
 			copyNode.LastDrain.AccessorID = accessorId
 		}
 		if drainMeta != nil {
-			// similarly, won't have metadata for drain complete; keep existing
-			copyNode.Meta = drainMeta
+			// similarly, won't have metadata for drain complete; keep the existing operator-provided metadata
+			copyNode.LastDrain.Meta = drainMeta
 		}
 		if drain == nil {
 			if drainCompleted {