provider: refresh schedule (libp2p#1131)

* provider: refresh schedule

* address review

Author: guillaumemichel

provider/datastore/keystore.go

@@ -369,6 +369,62 @@ func (s *KeyStore) Get(ctx context.Context, prefix bitstr.Key) ([]mh.Multihash,
 	return result, nil
 }
 
+// ContainsPrefix reports whether the KeyStore currently holds at least one
+// multihash whose kademlia identifier (bit256.Key) starts with the provided
+// bit-prefix.
+func (s *KeyStore) ContainsPrefix(ctx context.Context, prefix bitstr.Key) (bool, error) {
+	if s.closed() {
+		return false, ErrKeyStoreClosed
+	}
+	s.lk.Lock()
+	defer s.lk.Unlock()
+
+	// Case 1: shorter than bucket length — any bucket with this key-prefix is enough.
+	if len(prefix) < s.prefixLen {
+		rem := s.prefixLen - len(prefix)
+		limit := 1 << rem
+		for i := range limit {
+			suffix := fmt.Sprintf("%0*b", rem, i)
+			dsKey := s.dsKey(prefix + bitstr.Key(suffix))
+			exists, err := s.ds.Has(ctx, dsKey)
+			if err != nil {
+				return false, err
+			}
+			if exists {
+				return true, nil
+			}
+		}
+		return false, nil
+	}
+
+	// Case 2: at least a full bucket prefix — check that bucket's content.
+	dsKey := s.dsKey(bitstr.Key(prefix[:s.prefixLen]))
+
+	// Fast path when asking exactly for a bucket prefix: existence implies non-empty.
+	if len(prefix) == s.prefixLen {
+		return s.ds.Has(ctx, dsKey)
+	}
+
+	// Longer-than-bucket: must inspect entries and see if any starts with `prefix`.
+	data, err := s.ds.Get(ctx, dsKey)
+	if err != nil {
+		if err == ds.ErrNotFound {
+			return false, nil
+		}
+		return false, err
+	}
+	var stored []mh.Multihash
+	if err := json.Unmarshal(data, &stored); err != nil {
+		return false, err
+	}
+	for _, h := range stored {
+		if keyspace.IsPrefix(prefix, keyspace.MhToBit256(h)) {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
 // emptyLocked deletes all entries under the datastore prefix, assuming s.lk is
 // already held.
 func (s *KeyStore) emptyLocked(ctx context.Context) error {

provider/datastore/keystore_test.go

@@ -2,12 +2,15 @@ package datastore
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	"github.com/ipfs/go-cid"
 	ds "github.com/ipfs/go-datastore"
+	"github.com/ipfs/go-test/random"
 	"github.com/libp2p/go-libp2p-kad-dht/provider/internal/keyspace"
 	mh "github.com/multiformats/go-multihash"
+	"github.com/stretchr/testify/require"
 
 	"github.com/probe-lab/go-libdht/kad/key"
 	"github.com/probe-lab/go-libdht/kad/key/bitstr"
@@ -85,6 +88,60 @@ func TestKeyStoreStoreAndGet(t *testing.T) {
 	}
 }
 
+func genMultihashesMatchingPrefix(prefix bitstr.Key, n int) []mh.Multihash {
+	mhs := make([]mh.Multihash, 0, n)
+	for i := 0; len(mhs) < n; i++ {
+		h := random.Multihashes(1)[0]
+		k := keyspace.MhToBit256(h)
+		if keyspace.IsPrefix(prefix, k) {
+			mhs = append(mhs, h)
+		}
+	}
+	return mhs
+}
+
+func TestKeyStoreContainsPrefix(t *testing.T) {
+	ctx := context.Background()
+	store, err := NewKeyStore(ds.NewMapDatastore())
+	require.NoError(t, err)
+
+	ok, err := store.ContainsPrefix(ctx, bitstr.Key("0000"))
+	require.NoError(t, err)
+	require.False(t, ok)
+
+	generated := genMultihashesMatchingPrefix(bitstr.Key(strings.Repeat("0", DefaultKeyStorePrefixBits+4)), 1)
+	require.True(t, keyspace.IsPrefix(bitstr.Key("0000"), keyspace.MhToBit256(generated[0])))
+	store.Put(ctx, generated...)
+
+	ok, err = store.ContainsPrefix(ctx, bitstr.Key("0"))
+	require.NoError(t, err)
+	require.True(t, ok)
+
+	ok, err = store.ContainsPrefix(ctx, bitstr.Key("0000"))
+	require.NoError(t, err)
+	require.True(t, ok)
+
+	ok, err = store.ContainsPrefix(ctx, bitstr.Key(strings.Repeat("0", DefaultKeyStorePrefixBits)))
+	require.NoError(t, err)
+	require.True(t, ok)
+
+	ok, err = store.ContainsPrefix(ctx, bitstr.Key(strings.Repeat("0", DefaultKeyStorePrefixBits+4)))
+	require.NoError(t, err)
+	require.True(t, ok)
+
+	ok, err = store.ContainsPrefix(ctx, bitstr.Key("1"))
+	require.NoError(t, err)
+	require.False(t, ok)
+
+	ok, err = store.ContainsPrefix(ctx, bitstr.Key("0001"))
+	require.NoError(t, err)
+	require.False(t, ok)
+
+	ok, err = store.ContainsPrefix(ctx, bitstr.Key(strings.Repeat("0", DefaultKeyStorePrefixBits+2)+"1"))
+	require.NoError(t, err)
+	require.False(t, ok)
+}
+
 func TestKeyStoreReset(t *testing.T) {
 	store, err := NewKeyStore(ds.NewMapDatastore())
 	if err != nil {

provider/internal/keyspace/key.go

@@ -131,6 +131,42 @@ func ShortestCoveredPrefix(target bitstr.Key, peers []peer.ID) (bitstr.Key, []pe
 	return target[:coveredCpl], peers[:lastCoveredPeerIndex]
 }
 
+// ExtendBinaryPrefix returns all bitstrings of length n that start with prefix.
+// Example: prefix="1101", n=6 -> ["110100", "110101", "110110", "110111"].
+func ExtendBinaryPrefix(prefix bitstr.Key, n int) []bitstr.Key {
+	extraBits := n - len(prefix)
+	if n < 0 || extraBits < 0 {
+		return nil
+	}
+
+	extLen := 1 << extraBits // 2^extraBits
+	rd := make([]bitstr.Key, 0, extLen)
+	wr := make([]bitstr.Key, 1, extLen)
+	wr[0] = prefix
+
+	// Iteratively append bits until reaching length n.
+	for range extraBits {
+		rd, wr = wr, rd[:0]
+		for _, s := range rd {
+			wr = append(wr, s+"0", s+"1")
+		}
+	}
+	return wr
+}
+
+// SiblingPrefixes returns the prefixes of the sibling subtrees along the path
+// to key. Together with the subtree under `key` itself, these prefixes
+// partition the keyspace.
+//
+// For key "1100" it returns: ["0", "10", "111", "1101"].
+func SiblingPrefixes(key bitstr.Key) []bitstr.Key {
+	complements := make([]bitstr.Key, len(key))
+	for i := range key {
+		complements[i] = FlipLastBit(key[:i+1])
+	}
+	return complements
+}
+
 // PrefixAndKeys is a struct that holds a prefix and the multihashes whose
 // kademlia identifier share the same prefix.
 type PrefixAndKeys struct {

provider/internal/keyspace/key_test.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/ipfs/go-test/random"
 	kb "github.com/libp2p/go-libp2p-kbucket"
 	"github.com/libp2p/go-libp2p/core/peer"
 	mh "github.com/multiformats/go-multihash"
@@ -178,9 +179,7 @@ func TestShortestCoveredPrefix(t *testing.T) {
 	for range nIterations {
 		minCpl := KeyLen
 		largestCplCount := 0
-		for i := range peers {
-			peers[i] = genRandPeerID(t)
-		}
+		peers = random.Peers(nPeers)
 		peers = kb.SortClosestPeers(peers, target[:])
 		for i := range peers {
 			cpl = kb.CommonPrefixLen(kb.ConvertPeerID(peers[i]), target[:])
@@ -203,6 +202,30 @@ func TestShortestCoveredPrefix(t *testing.T) {
 	require.Empty(t, coveredPeers)
 }
 
+func TestExtendBinaryPrefix(t *testing.T) {
+	prefix := bitstr.Key("")
+	l := 1
+	require.Equal(t, []bitstr.Key{"0", "1"}, ExtendBinaryPrefix(prefix, l))
+	prefix = bitstr.Key("1101")
+	l = 6
+	require.Equal(t, []bitstr.Key{"110100", "110101", "110110", "110111"}, ExtendBinaryPrefix(prefix, l))
+}
+
+func TestSiblingPrefixes(t *testing.T) {
+	k := bitstr.Key("")
+	require.Empty(t, SiblingPrefixes(k))
+	k = bitstr.Key("0")
+	require.Equal(t, []bitstr.Key{"1"}, SiblingPrefixes(k))
+	k = bitstr.Key("1")
+	require.Equal(t, []bitstr.Key{"0"}, SiblingPrefixes(k))
+	k = bitstr.Key("00")
+	require.Equal(t, []bitstr.Key{"1", "01"}, SiblingPrefixes(k))
+	k = bitstr.Key("000")
+	require.Equal(t, []bitstr.Key{"1", "01", "001"}, SiblingPrefixes(k))
+	k = bitstr.Key("1100")
+	require.Equal(t, []bitstr.Key{"0", "10", "111", "1101"}, SiblingPrefixes(k))
+}
+
 func genMultihashes(n int) []mh.Multihash {
 	mhs := make([]mh.Multihash, n)
 	for i := range mhs {

provider/internal/keyspace/trie.go

@@ -3,6 +3,7 @@ package keyspace
 import (
 	"github.com/libp2p/go-libp2p/core/peer"
 	mh "github.com/multiformats/go-multihash"
+
 	"github.com/probe-lab/go-libdht/kad"
 	"github.com/probe-lab/go-libdht/kad/key"
 	"github.com/probe-lab/go-libdht/kad/key/bit256"
@@ -180,6 +181,47 @@ func pruneSubtrieAtDepth[K0 kad.Key[K0], K1 kad.Key[K1], D any](t *trie.Trie[K0,
 	return false
 }
 
+// TrieGaps returns all prefixes that aren't covered by a key (prefix) in the
+// trie. Combining the prefixes included in the trie with the gap prefixes
+// results in a full keyspace coverage.
+//
+// E.g Trie: ["00", "100"], GapsInTrie: ["01", "101", "11"]
+func TrieGaps[D any](t *trie.Trie[bitstr.Key, D]) []bitstr.Key {
+	if t.IsLeaf() {
+		if t.HasKey() {
+			return SiblingPrefixes(*t.Key())
+		}
+		return []bitstr.Key{""}
+	}
+	return trieGapsAtDepth(t, 0)
+}
+
+func trieGapsAtDepth[D any](t *trie.Trie[bitstr.Key, D], depth int) []bitstr.Key {
+	var gaps []bitstr.Key
+	for i := range 2 {
+		bstr := bitstr.Key(byte('0' + i))
+		if b := t.Branch(i); b == nil {
+			gaps = append(gaps, bstr)
+		} else if b.IsLeaf() {
+			if b.HasKey() {
+				k := *b.Key()
+				if len(k) > depth+1 {
+					for _, siblingPrefix := range SiblingPrefixes(k)[depth+1:] {
+						gaps = append(gaps, siblingPrefix[depth:])
+					}
+				}
+			} else {
+				gaps = append(gaps, bstr)
+			}
+		} else {
+			for _, gap := range trieGapsAtDepth(b, depth+1) {
+				gaps = append(gaps, bstr+gap)
+			}
+		}
+	}
+	return gaps
+}
+
 // mapMerge merges all key-value pairs from the source map into the destination
 // map. Values from the source are appended to existing slices in the
 // destination.