diff --git a/pstoremem/addr_book.go b/pstoremem/addr_book.go index 311c436..f0f8a2c 100644 --- a/pstoremem/addr_book.go +++ b/pstoremem/addr_book.go @@ -26,31 +26,47 @@ func (e *expiringAddr) ExpiredBy(t time.Time) bool { return t.After(e.Expires) } -type addrSlice []expiringAddr - var _ pstore.AddrBook = (*memoryAddrBook)(nil) // memoryAddrBook manages addresses. type memoryAddrBook struct { addrmu sync.Mutex - addrs map[peer.ID]addrSlice + addrs map[peer.ID]map[string]expiringAddr + + nextGC time.Time subManager *AddrSubManager } func NewAddrBook() pstore.AddrBook { return &memoryAddrBook{ - addrs: make(map[peer.ID]addrSlice), + addrs: make(map[peer.ID]map[string]expiringAddr), subManager: NewAddrSubManager(), } } +// gc garbage collects the in-memory address book. The caller *must* hold the addrmu lock. +func (mab *memoryAddrBook) gc() { + now := time.Now() + if !now.After(mab.nextGC) { + return + } + for p, amap := range mab.addrs { + for k, addr := range amap { + if addr.ExpiredBy(now) { + delete(amap, k) + } + } + if len(amap) == 0 { + delete(mab.addrs, p) + } + } + mab.nextGC = time.Now().Add(pstore.AddressTTL) +} + func (mab *memoryAddrBook) PeersWithAddrs() peer.IDSlice { mab.addrmu.Lock() defer mab.addrmu.Unlock() - if mab.addrs == nil { - return nil - } pids := make(peer.IDSlice, 0, len(mab.addrs)) for pid := range mab.addrs { @@ -76,20 +92,17 @@ func (mab *memoryAddrBook) AddAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Du return } - oldAddrs := mab.addrs[p] - amap := make(map[string]expiringAddr, len(oldAddrs)) - for _, ea := range oldAddrs { - amap[string(ea.Addr.Bytes())] = ea + amap := mab.addrs[p] + if amap == nil { + amap = make(map[string]expiringAddr, len(addrs)) + mab.addrs[p] = amap } - - // only expand ttls exp := time.Now().Add(ttl) for _, addr := range addrs { if addr == nil { log.Warningf("was passed nil multiaddr for %s", p) continue } - addrstr := string(addr.Bytes()) a, found := amap[addrstr] if !found || exp.After(a.Expires) { @@ -98,11 +111,7 @@ func (mab *memoryAddrBook) AddAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Du mab.subManager.BroadcastAddr(p, addr) } } - newAddrs := make([]expiringAddr, 0, len(amap)) - for _, ea := range amap { - newAddrs = append(newAddrs, ea) - } - mab.addrs[p] = newAddrs + mab.gc() } // SetAddr calls mgr.SetAddrs(p, addr, ttl) @@ -116,10 +125,10 @@ func (mab *memoryAddrBook) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Du mab.addrmu.Lock() defer mab.addrmu.Unlock() - oldAddrs := mab.addrs[p] - amap := make(map[string]expiringAddr, len(oldAddrs)) - for _, ea := range oldAddrs { - amap[string(ea.Addr.Bytes())] = ea + amap := mab.addrs[p] + if amap == nil { + amap = make(map[string]expiringAddr, len(addrs)) + mab.addrs[p] = amap } exp := time.Now().Add(ttl) @@ -129,21 +138,17 @@ func (mab *memoryAddrBook) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Du continue } // re-set all of them for new ttl. - addrs := string(addr.Bytes()) + addrstr := string(addr.Bytes()) if ttl > 0 { - amap[addrs] = expiringAddr{Addr: addr, Expires: exp, TTL: ttl} + amap[addrstr] = expiringAddr{Addr: addr, Expires: exp, TTL: ttl} mab.subManager.BroadcastAddr(p, addr) } else { - delete(amap, addrs) + delete(amap, addrstr) } } - newAddrs := make([]expiringAddr, 0, len(amap)) - for _, ea := range amap { - newAddrs = append(newAddrs, ea) - } - mab.addrs[p] = newAddrs + mab.gc() } // UpdateAddrs updates the addresses associated with the given peer that have @@ -152,22 +157,20 @@ func (mab *memoryAddrBook) UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL t mab.addrmu.Lock() defer mab.addrmu.Unlock() - if mab.addrs == nil { - return - } - - addrs, found := mab.addrs[p] + amap, found := mab.addrs[p] if !found { return } exp := time.Now().Add(newTTL) - for i := range addrs { - if aexp := &addrs[i]; oldTTL == aexp.TTL { - aexp.TTL = newTTL - aexp.Expires = exp + for k, addr := range amap { + if oldTTL == addr.TTL { + addr.TTL = newTTL + addr.Expires = exp + amap[k] = addr } } + mab.gc() } // Addresses returns all known (and valid) addresses for a given @@ -175,32 +178,21 @@ func (mab *memoryAddrBook) Addrs(p peer.ID) []ma.Multiaddr { mab.addrmu.Lock() defer mab.addrmu.Unlock() - // not initialized? nothing to give. - if mab.addrs == nil { - return nil - } - - maddrs, found := mab.addrs[p] + amap, found := mab.addrs[p] if !found { return nil } now := time.Now() - good := make([]ma.Multiaddr, 0, len(maddrs)) - cleaned := make([]expiringAddr, 0, len(maddrs)) - for _, m := range maddrs { + good := make([]ma.Multiaddr, 0, len(amap)) + for k, m := range amap { if !m.ExpiredBy(now) { - cleaned = append(cleaned, m) good = append(good, m.Addr) + } else { + delete(amap, k) } } - // clean up the expired ones. - if len(cleaned) == 0 { - delete(mab.addrs, p) - } else { - mab.addrs[p] = cleaned - } return good }