Merge branch 'master' into provider-stats

Author: guillaumemichel

provider/buffered/provider.go

@@ -29,7 +29,7 @@ const (
 
 var _ internal.Provider = (*SweepingProvider)(nil)
 
-// buffered.SweepingProvider is a wrapper around a SweepingProvider buffering
+// SweepingProvider (buffered) is a wrapper around a SweepingProvider buffering
 // requests, to allow core operations to return instantly. Operations are
 // queued and processed asynchronously in batches for improved performance.
 type SweepingProvider struct {

provider/dual/options.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"time"
 
-	ds "github.com/ipfs/go-datastore"
 	"github.com/libp2p/go-libp2p-kad-dht/amino"
 	"github.com/libp2p/go-libp2p-kad-dht/dual"
 	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
@@ -38,53 +37,44 @@ type config struct {
 
 type Option func(opt *config) error
 
-func (cfg *config) apply(opts ...Option) error {
-	for i, o := range opts {
-		if err := o(cfg); err != nil {
-			return fmt.Errorf("dual dht provider option %d failed: %w", i, err)
+// getOpts creates a config and applies Options to it.
+func getOpts(opts []Option, d *dual.DHT) (config, error) {
+	cfg := config{
+		reprovideInterval: [2]time.Duration{amino.DefaultReprovideInterval, amino.DefaultReprovideInterval},
+		maxReprovideDelay: [2]time.Duration{provider.DefaultMaxReprovideDelay, provider.DefaultMaxReprovideDelay},
+
+		offlineDelay:                    [2]time.Duration{provider.DefaultOfflineDelay, provider.DefaultOfflineDelay},
+		connectivityCheckOnlineInterval: [2]time.Duration{provider.DefaultConnectivityCheckOnlineInterval, provider.DefaultConnectivityCheckOnlineInterval},
+
+		maxWorkers:               [2]int{4, 4},
+		dedicatedPeriodicWorkers: [2]int{2, 2},
+		dedicatedBurstWorkers:    [2]int{1, 1},
+		maxProvideConnsPerWorker: [2]int{20, 20},
+	}
+
+	// Apply options
+	for i, opt := range opts {
+		if err := opt(&cfg); err != nil {
+			return config{}, fmt.Errorf("dual dht provider option %d failed: %w", i, err)
 		}
 	}
-	return nil
-}
 
-func (cfg *config) resolveDefaults(d *dual.DHT) {
+	// Resolve defaults
 	if cfg.msgSenders[lanID] == nil {
 		cfg.msgSenders[lanID] = d.LAN.MessageSender()
 	}
 	if cfg.msgSenders[wanID] == nil {
 		cfg.msgSenders[wanID] = d.WAN.MessageSender()
 	}
-}
 
-func (c *config) validate() error {
-	if c.dedicatedPeriodicWorkers[lanID]+c.dedicatedBurstWorkers[lanID] > c.maxWorkers[lanID] {
-		return errors.New("provider config: total dedicated workers exceed max workers")
+	// Validate config
+	if cfg.dedicatedPeriodicWorkers[lanID]+cfg.dedicatedBurstWorkers[lanID] > cfg.maxWorkers[lanID] {
+		return config{}, errors.New("provider config: total dedicated workers exceed max workers")
 	}
-	if c.dedicatedPeriodicWorkers[wanID]+c.dedicatedBurstWorkers[wanID] > c.maxWorkers[wanID] {
-		return errors.New("provider config: total dedicated workers exceed max workers")
+	if cfg.dedicatedPeriodicWorkers[wanID]+cfg.dedicatedBurstWorkers[wanID] > cfg.maxWorkers[wanID] {
+		return config{}, errors.New("provider config: total dedicated workers exceed max workers")
 	}
-	return nil
-}
-
-var DefaultConfig = func(cfg *config) error {
-	var err error
-	cfg.keystore, err = keystore.NewKeystore(ds.NewMapDatastore())
-	if err != nil {
-		return err
-	}
-
-	cfg.reprovideInterval = [2]time.Duration{amino.DefaultReprovideInterval, amino.DefaultReprovideInterval}
-	cfg.maxReprovideDelay = [2]time.Duration{provider.DefaultMaxReprovideDelay, provider.DefaultMaxReprovideDelay}
-
-	cfg.offlineDelay = [2]time.Duration{provider.DefaultOfflineDelay, provider.DefaultOfflineDelay}
-	cfg.connectivityCheckOnlineInterval = [2]time.Duration{provider.DefaultConnectivityCheckOnlineInterval, provider.DefaultConnectivityCheckOnlineInterval}
-
-	cfg.maxWorkers = [2]int{4, 4}
-	cfg.dedicatedPeriodicWorkers = [2]int{2, 2}
-	cfg.dedicatedBurstWorkers = [2]int{1, 1}
-	cfg.maxProvideConnsPerWorker = [2]int{20, 20}
-
-	return nil
+	return cfg, nil
 }
 
 func WithKeystore(ks keystore.Keystore) Option {

provider/dual/provider.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/ipfs/go-cid"
+	"github.com/ipfs/go-datastore"
 	dht "github.com/libp2p/go-libp2p-kad-dht"
 	"github.com/libp2p/go-libp2p-kad-dht/dual"
 	"github.com/libp2p/go-libp2p-kad-dht/provider"
@@ -23,6 +24,8 @@ type SweepingProvider struct {
 	LAN      *provider.SweepingProvider
 	WAN      *provider.SweepingProvider
 	keystore keystore.Keystore
+
+	cleanupFuncs []func() error
 }
 
 // New creates a new SweepingProvider that manages provides and reprovides for
@@ -32,15 +35,19 @@ func New(d *dual.DHT, opts ...Option) (*SweepingProvider, error) {
 		return nil, errors.New("cannot create sweeping provider for nil dual DHT")
 	}
 
-	var cfg config
-	err := cfg.apply(append([]Option{DefaultConfig}, opts...)...)
+	cfg, err := getOpts(opts, d)
 	if err != nil {
 		return nil, err
 	}
-	cfg.resolveDefaults(d)
-	err = cfg.validate()
-	if err != nil {
-		return nil, err
+	var cleanupFuncs []func() error
+	if cfg.keystore == nil {
+		ds := datastore.NewMapDatastore()
+		cfg.keystore, err = keystore.NewKeystore(ds)
+		if err != nil {
+			ds.Close()
+			return nil, fmt.Errorf("couldn't create a keystore: %w", err)
+		}
+		cleanupFuncs = []func() error{ds.Close, cfg.keystore.Close}
 	}
 
 	sweepingProviders := make([]*provider.SweepingProvider, 2)
@@ -74,10 +81,11 @@ func New(d *dual.DHT, opts ...Option) (*SweepingProvider, error) {
 	}
 
 	return &SweepingProvider{
-		dht:      d,
-		LAN:      sweepingProviders[0],
-		WAN:      sweepingProviders[1],
-		keystore: cfg.keystore,
+		dht:          d,
+		LAN:          sweepingProviders[0],
+		WAN:          sweepingProviders[1],
+		keystore:     cfg.keystore,
+		cleanupFuncs: cleanupFuncs,
 	}, nil
 }
 
@@ -102,9 +110,25 @@ func (s *SweepingProvider) runOnBoth(f func(*provider.SweepingProvider) error) e
 
 // Close stops both DHT providers and releases associated resources.
 func (s *SweepingProvider) Close() error {
-	return s.runOnBoth(func(p *provider.SweepingProvider) error {
+	err := s.runOnBoth(func(p *provider.SweepingProvider) error {
 		return p.Close()
 	})
+
+	if s.cleanupFuncs != nil {
+		// Cleanup keystore and datastore if we created them
+		var errs []error
+		for i := len(s.cleanupFuncs) - 1; i >= 0; i-- { // LIFO: last-added is cleaned up first
+			if f := s.cleanupFuncs[i]; f != nil {
+				if err := f(); err != nil {
+					errs = append(errs, err)
+				}
+			}
+		}
+		if len(errs) > 0 {
+			err = errors.Join(append(errs, err)...)
+		}
+	}
+	return err
 }
 
 // ProvideOnce sends provider records for the specified keys to both DHT swarms

provider/internal/connectivity/connectivity.go

@@ -55,8 +55,7 @@ type ConnectivityChecker struct {
 
 // New creates a new ConnectivityChecker instance.
 func New(checkFunc func() bool, opts ...Option) (*ConnectivityChecker, error) {
-	var cfg config
-	err := cfg.apply(append([]Option{DefaultConfig}, opts...)...)
+	cfg, err := getOpts(opts)
 	if err != nil {
 		return nil, err
 	}

provider/internal/connectivity/options.go

@@ -14,21 +14,21 @@ type config struct {
 	onOnline  func()
 }
 
-func (cfg *config) apply(opts ...Option) error {
-	for i, o := range opts {
-		if err := o(cfg); err != nil {
-			return fmt.Errorf("reprovider dht option %d failed: %w", i, err)
-		}
-	}
-	return nil
-}
-
 type Option func(opt *config) error
 
-var DefaultConfig = func(cfg *config) error {
-	cfg.onlineCheckInterval = 1 * time.Minute
-	cfg.offlineDelay = 2 * time.Hour
-	return nil
+// getOpts creates a config and applies Options to it.
+func getOpts(opts []Option) (config, error) {
+	cfg := config{
+		onlineCheckInterval: 1 * time.Minute,
+		offlineDelay:        2 * time.Hour,
+	}
+
+	for i, opt := range opts {
+		if err := opt(&cfg); err != nil {
+			return config{}, fmt.Errorf("connectivity option %d error: %s", i, err)
+		}
+	}
+	return cfg, nil
 }
 
 // WithOnlineCheckInterval sets the minimum interval between online checks.

provider/internal/keyspace/trie.go

@@ -1,6 +1,8 @@
 package keyspace
 
 import (
+	"slices"
+
 	"github.com/libp2p/go-libp2p/core/peer"
 	mh "github.com/multiformats/go-multihash"
 
@@ -185,46 +187,92 @@ func pruneSubtrieAtDepth[K0 kad.Key[K0], K1 kad.Key[K1], D any](t *trie.Trie[K0,
 }
 
 // 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.
+// trie, at the `target` location. Combining the prefixes included in the trie
+// with the gap prefixes results in a full keyspace coverage of the `target`
+// subtrie.
+//
+// Results are sorted according to the provided `order`.
 //
-// E.g Trie: ["00", "100"], GapsInTrie: ["01", "101", "11"]
-func TrieGaps[D any](t *trie.Trie[bitstr.Key, D]) []bitstr.Key {
+// Example:
+//   - Trie: ["0000", "0010", "100"]
+//   - Target: "0"
+//   - Order: "0000",
+//   - GapsInTrie: ["0001", "0011", "01"]
+func TrieGaps[K kad.Key[K], D any](t *trie.Trie[bitstr.Key, D], target bitstr.Key, order K) []bitstr.Key {
 	if t.IsLeaf() {
-		if t.HasKey() {
-			return SiblingPrefixes(*t.Key())
+		if k := t.Key(); k != nil {
+			if IsBitstrPrefix(target, *k) {
+				siblingPrefixes := SiblingPrefixes(*k)[len(target):]
+				sortBitstrKeysByOrder(siblingPrefixes, order)
+				return siblingPrefixes
+			}
+			if IsBitstrPrefix(*k, target) {
+				// The only key in the trie is a prefix of target, meaning the whole
+				// target is covered.
+				return nil
+			}
 		}
-		return []bitstr.Key{""}
+		return []bitstr.Key{target}
 	}
-	return trieGapsAtDepth(t, 0)
+	return trieGapsAtDepth(t, 0, target, order)
 }
 
-func trieGapsAtDepth[D any](t *trie.Trie[bitstr.Key, D], depth int) []bitstr.Key {
+func trieGapsAtDepth[K kad.Key[K], D any](t *trie.Trie[bitstr.Key, D], depth int, target bitstr.Key, order K) []bitstr.Key {
 	var gaps []bitstr.Key
-	for i := range 2 {
+	insideTarget := depth >= target.BitLen()
+	b := int(order.Bit(depth))
+	for _, i := range []int{b, 1 - b} {
+		if !insideTarget && i != int(target.Bit(depth)) {
+			continue
+		}
 		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:] {
+					siblingPrefixes := SiblingPrefixes(k)[depth+1:]
+					sortBitstrKeysByOrder(siblingPrefixes, order)
+					for _, siblingPrefix := range siblingPrefixes {
 						gaps = append(gaps, siblingPrefix[depth:])
 					}
 				}
 			} else {
 				gaps = append(gaps, bstr)
 			}
 		} else {
-			for _, gap := range trieGapsAtDepth(b, depth+1) {
+			for _, gap := range trieGapsAtDepth(b, depth+1, target, order) {
 				gaps = append(gaps, bstr+gap)
 			}
 		}
 	}
 	return gaps
 }
 
+// sortBitstrKeysByOrder sorts the provided bitstr keys according to the
+// provided order.
+func sortBitstrKeysByOrder[K kad.Key[K]](keys []bitstr.Key, order K) {
+	slices.SortFunc(keys, func(a, b bitstr.Key) int {
+		maxLen := min(len(a), len(b), order.BitLen())
+		for i := range maxLen {
+			if a[i] != b[i] {
+				if a.Bit(i) == order.Bit(i) {
+					return -1
+				}
+				return 1
+			}
+		}
+		if len(a) == len(b) || maxLen == order.BitLen() {
+			return 0
+		}
+		if len(a) < len(b) {
+			return 1
+		}
+		return -1
+	})
+}
+
 // 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.