From 76b913c5198ccc69303476dbef7b86d75b41458c Mon Sep 17 00:00:00 2001
From: Diptanu Choudhury <diptanuc@gmail.com>
Date: Tue, 12 Jul 2016 11:29:23 -0600
Subject: [PATCH] Handling allocations with client state pending

---
 nomad/node_endpoint.go          |  9 +++++++++
 nomad/state/state_store.go      |  3 ++-
 nomad/state/state_store_test.go | 15 +++++++++++++--
 nomad/structs/structs.go        |  4 ++++
 4 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/nomad/node_endpoint.go b/nomad/node_endpoint.go
index 25c8e118ef8c..51e041a7550f 100644
--- a/nomad/node_endpoint.go
+++ b/nomad/node_endpoint.go
@@ -66,6 +66,9 @@ func (n *Node) Register(args *structs.NodeRegisterRequest, reply *structs.NodeUp
 		return fmt.Errorf("invalid status for node")
 	}
 
+	// Set the timestamp when the node is registered
+	args.Node.StatusUpdatedAt = time.Now().Unix()
+
 	// Compute the node class
 	if err := args.Node.ComputeClass(); err != nil {
 		return fmt.Errorf("failed to computed node class: %v", err)
@@ -218,6 +221,9 @@ func (n *Node) UpdateStatus(args *structs.NodeUpdateStatusRequest, reply *struct
 		return fmt.Errorf("node not found")
 	}
 
+	// Update the timestamp of when the node status was updated
+	node.StatusUpdatedAt = time.Now().Unix()
+
 	// Commit this update via Raft
 	var index uint64
 	if node.Status != args.Status {
@@ -291,6 +297,9 @@ func (n *Node) UpdateDrain(args *structs.NodeUpdateDrainRequest,
 		return fmt.Errorf("node not found")
 	}
 
+	// Update the timestamp to
+	node.StatusUpdatedAt = time.Now().Unix()
+
 	// Commit this update via Raft
 	var index uint64
 	if node.Drain != args.Drain {
diff --git a/nomad/state/state_store.go b/nomad/state/state_store.go
index b3c7a6cdadbc..8f40e64a4fcc 100644
--- a/nomad/state/state_store.go
+++ b/nomad/state/state_store.go
@@ -201,7 +201,8 @@ func (s *StateStore) UpdateNodeStatus(index uint64, nodeID, status string) error
 		for _, alloc := range allocs {
 			copyAlloc := new(structs.Allocation)
 			*copyAlloc = *alloc
-			if alloc.ClientStatus == structs.AllocClientStatusRunning {
+			if alloc.ClientStatus == structs.AllocClientStatusPending ||
+				alloc.ClientStatus == structs.AllocClientStatusRunning {
 				copyAlloc.ClientStatus = structs.AllocClientStatusLost
 				if err := txn.Insert("allocs", copyAlloc); err != nil {
 					return fmt.Errorf("alloc insert failed: %v", err)
diff --git a/nomad/state/state_store_test.go b/nomad/state/state_store_test.go
index 800e8d5605ff..56b8cc304679 100644
--- a/nomad/state/state_store_test.go
+++ b/nomad/state/state_store_test.go
@@ -138,12 +138,15 @@ func TestStateStore_UpdateNodeStatus_Node(t *testing.T) {
 
 	alloc := mock.Alloc()
 	alloc1 := mock.Alloc()
+	alloc2 := mock.Alloc()
 	alloc.NodeID = node.ID
 	alloc1.NodeID = node.ID
+	alloc2.NodeID = node.ID
 	alloc.ClientStatus = structs.AllocClientStatusRunning
 	alloc1.ClientStatus = structs.AllocClientStatusFailed
+	alloc2.ClientStatus = structs.AllocClientStatusPending
 
-	if err = state.UpsertAllocs(1002, []*structs.Allocation{alloc, alloc1}); err != nil {
+	if err = state.UpsertAllocs(1002, []*structs.Allocation{alloc, alloc1, alloc2}); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	if err = state.UpdateNodeStatus(1003, node.ID, structs.NodeStatusDown); err != nil {
@@ -163,7 +166,15 @@ func TestStateStore_UpdateNodeStatus_Node(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 	if alloc1Out.ClientStatus != structs.AllocClientStatusFailed {
-		t.Fatalf("expected alloc status: %v, actual: %v", structs.AllocClientStatusLost, allocOut.ClientStatus)
+		t.Fatalf("expected alloc status: %v, actual: %v", structs.AllocClientStatusFailed, alloc1Out.ClientStatus)
+	}
+
+	alloc2Out, err := state.AllocByID(alloc2.ID)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if alloc2Out.ClientStatus != structs.AllocClientStatusLost {
+		t.Fatalf("expected alloc status: %v, actual: %v", structs.AllocClientStatusLost, alloc2Out.ClientStatus)
 	}
 
 	notify.verify(t)
diff --git a/nomad/structs/structs.go b/nomad/structs/structs.go
index 148c24b09d95..55d2af7073d6 100644
--- a/nomad/structs/structs.go
+++ b/nomad/structs/structs.go
@@ -634,6 +634,10 @@ type Node struct {
 	// StatusDescription is meant to provide more human useful information
 	StatusDescription string
 
+	// StatusUpdatedAt is the time stamp at which the state of the node was
+	// updated
+	StatusUpdatedAt int64
+
 	// Raft Indexes
 	CreateIndex uint64
 	ModifyIndex uint64