From fb40e8babe259632a3dbb0956266141d8b9ec8d4 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Wed, 7 Mar 2018 15:16:45 -0800
Subject: [PATCH] handle empty node case

---
 nomad/drainer_int_test.go      | 43 ++++++++++++++++++++++++++++++++++
 nomad/drainerv2/watch_nodes.go | 17 ++++++++++++--
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/nomad/drainer_int_test.go b/nomad/drainer_int_test.go
index 0c0372d16..39422a5a0 100644
--- a/nomad/drainer_int_test.go
+++ b/nomad/drainer_int_test.go
@@ -186,3 +186,46 @@ func TestDrainer_Simple_ServiceOnly(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	})
 }
+
+func TestDrainer_DrainEmptyNode(t *testing.T) {
+	t.Parallel()
+	require := require.New(t)
+	s1 := TestServer(t, nil)
+	defer s1.Shutdown()
+	codec := rpcClient(t, s1)
+	testutil.WaitForLeader(t, s1.RPC)
+
+	// Create a node
+	n1 := mock.Node()
+	nodeReg := &structs.NodeRegisterRequest{
+		Node:         n1,
+		WriteRequest: structs.WriteRequest{Region: "global"},
+	}
+	var nodeResp structs.NodeUpdateResponse
+	require.Nil(msgpackrpc.CallWithCodec(codec, "Node.Register", nodeReg, &nodeResp))
+
+	// Drain the node
+	drainReq := &structs.NodeUpdateDrainRequest{
+		NodeID: n1.ID,
+		DrainStrategy: &structs.DrainStrategy{
+			DrainSpec: structs.DrainSpec{
+				Deadline: 10 * time.Minute,
+			},
+		},
+		WriteRequest: structs.WriteRequest{Region: "global"},
+	}
+	var drainResp structs.NodeDrainUpdateResponse
+	require.Nil(msgpackrpc.CallWithCodec(codec, "Node.UpdateDrain", drainReq, &drainResp))
+
+	// Check that the node drain is removed
+	state := s1.State()
+	testutil.WaitForResult(func() (bool, error) {
+		node, err := state.NodeByID(nil, n1.ID)
+		if err != nil {
+			return false, err
+		}
+		return node.DrainStrategy == nil, fmt.Errorf("has drain strategy still set")
+	}, func(err error) {
+		t.Fatalf("err: %v", err)
+	})
+}
diff --git a/nomad/drainerv2/watch_nodes.go b/nomad/drainerv2/watch_nodes.go
index 9b6b32b3a..289767d34 100644
--- a/nomad/drainerv2/watch_nodes.go
+++ b/nomad/drainerv2/watch_nodes.go
@@ -78,9 +78,22 @@ func (n *NodeDrainer) Update(node *structs.Node) {
 		n.jobWatcher.RegisterJob(job)
 	}
 
-	// TODO we need to check if the node is done such that if an operator drains
-	// a node with nothing on it we unset drain
+	// Check if the node is done such that if an operator drains a node with
+	// nothing on it we unset drain
+	done, err := draining.IsDone()
+	if err != nil {
+		n.logger.Printf("[ERR] nomad.drain: failed to check if node %q is done draining: %v", node.ID, err)
+		return
+	}
 
+	if done {
+		index, err := n.raft.NodeDrainComplete(node.ID)
+		if err != nil {
+			n.logger.Printf("[ERR] nomad.drain: failed to unset drain for node %q: %v", node.ID, err)
+		} else {
+			n.logger.Printf("[INFO] nomad.drain: node %q completed draining at index %d", node.ID, index)
+		}
+	}
 }
 
 // nodeDrainWatcher is used to watch nodes that are entering, leaving or