provider: minor fixes (#1133)

* provider: minor fixes

* adjusting for kubo PR

* address review

Author: guillaumemichel

dual/provider/provider.go

@@ -3,7 +3,6 @@ package provider
 import (
 	"context"
 	"errors"
-	"sync"
 	"sync/atomic"
 
 	"github.com/ipfs/go-cid"
@@ -84,18 +83,19 @@ func New(d *dual.DHT, opts ...Option) (*SweepingProvider, error) {
 
 // runOnBoth runs the provided function on both the LAN and WAN providers in
 // parallel and waits for both to complete.
-func (s *SweepingProvider) runOnBoth(f func(*provider.SweepingProvider)) {
-	wg := sync.WaitGroup{}
-	wg.Add(2)
-	go func() {
-		defer wg.Done()
-		f(s.LAN)
-	}()
-	go func() {
-		defer wg.Done()
+func (s *SweepingProvider) runOnBoth(wait bool, f func(*provider.SweepingProvider)) {
+	if wait {
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			f(s.LAN)
+		}()
 		f(s.WAN)
-	}()
-	wg.Wait()
+		<-done
+		return
+	}
+	go f(s.LAN)
+	go f(s.WAN)
 }
 
 // ProvideOnce sends provider records for the specified keys to both DHT swarms
@@ -104,7 +104,7 @@ func (s *SweepingProvider) runOnBoth(f func(*provider.SweepingProvider)) {
 // Add the supplied multihashes to the provide queue, and return immediately.
 // The provide operation happens asynchronously.
 func (s *SweepingProvider) ProvideOnce(keys ...mh.Multihash) {
-	s.runOnBoth(func(p *provider.SweepingProvider) {
+	s.runOnBoth(false, func(p *provider.SweepingProvider) {
 		p.ProvideOnce(keys...)
 	})
 }
@@ -128,7 +128,7 @@ func (s *SweepingProvider) StartProviding(force bool, keys ...mh.Multihash) {
 		return
 	}
 
-	s.runOnBoth(func(p *provider.SweepingProvider) {
+	s.runOnBoth(false, func(p *provider.SweepingProvider) {
 		p.AddToSchedule(newKeys...)
 	})
 
@@ -160,7 +160,7 @@ func (s *SweepingProvider) StopProviding(keys ...mh.Multihash) {
 // reprovided as scheduled.
 func (s *SweepingProvider) Clear() int {
 	var total atomic.Int32
-	s.runOnBoth(func(p *provider.SweepingProvider) {
+	s.runOnBoth(true, func(p *provider.SweepingProvider) {
 		total.Add(int32(p.Clear()))
 	})
 	return int(total.Load())
@@ -174,7 +174,7 @@ func (s *SweepingProvider) Clear() int {
 // This is done automatically during the reprovide operation if a region has no
 // keys.
 func (s *SweepingProvider) RefreshSchedule() {
-	s.runOnBoth(func(p *provider.SweepingProvider) {
+	go s.runOnBoth(false, func(p *provider.SweepingProvider) {
 		p.RefreshSchedule()
 	})
 }

fullrt/dht.go

@@ -150,6 +150,7 @@ func NewFullRT(h host.Host, protocolPrefix protocol.ID, options ...Option) (*Ful
 		EnableProviders:  true,
 		EnableValues:     true,
 		ProtocolPrefix:   protocolPrefix,
+		MsgSenderBuilder: net.NewMessageSenderImpl,
 	}
 
 	if err := dhtcfg.Apply(fullrtcfg.dhtOpts...); err != nil {
@@ -163,7 +164,7 @@ func NewFullRT(h host.Host, protocolPrefix protocol.ID, options ...Option) (*Ful
 		return nil, err
 	}
 
-	ms := net.NewMessageSenderImpl(h, amino.Protocols)
+	ms := dhtcfg.MsgSenderBuilder(h, amino.Protocols)
 	protoMessenger, err := dht_pb.NewProtocolMessenger(ms)
 	if err != nil {
 		return nil, err

provider/datastore/keystore.go

@@ -46,6 +46,7 @@ type KeyStore struct {
 	prefixLen int
 
 	gcFunc      KeyChanFunc // optional function to get keys for garbage collection
+	gcFuncLk    sync.Mutex
 	gcInterval  time.Duration
 	gcBatchSize int
 }
@@ -103,17 +104,6 @@ func WithDatastorePrefix(base string) KeyStoreOption {
 	}
 }
 
-// WithGCFunc sets the function periodically used to garbage collect keys that
-// shouldn't be provided anymore. During the garbage collection process, the
-// store is entirely purged, and is repopulated from the keys supplied by the
-// KeyChanFunc.
-func WithGCFunc(gcFunc KeyChanFunc) KeyStoreOption {
-	return func(cfg *keyStoreCfg) error {
-		cfg.gcFunc = gcFunc
-		return nil
-	}
-}
-
 // WithGCInterval defines the interval at which the KeyStore is garbage
 // collected. During the garbage collection process, the store is entirely
 // purged, and is repopulated from the keys supplied by the gcFunc.
@@ -141,6 +131,33 @@ func WithGCBatchSize(size int) KeyStoreOption {
 	}
 }
 
+// WithGCFunc sets the function periodically used to garbage collect keys that
+// shouldn't be provided anymore. During the garbage collection process, the
+// store is entirely purged, and is repopulated from the keys supplied by the
+// KeyChanFunc.
+func WithGCFunc(gcFunc KeyChanFunc) KeyStoreOption {
+	return func(cfg *keyStoreCfg) error {
+		cfg.gcFunc = gcFunc
+		return nil
+	}
+}
+
+func (s *KeyStore) SetGCFunc(gcFunc KeyChanFunc) {
+	if gcFunc == nil {
+		return
+	}
+	s.gcFuncLk.Lock()
+	defer s.gcFuncLk.Unlock()
+
+	startGC := s.gcFunc == nil && s.gcInterval > 0
+	s.gcFunc = gcFunc
+
+	if startGC {
+		s.wg.Add(1)
+		go s.runGC()
+	}
+}
+
 // NewKeyStore creates a new KeyStore backed by the provided datastore.
 func NewKeyStore(d ds.Batching, opts ...KeyStoreOption) (*KeyStore, error) {
 	var cfg keyStoreCfg
@@ -161,6 +178,7 @@ func NewKeyStore(d ds.Batching, opts ...KeyStoreOption) (*KeyStore, error) {
 		gcBatchSize: cfg.gcBatchSize,
 	}
 	if cfg.gcFunc != nil && cfg.gcInterval > 0 {
+		keyStore.wg.Add(1)
 		go keyStore.runGC()
 	}
 	return &keyStore, nil
@@ -188,15 +206,16 @@ func (s *KeyStore) closed() bool {
 // resets the state of the KeyStore to match the state returned by the GC
 // function.
 func (s *KeyStore) runGC() {
-	s.wg.Add(1)
 	defer s.wg.Done()
 	ticker := time.NewTicker(s.gcInterval)
 	for {
 		select {
 		case <-s.done:
 			return
 		case <-ticker.C:
+			s.gcFuncLk.Lock()
 			keysChan, err := s.gcFunc(context.Background())
+			s.gcFuncLk.Unlock()
 			if err != nil {
 				logger.Errorf("garbage collection failed: %v", err)
 				continue
@@ -218,17 +237,14 @@ func (s *KeyStore) ResetCids(ctx context.Context, keysChan <-chan cid.Cid) error
 	if err != nil {
 		return fmt.Errorf("KeyStore empty failed during reset: %w", err)
 	}
-	keys := make([]mh.Multihash, s.gcBatchSize)
-	i := 0
+	keys := make([]mh.Multihash, 0, s.gcBatchSize)
 	for c := range keysChan {
-		keys[i] = c.Hash()
-		i++
-		if i == s.gcBatchSize {
+		keys = append(keys, c.Hash())
+		if len(keys) == s.gcBatchSize {
 			_, err = s.Put(ctx, keys...)
 			if err != nil {
 				return fmt.Errorf("KeyStore put failed during reset: %w", err)
 			}
-			i = 0
 			keys = keys[:0]
 		}
 	}
@@ -245,7 +261,7 @@ func (s *KeyStore) dsKey(prefix bitstr.Key) ds.Key {
 
 // putLocked stores the provided keys while assuming s.lk is already held.
 func (s *KeyStore) putLocked(ctx context.Context, keys ...mh.Multihash) ([]mh.Multihash, error) {
-	groups := make(map[bitstr.Key][]mh.Multihash)
+	groups := make(map[bitstr.Key][]mh.Multihash, len(keys))
 	for _, h := range keys {
 		k := keyspace.MhToBit256(h)
 		bs := bitstr.Key(key.BitString(k)[:s.prefixLen])

provider/internal/keyspace/trie.go

@@ -50,13 +50,16 @@ func FindPrefixOfKey[K0 kad.Key[K0], K1 kad.Key[K1], D any](t *trie.Trie[K0, D],
 }
 
 func findPrefixOfKeyAtDepth[K0 kad.Key[K0], K1 kad.Key[K1], D any](t *trie.Trie[K0, D], k K1, depth int) (K0, bool) {
+	var zero K0
 	if t.IsLeaf() {
 		if !t.HasKey() {
-			var zero K0
 			return zero, false
 		}
 		return *t.Key(), key.CommonPrefixLength(*t.Key(), k) == (*t.Key()).BitLen()
 	}
+	if depth == k.BitLen() {
+		return zero, false
+	}
 	b := int(k.Bit(depth))
 	return findPrefixOfKeyAtDepth(t.Branch(b), k, depth+1)
 }

provider/internal/keyspace/trie_test.go

@@ -89,6 +89,21 @@ func TestFindPrefixOfKey(t *testing.T) {
 	require.False(t, ok)
 }
 
+func TestFindPrefixOfTooShortKey(t *testing.T) {
+	tr := trie.New[bitstr.Key, struct{}]()
+	keys := []bitstr.Key{
+		"0000",
+		"0001",
+		"0010",
+		"0011",
+	}
+	for _, k := range keys {
+		tr.Add(k, struct{}{})
+	}
+	_, ok := FindPrefixOfKey(tr, bitstr.Key("000"))
+	require.False(t, ok)
+}
+
 func TestFindSubtrie(t *testing.T) {
 	keys := []bitstr.Key{
 		"0000",

provider/provider.go

@@ -46,6 +46,10 @@ const (
 	// the measurement, but adds network load and latency before the initial
 	// provide request can be performed.
 	initialGetClosestPeersCount = 4
+	// defaultPrefixLenValidity is the validity of the cached average region
+	// prefix length computed from the schedule. It allows to avoid recomputing
+	// the average everytime we need to average prefix length.
+	defaultPrefixLenValidity = 5 * time.Minute
 
 	// retryInterval is the interval at which the provider tries to perform any
 	// previously failed work (provide or reprovide).
@@ -197,7 +201,7 @@ func New(opts ...Option) (*SweepingProvider, error) {
 		}),
 		maxProvideConnsPerWorker: cfg.maxProvideConnsPerWorker,
 
-		avgPrefixLenValidity: 5 * time.Minute,
+		avgPrefixLenValidity: defaultPrefixLenValidity,
 		cachedAvgPrefixLen:   -1,
 		avgPrefixLenReady:    make(chan struct{}),
 
@@ -380,7 +384,7 @@ func (s *SweepingProvider) unscheduleSubsumedPrefixesNoLock(prefix bitstr.Key) {
 		} else {
 			timeUntilReprovide := s.timeUntil(next.Data)
 			s.scheduleNextReprovideNoLock(next.Key, timeUntilReprovide)
-			logger.Warnf("next scheduled prefix now is %s", s.scheduleCursor)
+			logger.Debugf("next scheduled prefix now is %s", s.scheduleCursor)
 		}
 	}
 }
@@ -1083,7 +1087,7 @@ func (s *SweepingProvider) batchProvide(prefix bitstr.Key, keys []mh.Multihash)
 	regions = keyspace.AssignKeysToRegions(regions, keys)
 
 	if !s.provideRegions(regions, addrInfo, false, false) {
-		logger.Errorf("failed to provide any region for prefix %s", prefix)
+		logger.Warnf("failed to provide any region for prefix %s", prefix)
 	}
 }
 
@@ -1155,20 +1159,20 @@ func (s *SweepingProvider) batchReprovide(prefix bitstr.Key, periodicReprovide b
 	regions = keyspace.AssignKeysToRegions(regions, keys)
 
 	if !s.provideRegions(regions, addrInfo, true, periodicReprovide) {
-		logger.Errorf("failed to reprovide any region for prefix %s", prefix)
+		logger.Warnf("failed to reprovide any region for prefix %s", prefix)
 	}
 }
 
 func (s *SweepingProvider) failedProvide(prefix bitstr.Key, keys []mh.Multihash, err error) {
-	logger.Error(err)
+	logger.Warn(err)
 	// Put keys back to the provide queue.
 	s.provideQueue.Enqueue(prefix, keys...)
 
 	s.connectivity.TriggerCheck()
 }
 
 func (s *SweepingProvider) failedReprovide(prefix bitstr.Key, err error) {
-	logger.Error(err)
+	logger.Warn(err)
 	// Put prefix in the reprovide queue.
 	s.reprovideQueue.Enqueue(prefix)