fix heap property, fix after check

p2p/host/peerstore/pstoremem/addr_book.go

@@ -114,11 +114,12 @@ func (pa *peerAddrs) NextExpiry() time.Time {
 	if len(pa.expiringHeap) == 0 {
 		return time.Time{}
 	}
-	return pa.expiringHeap[len(pa.expiringHeap)-1].Expires
+	return pa.expiringHeap[0].Expires
 }
 
 func (pa *peerAddrs) PopIfExpired(now time.Time) (*expiringAddr, bool) {
-	if len(pa.expiringHeap) > 0 && now.After(pa.NextExpiry()) {
+	// Use !Before and not After to ensure that we expire *at* now and not just after now.
+	if len(pa.expiringHeap) > 0 && !now.Before(pa.NextExpiry()) {
 		a := heap.Pop(pa)
 		return a.(*expiringAddr), true
 	}

p2p/host/peerstore/pstoremem/addr_book_test.go

@@ -0,0 +1,53 @@
+package pstoremem
+
+import (
+	"container/heap"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/libp2p/go-libp2p/core/peer"
+	ma "github.com/multiformats/go-multiaddr"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPeerAddrsNextExpiry(t *testing.T) {
+	paa := newPeerAddrs()
+	pa := &paa
+	a1 := ma.StringCast("/ip4/1.2.3.4/udp/1/quic-v1")
+	a2 := ma.StringCast("/ip4/1.2.3.4/udp/2/quic-v1")
+
+	// t1 is before t2
+	t1 := time.Time{}.Add(1 * time.Second)
+	t2 := time.Time{}.Add(2 * time.Second)
+	heap.Push(pa, &expiringAddr{Addr: a1, Expires: t1, TTL: 10 * time.Second, Peer: "p1"})
+	heap.Push(pa, &expiringAddr{Addr: a2, Expires: t2, TTL: 10 * time.Second, Peer: "p2"})
+
+	if pa.NextExpiry() != t1 {
+		t.Fatal("expiry should be set to t1, got", pa.NextExpiry())
+	}
+}
+
+func TestPeerAddrsHeapProperty(t *testing.T) {
+	paa := newPeerAddrs()
+	pa := &paa
+	addrs := []ma.Multiaddr{}
+	expiries := []time.Time{}
+	for i := 0; i < 1000; i++ {
+		addrs = append(addrs, ma.StringCast(fmt.Sprintf("/ip4/1.2.3.4/udp/%d/quic-v1", i)))
+		expiries = append(expiries, time.Time{}.Add(time.Duration(i)*time.Second))
+		pid := peer.ID(fmt.Sprintf("p%d", i))
+		heap.Push(pa, &expiringAddr{Addr: addrs[i], Expires: expiries[i], TTL: 10 * time.Second, Peer: pid})
+	}
+
+	for i := 0; i < 1000; i++ {
+		ea, ok := pa.PopIfExpired(expiries[i])
+		require.True(t, ok)
+		require.Equal(t, ea.Addr, addrs[i])
+
+		ea, ok = pa.PopIfExpired(expiries[i])
+		require.False(t, ok)
+		require.Nil(t, ea)
+	}
+
+}