Fix PeerSet.Remove() bug.

p2p/addrbook.go

@@ -334,7 +334,7 @@ func (a *AddrBook) loadFromFile(filePath string) {
 	// If doesn't exist, do nothing.
 	_, err := os.Stat(filePath)
 	if os.IsNotExist(err) {
-		return
+		panic(Fmt("File does not exist: %v", filePath))
 	}
 
 	// Load addrBookJSON{}

p2p/peer_set.go

@@ -70,12 +70,13 @@ func (ps *PeerSet) Remove(peer *Peer) {
 	// If it's the last peer, that's an easy special case.
 	if index == len(ps.list)-1 {
 		ps.list = newList
+		delete(ps.lookup, peer.Key)
 		return
 	}
 	// Move the last item from ps.list to "index" in list.
 	lastPeer := ps.list[len(ps.list)-1]
-	lastPeerAddr := lastPeer.mconn.RemoteAddress.String()
-	lastPeerItem := ps.lookup[lastPeerAddr]
+	lastPeerKey := lastPeer.Key
+	lastPeerItem := ps.lookup[lastPeerKey]
 	newList[index] = lastPeer
 	lastPeerItem.index = index
 	ps.list = newList

p2p/peer_set_test.go

@@ -0,0 +1,63 @@
+package p2p
+
+import (
+	"math/rand"
+	"testing"
+
+	. "github.com/tendermint/tendermint/common"
+)
+
+// Returns an empty dummy peer
+func randPeer() *Peer {
+	return &Peer{
+		Key: Fmt("%v.%v.%v.%v:%v", rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%10000+80),
+	}
+}
+
+func TestAddRemoveOne(t *testing.T) {
+	peerSet := NewPeerSet()
+
+	peer := randPeer()
+	added := peerSet.Add(peer)
+	if !added {
+		t.Errorf("Failed to add new peer")
+	}
+	if peerSet.Size() != 1 {
+		t.Errorf("Failed to add new peer and increment size")
+	}
+
+	peerSet.Remove(peer)
+	if peerSet.Has(peer.Key) {
+		t.Errorf("Failed to remove peer")
+	}
+	if peerSet.Size() != 0 {
+		t.Errorf("Failed to remove peer and decrement size")
+	}
+}
+
+func TestAddRemoveMany(t *testing.T) {
+	peerSet := NewPeerSet()
+
+	peers := []*Peer{}
+	for i := 0; i < 100; i++ {
+		peer := randPeer()
+		added := peerSet.Add(peer)
+		if !added {
+			t.Errorf("Failed to add new peer")
+		}
+		if peerSet.Size() != i+1 {
+			t.Errorf("Failed to add new peer and increment size")
+		}
+		peers = append(peers, peer)
+	}
+
+	for i, peer := range peers {
+		peerSet.Remove(peer)
+		if peerSet.Has(peer.Key) {
+			t.Errorf("Failed to remove peer")
+		}
+		if peerSet.Size() != len(peers)-i-1 {
+			t.Errorf("Failed to remove peer and decrement size")
+		}
+	}
+}