Merge pull request #20 from libp2p/feat/shrink-addrs

add a method for atomically updating TTLs

addr_manager.go

@@ -2,6 +2,7 @@ package peerstore
 
 import (
 	"context"
+	"math"
 	"sort"
 	"sync"
 	"time"
@@ -27,25 +28,27 @@ const (
 
 	// OwnObservedAddrTTL is used for our own external addresses observed by peers.
 	OwnObservedAddrTTL = time.Minute * 10
+)
 
-	// PermanentAddrTTL is the ttl for a "permanent address" (e.g. bootstrap nodes)
+// Permanent TTLs (distinct so we can distinguish between them)
-	// if we haven't shipped you an update to ipfs in 356 days
+const (
-	// we probably arent running the same bootstrap nodes...
+	// PermanentAddrTTL is the ttl for a "permanent address" (e.g. bootstrap nodes).
-	PermanentAddrTTL = time.Hour * 24 * 356
+	PermanentAddrTTL = math.MaxInt64 - iota
 
 	// ConnectedAddrTTL is the ttl used for the addresses of a peer to whom
 	// we're connected directly. This is basically permanent, as we will
 	// clear them + re-add under a TempAddrTTL after disconnecting.
-	ConnectedAddrTTL = PermanentAddrTTL
+	ConnectedAddrTTL
 )
 
 type expiringAddr struct {
-	Addr ma.Multiaddr
+	Addr    ma.Multiaddr
-	TTL  time.Time
+	TTL     time.Duration
+	Expires time.Time
 }
 
 func (e *expiringAddr) ExpiredBy(t time.Time) bool {
-	return t.After(e.TTL)
+	return t.After(e.Expires)
 }
 
 type addrSet map[string]expiringAddr
@@ -122,8 +125,8 @@ func (mgr *AddrManager) AddAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Durat
 
 		addrstr := string(addr.Bytes())
 		a, found := amap[addrstr]
-		if !found || exp.After(a.TTL) {
+		if !found || exp.After(a.Expires) {
-			amap[addrstr] = expiringAddr{Addr: addr, TTL: exp}
+			amap[addrstr] = expiringAddr{Addr: addr, Expires: exp, TTL: ttl}
 
 			for _, sub := range subs {
 				sub.pubAddr(addr)
@@ -164,7 +167,7 @@ func (mgr *AddrManager) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Durat
 		addrs := string(addr.Bytes())
 
 		if ttl > 0 {
-			amap[addrs] = expiringAddr{Addr: addr, TTL: exp}
+			amap[addrs] = expiringAddr{Addr: addr, Expires: exp, TTL: ttl}
 
 			for _, sub := range subs {
 				sub.pubAddr(addr)
@@ -175,6 +178,31 @@ func (mgr *AddrManager) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Durat
 	}
 }
 
+// UpdateAddrs updates the addresses associated with the given peer that have
+// the given oldTTL to have the given newTTL.
+func (mgr *AddrManager) UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL time.Duration) {
+	mgr.addrmu.Lock()
+	defer mgr.addrmu.Unlock()
+
+	if mgr.addrs == nil {
+		return
+	}
+
+	amap, found := mgr.addrs[p]
+	if !found {
+		return
+	}
+
+	exp := time.Now().Add(newTTL)
+	for addrstr, aexp := range amap {
+		if oldTTL == aexp.TTL {
+			aexp.TTL = newTTL
+			aexp.Expires = exp
+			amap[addrstr] = aexp
+		}
+	}
+}
+
 // Addresses returns all known (and valid) addresses for a given
 func (mgr *AddrManager) Addrs(p peer.ID) []ma.Multiaddr {
 	mgr.addrmu.Lock()

addr_manager_test.go

@@ -9,6 +9,7 @@ import (
 )
 
 func IDS(t *testing.T, ids string) peer.ID {
+	t.Helper()
 	id, err := peer.IDB58Decode(ids)
 	if err != nil {
 		t.Fatalf("id %q is bad: %s", ids, err)
@@ -17,6 +18,7 @@ func IDS(t *testing.T, ids string) peer.ID {
 }
 
 func MA(t *testing.T, m string) ma.Multiaddr {
+	t.Helper()
 	maddr, err := ma.NewMultiaddr(m)
 	if err != nil {
 		t.Fatal(err)
@@ -25,6 +27,7 @@ func MA(t *testing.T, m string) ma.Multiaddr {
 }
 
 func testHas(t *testing.T, exp, act []ma.Multiaddr) {
+	t.Helper()
 	if len(exp) != len(act) {
 		t.Fatal("lengths not the same")
 	}
@@ -202,6 +205,52 @@ func TestSetNegativeTTLClears(t *testing.T) {
 	testHas(t, nil, m.Addrs(id1))
 }
 
+func TestUpdateTTLs(t *testing.T) {
+	id1 := IDS(t, "QmcNstKuwBBoVTpSCSDrwzjgrRcaYXK833Psuz2EMHwyQN")
+	id2 := IDS(t, "QmcNstKuwBBoVTpSCSDrwzjgrRcaYXK833Psuz2EMHwyQM")
+	ma11 := MA(t, "/ip4/1.2.3.1/tcp/1111")
+	ma12 := MA(t, "/ip4/1.2.3.1/tcp/1112")
+	ma21 := MA(t, "/ip4/1.2.3.1/tcp/1121")
+	ma22 := MA(t, "/ip4/1.2.3.1/tcp/1122")
+
+	m := AddrManager{}
+
+	// Shouldn't panic.
+	m.UpdateAddrs(id1, time.Hour, time.Minute)
+
+	m.SetAddr(id1, ma11, time.Hour)
+	m.SetAddr(id1, ma12, time.Minute)
+
+	// Shouldn't panic.
+	m.UpdateAddrs(id2, time.Hour, time.Minute)
+
+	m.SetAddr(id2, ma21, time.Hour)
+	m.SetAddr(id2, ma22, time.Minute)
+
+	testHas(t, []ma.Multiaddr{ma11, ma12}, m.Addrs(id1))
+	testHas(t, []ma.Multiaddr{ma21, ma22}, m.Addrs(id2))
+
+	m.UpdateAddrs(id1, time.Hour, time.Millisecond)
+
+	testHas(t, []ma.Multiaddr{ma11, ma12}, m.Addrs(id1))
+	testHas(t, []ma.Multiaddr{ma21, ma22}, m.Addrs(id2))
+
+	time.Sleep(time.Millisecond)
+
+	testHas(t, []ma.Multiaddr{ma12}, m.Addrs(id1))
+	testHas(t, []ma.Multiaddr{ma21, ma22}, m.Addrs(id2))
+
+	m.UpdateAddrs(id2, time.Hour, time.Millisecond)
+
+	testHas(t, []ma.Multiaddr{ma12}, m.Addrs(id1))
+	testHas(t, []ma.Multiaddr{ma21, ma22}, m.Addrs(id2))
+
+	time.Sleep(time.Millisecond)
+
+	testHas(t, []ma.Multiaddr{ma12}, m.Addrs(id1))
+	testHas(t, []ma.Multiaddr{ma22}, m.Addrs(id2))
+}
+
 func TestNilAddrsDontBreak(t *testing.T) {
 	id1 := IDS(t, "QmcNstKuwBBoVTpSCSDrwzjgrRcaYXK833Psuz2EMHwyQN")
 	m := AddrManager{}

peerstore.go

@@ -69,6 +69,10 @@ type AddrBook interface {
 	// This is used when we receive the best estimate of the validity of an address.
 	SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration)
 
+	// UpdateAddrs updates the addresses associated with the given peer that have
+	// the given oldTTL to have the given newTTL.
+	UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL time.Duration)
+
 	// Addresses returns all known (and valid) addresses for a given peer
 	Addrs(p peer.ID) []ma.Multiaddr