p2p/peer: addressbook can now clear addrs

p2p/peer/addr/addrsrcs.go

+// Package addr provides utility functions to handle peer addresses.
+package addr
+
+import (
+	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
+)
+
+// AddrSource is a source of addresses. It allows clients to retrieve
+// a set of addresses at a last possible moment in time. It is used
+// to query a set of addresses that may change over time, as a result
+// of the network changing interfaces or mappings.
+type Source interface {
+	Addrs() []ma.Multiaddr
+}
+
+// CombineSources returns a new AddrSource which is the
+// concatenation of all input AddrSources:
+//
+//   combined := CombinedSources(a, b)
+//   combined.Addrs() // append(a.Addrs(), b.Addrs()...)
+//
+func CombineSources(srcs ...Source) Source {
+	return combinedAS(srcs)
+}
+
+type combinedAS []Source
+
+func (cas combinedAS) Addrs() []ma.Multiaddr {
+	var addrs []ma.Multiaddr
+	for _, s := range cas {
+		addrs = append(addrs, s.Addrs()...)
+	}
+	return addrs
+}
+
+// UniqueSource returns a new AddrSource which omits duplicate
+// addresses from the inputs:
+//
+//   unique := UniqueSource(a, b)
+//   unique.Addrs() // append(a.Addrs(), b.Addrs()...)
+//                  // but only adds each addr once.
+//
+func UniqueSource(srcs ...Source) Source {
+	return uniqueAS(srcs)
+}
+
+type uniqueAS []Source
+
+func (uas uniqueAS) Addrs() []ma.Multiaddr {
+	seen := make(map[string]struct{})
+	var addrs []ma.Multiaddr
+	for _, s := range uas {
+		for _, a := range s.Addrs() {
+			s := a.String()
+			if _, found := seen[s]; !found {
+				addrs = append(addrs, a)
+				seen[s] = struct{}{}
+			}
+		}
+	}
+	return addrs
+}
+
+// Slice is a simple slice of addresses that implements
+// the AddrSource interface.
+type Slice []ma.Multiaddr
+
+func (as Slice) Addrs() []ma.Multiaddr {
+	return as
+}

p2p/peer/addr/addrsrcs_test.go

+package addr
+
+import (
+	"fmt"
+	"testing"
+
+	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
+)
+
+func newAddrOrFatal(t *testing.T, s string) ma.Multiaddr {
+	a, err := ma.NewMultiaddr(s)
+	if err != nil {
+		t.Fatal("error parsing multiaddr", err)
+	}
+	return a
+}
+
+func newAddrs(t *testing.T, n int) []ma.Multiaddr {
+	addrs := make([]ma.Multiaddr, n)
+	for i := 0; i < n; i++ {
+		s := fmt.Sprintf("/ip4/1.2.3.4/tcp/%d", i)
+		addrs[i] = newAddrOrFatal(t, s)
+	}
+	return addrs
+}
+
+func addrSetsSame(a, b []ma.Multiaddr) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i, aa := range a {
+		bb := b[i]
+		if !aa.Equal(bb) {
+			return false
+		}
+	}
+	return true
+}
+
+func addrSourcesSame(a, b Source) bool {
+	return addrSetsSame(a.Addrs(), b.Addrs())
+}
+
+func TestAddrCombine(t *testing.T) {
+	addrs := newAddrs(t, 30)
+	a := Slice(addrs[0:10])
+	b := Slice(addrs[10:20])
+	c := Slice(addrs[20:30])
+	d := CombineSources(a, b, c)
+	if !addrSetsSame(addrs, d.Addrs()) {
+		t.Error("addrs differ")
+	}
+	if !addrSourcesSame(Slice(addrs), d) {
+		t.Error("addrs differ")
+	}
+}
+
+func TestAddrUnique(t *testing.T) {
+
+	addrs := newAddrs(t, 40)
+	a := Slice(addrs[0:20])
+	b := Slice(addrs[10:30])
+	c := Slice(addrs[20:40])
+	d := CombineSources(a, b, c)
+	e := UniqueSource(a, b, c)
+	if addrSetsSame(addrs, d.Addrs()) {
+		t.Error("addrs same")
+	}
+	if addrSourcesSame(Slice(addrs), d) {
+		t.Error("addrs same")
+	}
+	if !addrSetsSame(addrs, e.Addrs()) {
+		t.Error("addrs differ", addrs, "\n\n", e.Addrs(), "\n\n")
+	}
+	if !addrSourcesSame(Slice(addrs), e) {
+		t.Error("addrs differ", addrs, "\n\n", e.Addrs(), "\n\n")
+	}
+}

