custom ProviderManager that brokers AddrInfos

dht.go

@@ -84,8 +84,8 @@ type IpfsDHT struct {
  datastore ds.Datastore // Local data
 
  routingTable *kb.RoutingTable // Array of routing tables for differently distanced nodes
- // ProviderManager stores & manages the provider records for this Dht peer.
- ProviderManager *providers.ProviderManager
+ // providerStore stores & manages the provider records for this Dht peer.
+ providerStore providers.ProviderStore
 
  // manages Routing Table refresh
  rtRefreshManager *rtrefresh.RtRefreshManager
@@ -221,7 +221,9 @@ func New(ctx context.Context, h host.Host, options ...Option) (*IpfsDHT, error)
  }
  dht.proc.Go(sn.subscribe)
  // handle providers
- dht.proc.AddChild(dht.ProviderManager.Process())
+ if mgr, ok := dht.providerStore.(interface{ Process() goprocess.Process }); ok {
+ dht.proc.AddChild(mgr.Process())
+ }
 
  // go-routine to make sure we ALWAYS have RT peer addresses in the peerstore
  // since RT membership is decoupled from connectivity
@@ -338,11 +340,14 @@ func makeDHT(ctx context.Context, h host.Host, cfg dhtcfg.Config) (*IpfsDHT, err
  // the DHT context should be done when the process is closed
  dht.ctx = goprocessctx.WithProcessClosing(ctxTags, dht.proc)
 
- pm, err := providers.NewProviderManager(dht.ctx, h.ID(), cfg.Datastore, cfg.ProvidersOptions...)
- if err != nil {
- return nil, err
+ if cfg.ProviderStore != nil {
+ dht.providerStore = cfg.ProviderStore
+ } else {
+ dht.providerStore, err = providers.NewProviderManager(dht.ctx, h.ID(), dht.peerstore, cfg.Datastore, cfg.ProvidersOptions...)
+ if err != nil {
+ return nil, fmt.Errorf("initializing default provider manager (%v)", err)
+ }
  }
- dht.ProviderManager = pm
 
  dht.rtFreezeTimeout = rtFreezeTimeout
 
@@ -413,6 +418,11 @@ func makeRoutingTable(dht *IpfsDHT, cfg dhtcfg.Config, maxLastSuccessfulOutbound
  return rt, err
 }
 
+// ProviderStore returns the provider storage object for storing and retrieving provider records.
+func (dht *IpfsDHT) ProviderStore() providers.ProviderStore {
+ return dht.providerStore
+}
+
 // GetRoutingTableDiversityStats returns the diversity stats for the Routing Table.
 func (dht *IpfsDHT) GetRoutingTableDiversityStats() []peerdiversity.CplDiversityStats {
  return dht.routingTable.GetDiversityStats()

dht_options.go

@@ -36,6 +36,14 @@ const DefaultPrefix protocol.ID = "/ipfs"
 
 type Option = dhtcfg.Option
 
+// ProviderStore sets the provider storage manager.
+func ProviderStore(ps providers.ProviderStore) Option {
+ return func(c *dhtcfg.Config) error {
+ c.ProviderStore = ps
+ return nil
+ }
+}
+
 // RoutingTableLatencyTolerance sets the maximum acceptable latency for peers
 // in the routing table's cluster.
 func RoutingTableLatencyTolerance(latency time.Duration) Option {

handlers.go

@@ -8,7 +8,6 @@ import (
  "time"
 
  "github.com/libp2p/go-libp2p-core/peer"
- "github.com/libp2p/go-libp2p-core/peerstore"
  pstore "github.com/libp2p/go-libp2p-peerstore"
 
  "github.com/gogo/protobuf/proto"
@@ -318,13 +317,8 @@ func (dht *IpfsDHT) handleGetProviders(ctx context.Context, p peer.ID, pmes *pb.
  resp := pb.NewMessage(pmes.GetType(), pmes.GetKey(), pmes.GetClusterLevel())
 
  // setup providers
- providers := dht.ProviderManager.GetProviders(ctx, key)
-
- if len(providers) > 0 {
- // TODO: pstore.PeerInfos should move to core (=> peerstore.AddrInfos).
- infos := pstore.PeerInfos(dht.peerstore, providers)
- resp.ProviderPeers = pb.PeerInfosToPBPeers(dht.host.Network(), infos)
- }
+ providers := dht.providerStore.GetProviders(ctx, key)
+ resp.ProviderPeers = pb.PeerInfosToPBPeers(dht.host.Network(), providers)
 
  // Also send closer peers.
  closer := dht.betterPeersToQuery(pmes, p, dht.bucketSize)
@@ -362,11 +356,7 @@ func (dht *IpfsDHT) handleAddProvider(ctx context.Context, p peer.ID, pmes *pb.M
  continue
  }
 
- if pi.ID != dht.self { // don't add own addrs.
- // add the received addresses to our peerstore.
- dht.peerstore.AddAddrs(pi.ID, pi.Addrs, peerstore.ProviderAddrTTL)
- }
- dht.ProviderManager.AddProvider(ctx, key, p)
+ dht.providerStore.AddProvider(ctx, key, pstore.PeerInfo{ID: p})
  }
 
  return nil, nil

internal/config/config.go

@@ -44,7 +44,8 @@ type Config struct {
  MaxRecordAge time.Duration
  EnableProviders bool
  EnableValues bool
- ProvidersOptions []providers.Option
+ ProvidersOptions []providers.Option //XXX: get rid?
+ ProviderStore providers.ProviderStore
  QueryPeerFilter QueryFilterFunc
 
  RoutingTable struct {

providers/providers_manager.go

@@ -8,6 +8,7 @@ import (
  "time"
 
  "github.com/libp2p/go-libp2p-core/peer"
+ peerstore "github.com/libp2p/go-libp2p-peerstore"
 
  lru "github.com/hashicorp/golang-lru/simplelru"
  ds "github.com/ipfs/go-datastore"
@@ -30,12 +31,20 @@ var lruCacheSize = 256
 var batchBufferSize = 256
 var log = logging.Logger("providers")
 
+// ProviderStore represents a store that associates peers and their addresses to keys.
+type ProviderStore interface {
+ AddProvider(ctx context.Context, key []byte, prov peer.AddrInfo)
+ GetProviders(ctx context.Context, key []byte) []peer.AddrInfo
+}
+
 // ProviderManager adds and pulls providers out of the datastore,
 // caching them in between
 type ProviderManager struct {
+ self peer.ID
  // all non channel fields are meant to be accessed only within
  // the run method
  cache lru.LRUCache
+ pstore peerstore.Peerstore
  dstore *autobatch.Datastore
 
  newprovs chan *addProv
@@ -86,10 +95,12 @@ type getProv struct {
 }
 
 // NewProviderManager constructor
-func NewProviderManager(ctx context.Context, local peer.ID, dstore ds.Batching, opts ...Option) (*ProviderManager, error) {
+func NewProviderManager(ctx context.Context, local peer.ID, ps peerstore.Peerstore, dstore ds.Batching, opts ...Option) (*ProviderManager, error) {
  pm := new(ProviderManager)
+ pm.self = local
  pm.getprovs = make(chan *getProv)
  pm.newprovs = make(chan *addProv)
+ pm.pstore = ps
  pm.dstore = autobatch.NewAutoBatching(dstore, batchBufferSize)
  cache, err := lru.NewLRU(lruCacheSize, nil)
  if err != nil {
@@ -214,10 +225,13 @@ func (pm *ProviderManager) run(proc goprocess.Process) {
 }
 
 // AddProvider adds a provider
-func (pm *ProviderManager) AddProvider(ctx context.Context, k []byte, val peer.ID) {
+func (pm *ProviderManager) AddProvider(ctx context.Context, k []byte, provInfo peer.AddrInfo) {
+ if provInfo.ID != pm.self { // don't add own addrs.
+ pm.pstore.AddAddrs(provInfo.ID, provInfo.Addrs, peerstore.ProviderAddrTTL)
+ }
  prov := &addProv{
  key: k,
- val: val,
+ val: provInfo.ID,
  }
  select {
  case pm.newprovs <- prov:
@@ -255,7 +269,7 @@ func mkProvKey(k []byte) string {
 
 // GetProviders returns the set of providers for the given key.
 // This method _does not_ copy the set. Do not modify it.
-func (pm *ProviderManager) GetProviders(ctx context.Context, k []byte) []peer.ID {
+func (pm *ProviderManager) GetProviders(ctx context.Context, k []byte) []peer.AddrInfo {
  gp := &getProv{
  key: k,
  resp: make(chan []peer.ID, 1), // buffered to prevent sender from blocking
@@ -269,7 +283,7 @@ func (pm *ProviderManager) GetProviders(ctx context.Context, k []byte) []peer.ID
  case <-ctx.Done():
  return nil
  case peers := <-gp.resp:
- return peers
+ return peerstore.PeerInfos(pm.pstore, peers)
  }
 }
 

providers/providers_manager_test.go

@@ -21,6 +21,10 @@ import (
  // lds "github.com/ipfs/go-ds-leveldb"
 )
 
+func TestProviderManagerImplementsProviderStore(t *testing.T) {
+ var _ ProviderStore = (*ProviderManager)(nil)
+}
+
 func TestProviderManager(t *testing.T) {
  ctx, cancel := context.WithCancel(context.Background())
  defer cancel()

routing.go

@@ -380,7 +380,7 @@ func (dht *IpfsDHT) Provide(ctx context.Context, key cid.Cid, brdcst bool) (err
  logger.Debugw("providing", "cid", key, "mh", internal.LoggableProviderRecordBytes(keyMH))
 
  // add self locally
- dht.ProviderManager.AddProvider(ctx, keyMH, dht.self)
+ dht.providerStore.AddProvider(ctx, keyMH, peer.AddrInfo{ID: dht.self})
  if !brdcst {
  return nil
  }
@@ -492,13 +492,13 @@ func (dht *IpfsDHT) findProvidersAsyncRoutine(ctx context.Context, key multihash
  ps = peer.NewLimitedSet(count)
  }
 
- provs := dht.ProviderManager.GetProviders(ctx, key)
+ provs := dht.providerStore.GetProviders(ctx, key)
  for _, p := range provs {
  // NOTE: Assuming that this list of peers is unique
- if ps.TryAdd(p) {
- pi := dht.peerstore.PeerInfo(p)
+ if ps.TryAdd(p.ID) {
+ pi := dht.peerstore.PeerInfo(p.ID)
  select {
- case peerOut <- pi:
+ case peerOut <- peer.AddrInfo{ID: p.ID, Addrs: append(pi.Addrs, p.Addrs...)}: // XXX: dedup addresses?
  case <-ctx.Done():
  return
  }