From af9668f130f6b115a2d7b9a38c1384a90c19cdf8 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Thu, 25 Jun 2015 12:47:38 +0100
Subject: [PATCH] Keep trying peers to ask for space until we find one that
 doesn't error

---
 ipam/allocate.go       |  7 +++++--
 ipam/ring/ring.go      | 31 +++++++++++++++----------------
 ipam/ring/ring_test.go | 37 ++++++++++++++++++++-----------------
 3 files changed, 40 insertions(+), 35 deletions(-)

diff --git a/ipam/allocate.go b/ipam/allocate.go
index 130380e9..c1dcda66 100644
--- a/ipam/allocate.go
+++ b/ipam/allocate.go
@@ -42,9 +42,12 @@ func (g *allocate) Try(alloc *Allocator) bool {
 	}
 
 	// out of space
-	if donor, err := alloc.ring.ChoosePeerToAskForSpace(g.r.Start, g.r.End); err == nil {
+	donors := alloc.ring.ChoosePeersToAskForSpace(g.r.Start, g.r.End)
+	for _, donor := range donors {
 		alloc.debugln("Decided to ask peer", donor, "for space in range", g.r)
-		alloc.sendSpaceRequest(donor, g.r)
+		if err := alloc.sendSpaceRequest(donor, g.r); err == nil {
+			break
+		}
 	}
 
 	return false
diff --git a/ipam/ring/ring.go b/ipam/ring/ring.go
index fcdb362e..4908e571 100644
--- a/ipam/ring/ring.go
+++ b/ipam/ring/ring.go
@@ -392,9 +392,9 @@ func (r *Ring) ReportFree(freespace map[address.Address]address.Offset) {
 	}
 }
 
-// ChoosePeerToAskForSpace chooses a weighted-random peer to ask
-// for space in the range [start, end).  Assumes start<end.
-func (r *Ring) ChoosePeerToAskForSpace(start, end address.Address) (result router.PeerName, err error) {
+// ChoosePeerToAskForSpace returns all peers we can ask for space in
+// the range [start, end), in weighted-random order.  Assumes start<end.
+func (r *Ring) ChoosePeersToAskForSpace(start, end address.Address) (result []router.PeerName) {
 	var (
 		sum               address.Offset
 		totalSpacePerPeer = make(map[router.PeerName]address.Offset) // Compute total free space per peer
@@ -423,21 +423,20 @@ func (r *Ring) ChoosePeerToAskForSpace(start, end address.Address) (result route
 		sum += entry.Free
 	}
 
-	if sum == 0 {
-		err = ErrNoFreeSpace
-		return
-	}
-
-	// Pick random peer, weighted by total free space
-	rn := rand.Int63n(int64(sum))
-	for peername, space := range totalSpacePerPeer {
-		rn -= int64(space)
-		if rn < 0 {
-			return peername, nil
+	for len(totalSpacePerPeer) > 0 {
+		// Pick random peer, weighted by total free space
+		rn := rand.Int63n(int64(sum)) // note using 64-bit because rand.Intn uses signed int
+		for peername, space := range totalSpacePerPeer {
+			rn -= int64(space)
+			if rn < 0 {
+				result = append(result, peername)
+				sum -= space
+				delete(totalSpacePerPeer, peername)
+				break
+			}
 		}
 	}
-
-	panic("Should never reach this point")
+	return
 }
 
 func (r *Ring) PickPeerForTransfer() router.PeerName {
diff --git a/ipam/ring/ring_test.go b/ipam/ring/ring_test.go
index ce63a38e..7c23ccdd 100644
--- a/ipam/ring/ring_test.go
+++ b/ipam/ring/ring_test.go
@@ -345,39 +345,42 @@ func TestGossip(t *testing.T) {
 func TestFindFree(t *testing.T) {
 	ring1 := New(start, end, peer1name)
 
-	_, err := ring1.ChoosePeerToAskForSpace(start, end)
-	wt.AssertTrue(t, err == ErrNoFreeSpace, "Expected ErrNoFreeSpace")
+	peers := ring1.ChoosePeersToAskForSpace(start, end)
+	wt.AssertTrue(t, len(peers) == 0, "Expected no peers")
 
 	ring1.Entries = []*entry{{Token: start, Peer: peer1name}}
-	_, err = ring1.ChoosePeerToAskForSpace(start, end)
-	wt.AssertTrue(t, err == ErrNoFreeSpace, "Expected ErrNoFreeSpace")
+	peers = ring1.ChoosePeersToAskForSpace(start, end)
+	wt.AssertTrue(t, len(peers) == 0, "Expected no peers")
 
 	// We shouldn't return outselves
 	ring1.ReportFree(map[address.Address]address.Offset{start: 10})
-	_, err = ring1.ChoosePeerToAskForSpace(start, end)
-	wt.AssertTrue(t, err == ErrNoFreeSpace, "Expected ErrNoFreeSpace")
+	peers = ring1.ChoosePeersToAskForSpace(start, end)
+	wt.AssertTrue(t, len(peers) == 0, "Expected no peers")
 
 	ring1.Entries = []*entry{{Token: start, Peer: peer1name, Free: 1},
 		{Token: middle, Peer: peer1name, Free: 1}}
-	_, err = ring1.ChoosePeerToAskForSpace(start, end)
-	wt.AssertTrue(t, err == ErrNoFreeSpace, "Expected ErrNoFreeSpace")
+	peers = ring1.ChoosePeersToAskForSpace(start, end)
+	wt.AssertTrue(t, len(peers) == 0, "Expected no peers")
 	ring1.assertInvariants()
 
 	// We should return others
 	ring1.Entries = []*entry{{Token: start, Peer: peer2name, Free: 1}}
-	peer, err := ring1.ChoosePeerToAskForSpace(start, end)
-	wt.AssertSuccess(t, err)
-	wt.AssertEquals(t, peer, peer2name)
+	peers = ring1.ChoosePeersToAskForSpace(start, end)
+	wt.AssertTrue(t, len(peers) == 1, "Expected one peer")
+	wt.AssertEquals(t, peers[0], peer2name)
 
 	ring1.Entries = []*entry{{Token: start, Peer: peer2name, Free: 1},
 		{Token: middle, Peer: peer3name, Free: 1}}
-	peer, err = ring1.ChoosePeerToAskForSpace(start, middle)
-	wt.AssertSuccess(t, err)
-	wt.AssertEquals(t, peer, peer2name)
+	peers = ring1.ChoosePeersToAskForSpace(start, middle)
+	wt.AssertTrue(t, len(peers) == 1, "Expected one peer")
+	wt.AssertEquals(t, peers[0], peer2name)
 
-	peer, err = ring1.ChoosePeerToAskForSpace(middle, end)
-	wt.AssertSuccess(t, err)
-	wt.AssertEquals(t, peer, peer3name)
+	peers = ring1.ChoosePeersToAskForSpace(middle, end)
+	wt.AssertTrue(t, len(peers) == 1, "Expected one peer")
+	wt.AssertEquals(t, peers[0], peer3name)
+
+	peers = ring1.ChoosePeersToAskForSpace(start, end)
+	wt.AssertTrue(t, len(peers) == 2, "Expected two peers")
 	ring1.assertInvariants()
 }