From 14e7306524cc1700930b4405c41b6865d7b9d85e Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Sat, 29 Aug 2015 14:33:30 -0700
Subject: [PATCH] client: test watching allocations

---
 client/client.go      |  7 +++--
 client/client_test.go | 68 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/client/client.go b/client/client.go
index cbb242335..3ac5a7f5c 100644
--- a/client/client.go
+++ b/client/client.go
@@ -473,7 +473,7 @@ func (c *Client) watchAllocations(allocUpdates chan []*structs.Allocation) {
 func (c *Client) runAllocs(updated []*structs.Allocation) {
 	// Get the existing allocs
 	c.allocLock.RLock()
-	exist := make([]*structs.Allocation, len(c.allocs))
+	exist := make([]*structs.Allocation, 0, len(c.allocs))
 	for _, ctx := range c.allocs {
 		exist = append(exist, ctx.Alloc())
 	}
@@ -515,14 +515,15 @@ func (c *Client) runAllocs(updated []*structs.Allocation) {
 
 // removeAlloc is invoked when we should remove an allocation
 func (c *Client) removeAlloc(alloc *structs.Allocation) error {
-	c.allocLock.RLock()
-	defer c.allocLock.RUnlock()
+	c.allocLock.Lock()
+	defer c.allocLock.Unlock()
 	ctx, ok := c.allocs[alloc.ID]
 	if !ok {
 		c.logger.Printf("[WARN] client: missing context for alloc '%s'", alloc.ID)
 		return nil
 	}
 	ctx.Destroy()
+	delete(c.allocs, alloc.ID)
 	return nil
 }
 
diff --git a/client/client_test.go b/client/client_test.go
index e05184106..4f05a5e48 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -238,3 +238,71 @@ func TestClient_UpdateAllocStatus(t *testing.T) {
 		t.Fatalf("bad: %#v", out)
 	}
 }
+
+func TestClient_WatchAllocs(t *testing.T) {
+	s1, _ := testServer(t, nil)
+	defer s1.Shutdown()
+	testutil.WaitForLeader(t, s1.RPC)
+
+	c1 := testClient(t, func(c *config.Config) {
+		c.RPCHandler = s1
+	})
+	defer c1.Shutdown()
+
+	// Create mock allocations
+	alloc1 := mock.Alloc()
+	alloc1.NodeID = c1.Node().ID
+	alloc2 := mock.Alloc()
+	alloc2.NodeID = c1.Node().ID
+
+	state := s1.State()
+	err := state.UpdateAllocations(100,
+		[]*structs.Allocation{alloc1, alloc2})
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Both allocations should get registered
+	testutil.WaitForResult(func() (bool, error) {
+		c1.allocLock.RLock()
+		num := len(c1.allocs)
+		c1.allocLock.RUnlock()
+		return num == 2, nil
+	}, func(err error) {
+		t.Fatalf("err: %v", err)
+	})
+
+	// Delete one allocation
+	err = state.DeleteEval(101, nil, []string{alloc1.ID})
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Update the other allocation
+	alloc2.DesiredStatus = structs.AllocDesiredStatusStop
+	err = state.UpdateAllocations(102,
+		[]*structs.Allocation{alloc2})
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// One allocations should get de-registered
+	testutil.WaitForResult(func() (bool, error) {
+		c1.allocLock.RLock()
+		num := len(c1.allocs)
+		c1.allocLock.RUnlock()
+		return num == 1, nil
+	}, func(err error) {
+		t.Fatalf("err: %v", err)
+	})
+
+	// One allocations should get updated
+	testutil.WaitForResult(func() (bool, error) {
+		c1.allocLock.RLock()
+		ar := c1.allocs[alloc2.ID]
+		c1.allocLock.RUnlock()
+		return ar.Alloc().DesiredStatus == structs.AllocDesiredStatusStop, nil
+	}, func(err error) {
+		t.Fatalf("err: %v", err)
+	})
+}