p2p/peer/peerstore.go

@@ -38,9 +38,10 @@ type Peerstore interface {
 
 // AddressBook tracks the addresses of Peers
 type AddressBook interface {
-	Addresses(ID) []ma.Multiaddr
-	AddAddress(ID, ma.Multiaddr)
-	AddAddresses(ID, []ma.Multiaddr)
+	Addresses(ID) []ma.Multiaddr     // returns addresses for ID
+	AddAddress(ID, ma.Multiaddr)     // Adds given addr for ID
+	AddAddresses(ID, []ma.Multiaddr) // Adds given addrs for ID
+	SetAddresses(ID, []ma.Multiaddr) // Sets given addrs for ID (clears previously stored)
 }
 
 type addressMap map[string]ma.Multiaddr
@@ -81,27 +82,32 @@ func (ab *addressbook) Addresses(p ID) []ma.Multiaddr {
 }
 
 func (ab *addressbook) AddAddress(p ID, m ma.Multiaddr) {
+	ab.AddAddresses(p, []ma.Multiaddr{m})
+}
+
+func (ab *addressbook) AddAddresses(p ID, ms []ma.Multiaddr) {
 	ab.Lock()
 	defer ab.Unlock()
 
-	_, found := ab.addrs[p]
+	amap, found := ab.addrs[p]
 	if !found {
-		ab.addrs[p] = addressMap{}
+		amap = addressMap{}
+		ab.addrs[p] = amap
+	}
+	for _, m := range ms {
+		amap[m.String()] = m
 	}
-	ab.addrs[p][m.String()] = m
 }
 
-func (ab *addressbook) AddAddresses(p ID, ms []ma.Multiaddr) {
+func (ab *addressbook) SetAddresses(p ID, ms []ma.Multiaddr) {
 	ab.Lock()
 	defer ab.Unlock()
 
+	amap := addressMap{}
 	for _, m := range ms {
-		_, found := ab.addrs[p]
-		if !found {
-			ab.addrs[p] = addressMap{}
-		}
-		ab.addrs[p][m.String()] = m
+		amap[m.String()] = m
 	}
+	ab.addrs[p] = amap // clear what was there before
 }
 
 // KeyBook tracks the Public keys of Peers.

p2p/peer/peerstore_test.go

@@ -29,20 +29,37 @@ func TestAddresses(t *testing.T) {
 	id1 := IDS(t, "QmcNstKuwBBoVTpSCSDrwzjgrRcaYXK833Psuz2EMHwyQN")
 	id2 := IDS(t, "QmRmPL3FDZKE3Qiwv1RosLdwdvbvg17b2hB39QPScgWKKZ")
 	id3 := IDS(t, "QmPhi7vBsChP7sjRoZGgg7bcKqF6MmCcQwvRbDte8aJ6Kn")
+	id4 := IDS(t, "QmPhi7vBsChP7sjRoZGgg7bcKqF6MmCcQwvRbDte8aJ5Kn")
+	id5 := IDS(t, "QmPhi7vBsChP7sjRoZGgg7bcKqF6MmCcQwvRbDte8aJ5Km")
 
 	ma11 := MA(t, "/ip4/1.2.3.1/tcp/1111")
-	ma21 := MA(t, "/ip4/1.2.3.2/tcp/1111")
-	ma22 := MA(t, "/ip4/1.2.3.2/tcp/2222")
-	ma31 := MA(t, "/ip4/1.2.3.3/tcp/1111")
-	ma32 := MA(t, "/ip4/1.2.3.3/tcp/2222")
-	ma33 := MA(t, "/ip4/1.2.3.3/tcp/3333")
+	ma21 := MA(t, "/ip4/2.2.3.2/tcp/1111")
+	ma22 := MA(t, "/ip4/2.2.3.2/tcp/2222")
+	ma31 := MA(t, "/ip4/3.2.3.3/tcp/1111")
+	ma32 := MA(t, "/ip4/3.2.3.3/tcp/2222")
+	ma33 := MA(t, "/ip4/3.2.3.3/tcp/3333")
+	ma41 := MA(t, "/ip4/4.2.3.3/tcp/1111")
+	ma42 := MA(t, "/ip4/4.2.3.3/tcp/2222")
+	ma43 := MA(t, "/ip4/4.2.3.3/tcp/3333")
+	ma44 := MA(t, "/ip4/4.2.3.3/tcp/4444")
+	ma51 := MA(t, "/ip4/5.2.3.3/tcp/1111")
+	ma52 := MA(t, "/ip4/5.2.3.3/tcp/2222")
+	ma53 := MA(t, "/ip4/5.2.3.3/tcp/3333")
+	ma54 := MA(t, "/ip4/5.2.3.3/tcp/4444")
+	ma55 := MA(t, "/ip4/5.2.3.3/tcp/5555")
 
 	ps.AddAddress(id1, ma11)
-	ps.AddAddress(id2, ma21)
-	ps.AddAddress(id2, ma22)
+	ps.AddAddresses(id2, []ma.Multiaddr{ma21, ma22})
+	ps.AddAddresses(id2, []ma.Multiaddr{ma21, ma22}) // idempotency
 	ps.AddAddress(id3, ma31)
 	ps.AddAddress(id3, ma32)
 	ps.AddAddress(id3, ma33)
+	ps.AddAddress(id3, ma33) // idempotency
+	ps.AddAddress(id3, ma33)
+	ps.AddAddresses(id4, []ma.Multiaddr{ma41, ma42, ma43, ma44})       // multiple
+	ps.AddAddresses(id5, []ma.Multiaddr{ma21, ma22})                   // clearing
+	ps.AddAddresses(id5, []ma.Multiaddr{ma41, ma42, ma43, ma44})       // clearing
+	ps.SetAddresses(id5, []ma.Multiaddr{ma51, ma52, ma53, ma54, ma55}) // clearing
 
 	test := func(exp, act []ma.Multiaddr) {
 		if len(exp) != len(act) {
@@ -69,9 +86,13 @@ func TestAddresses(t *testing.T) {
 	test([]ma.Multiaddr{ma11}, ps.Addresses(id1))
 	test([]ma.Multiaddr{ma21, ma22}, ps.Addresses(id2))
 	test([]ma.Multiaddr{ma31, ma32, ma33}, ps.Addresses(id3))
+	test([]ma.Multiaddr{ma41, ma42, ma43, ma44}, ps.Addresses(id4))
+	test([]ma.Multiaddr{ma51, ma52, ma53, ma54, ma55}, ps.Addresses(id5))
 
 	// test also the PeerInfo return
 	test([]ma.Multiaddr{ma11}, ps.PeerInfo(id1).Addrs)
 	test([]ma.Multiaddr{ma21, ma22}, ps.PeerInfo(id2).Addrs)
 	test([]ma.Multiaddr{ma31, ma32, ma33}, ps.PeerInfo(id3).Addrs)
+	test([]ma.Multiaddr{ma41, ma42, ma43, ma44}, ps.PeerInfo(id4).Addrs)
+	test([]ma.Multiaddr{ma51, ma52, ma53, ma54, ma55}, ps.PeerInfo(id5).Addrs)
 }