feat(share/discovery): persist discovered peers

share/availability/discovery/discovery.go

@@ -2,9 +2,12 @@ package discovery
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"sync"
 	"time"
 
+	"github.com/ipfs/go-datastore"
 	logging "github.com/ipfs/go-log/v2"
 	"github.com/libp2p/go-libp2p/core/discovery"
 	"github.com/libp2p/go-libp2p/core/event"
@@ -82,7 +85,9 @@ func NewDiscovery(
 	}
 }
 
-func (d *Discovery) Start(context.Context) error {
+func (d *Discovery) Start(ctx context.Context) error {
+	d.loadPeers(ctx)
+
 	ctx, cancel := context.WithCancel(context.Background())
 	d.cancel = cancel
 
@@ -102,7 +107,8 @@ func (d *Discovery) Start(context.Context) error {
 	return nil
 }
 
-func (d *Discovery) Stop(context.Context) error {
+func (d *Discovery) Stop(ctx context.Context) error {
+	d.dumpPeers(ctx)
 	d.cancel()
 	return nil
 }
@@ -309,6 +315,45 @@ func (d *Discovery) discover(ctx context.Context) bool {
 	}
 }
 
+// loadPeers loads and connects to persisted peers from datastore if registered
+func (d *Discovery) loadPeers(ctx context.Context) {
+	if d.params.datastore == nil {
+		return
+	}
+
+	peerInfos, err := loadPeers(ctx, d.params.datastore)
+	if err != nil {
+		log.Warnw("failed to load persisted peers", "error", err)
+		return
+	}
+
+	// TODO(@Wondertan): To wait or not to wait?
+	var wg sync.WaitGroup
+	wg.Add(len(peerInfos))
+	defer wg.Wait()
+	for _, pi := range peerInfos {
+		go func(pi peer.AddrInfo) {
+			d.handleDiscoveredPeer(ctx, pi)
+			wg.Done()
+		}(pi)
+	}
+}
+
+// dumpPeers dumps discovered peers to datastore if registered
+func (d *Discovery) dumpPeers(ctx context.Context) {
+	if d.params.datastore == nil {
+		return
+	}
+
+	// TODO(@Wondertan): We should be using persistent peerstore instead
+	//  and store here only peer ID pointers
+	peerInfos := d.set.peerInfos(d.host.Peerstore())
+	err := dumpPeers(ctx, d.params.datastore, peerInfos)
+	if err != nil {
+		log.Warnw("failed to persist peers", "error", err)
+	}
+}
+
 // handleDiscoveredPeer adds peer to the internal if can connect or is connected.
 // Report whether it succeeded.
 func (d *Discovery) handleDiscoveredPeer(ctx context.Context, peer peer.AddrInfo) bool {
@@ -372,3 +417,24 @@ func drainChannel(c <-chan time.Time) {
 		}
 	}
 }
+
+var dsKey = datastore.RawKey("/discovery/dump")
+
+func dumpPeers(ctx context.Context, ds datastore.Batching, infos []peer.AddrInfo) error {
+	dump, err := json.Marshal(infos)
+	if err != nil {
+		return err
+	}
+
+	return ds.Put(ctx, dsKey, dump)
+}
+
+func loadPeers(ctx context.Context, ds datastore.Batching) ([]peer.AddrInfo, error) {
+	dump, err := ds.Get(ctx, dsKey)
+	if err != nil {
+		return nil, err
+	}
+
+	var infos []peer.AddrInfo
+	return infos, json.Unmarshal(dump, &infos)
+}

share/availability/discovery/discovery_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ipfs/go-datastore"
 	dht "github.com/libp2p/go-libp2p-kad-dht"
 	"github.com/libp2p/go-libp2p/core/discovery"
 	"github.com/libp2p/go-libp2p/core/host"
@@ -72,6 +73,58 @@ func TestDiscovery(t *testing.T) {
 	assert.EqualValues(t, 0, peerA.set.Size())
 }
 
+func TestDiscoveryPersistence(t *testing.T) {
+	const nodes = 2
+
+	discoveryRetryTimeout = time.Millisecond * 100 // defined in discovery.go
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+	t.Cleanup(cancel)
+
+	tn := newTestnet(ctx, t)
+
+	ds := datastore.NewMapDatastore()
+	peerA := tn.discovery(
+		WithPeersLimit(nodes),
+		WithAdvertiseInterval(-1),
+		WithPersistence(ds),
+	)
+
+	evt := make(chan struct{})
+	peerA.WithOnPeersUpdate(func(peerID peer.ID, isAdded bool) {
+		evt <- struct{}{}
+	})
+
+	discs := make([]*Discovery, nodes)
+	for range discs {
+		tn.discovery(
+			WithPeersLimit(nodes),
+			WithAdvertiseInterval(time.Millisecond*100),
+		)
+
+		select {
+		case <-evt:
+		case <-ctx.Done():
+			t.Fatal("did not discover peer in time")
+		}
+	}
+
+	err := peerA.Stop(ctx)
+	require.NoError(t, err)
+
+	// make new borker discovery that won't discover anything
+	// over datastore with persisted peers
+	peerANew := tn.borkenDiscovery(
+		WithPeersLimit(nodes),
+		WithAdvertiseInterval(-1),
+		WithPersistence(ds),
+	)
+
+	peers, err := peerANew.Peers(ctx)
+	require.NoError(t, err)
+	assert.Len(t, peers, nodes)
+}
+
 type testnet struct {
 	ctx context.Context
 	T   *testing.T
@@ -109,6 +162,18 @@ func (t *testnet) discovery(opts ...Option) *Discovery {
 	return disc
 }
 
+func (t *testnet) borkenDiscovery(opts ...Option) *Discovery {
+	hst, _ := t.peer()
+	disc := NewDiscovery(hst, &dummyDiscovery{}, opts...)
+	err := disc.Start(t.ctx)
+	require.NoError(t.T, err)
+	t.T.Cleanup(func() {
+		err := disc.Stop(t.ctx)
+		require.NoError(t.T, err)
+	})
+	return disc
+}
+
 func (t *testnet) peer() (host.Host, discovery.Discovery) {
 	swarm := swarmt.GenSwarm(t.T, swarmt.OptDisableTCP)
 	hst, err := basic.NewHost(swarm, &basic.HostOpts{})
@@ -131,3 +196,18 @@ func (t *testnet) peer() (host.Host, discovery.Discovery) {
 
 	return hst, routing.NewRoutingDiscovery(dht)
 }
+
+type dummyDiscovery struct{}
+
+func (d *dummyDiscovery) Advertise(context.Context, string, ...discovery.Option) (time.Duration, error) {
+	return time.Hour, nil
+}
+
+func (d *dummyDiscovery) FindPeers(context.Context, string, ...discovery.Option) (<-chan peer.AddrInfo, error) {
+	retCh := make(chan peer.AddrInfo)
+	go func() {
+		time.Sleep(time.Second)
+		close(retCh)
+	}()
+	return retCh, nil
+}

share/availability/discovery/options.go

@@ -3,6 +3,8 @@ package discovery
 import (
 	"fmt"
 	"time"
+
+	"github.com/ipfs/go-datastore"
 )
 
 // Parameters is the set of Parameters that must be configured for the Discovery module
@@ -14,6 +16,8 @@ type Parameters struct {
 	// Set -1 to disable.
 	// NOTE: only full and bridge can advertise themselves.
 	AdvertiseInterval time.Duration
+
+	datastore datastore.Batching
 }
 
 // Option is a function that configures Discovery Parameters
@@ -57,3 +61,9 @@ func WithAdvertiseInterval(advInterval time.Duration) Option {
 		p.AdvertiseInterval = advInterval
 	}
 }
+
+func WithPersistence(ds datastore.Batching) Option {
+	return func(p *Parameters) {
+		p.datastore = ds
+	}
+}

share/availability/discovery/set.go

@@ -5,10 +5,10 @@ import (
 	"sync"
 
 	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/core/peerstore"
 )
 
 // limitedSet is a thread safe set of peers with given limit.
-// Inspired by libp2p peer.Set but extended with Remove method.
 type limitedSet struct {
 	lk sync.RWMutex
 	ps map[peer.ID]struct{}
@@ -91,3 +91,14 @@ func (ps *limitedSet) Peers(ctx context.Context) ([]peer.ID, error) {
 		return []peer.ID{p}, nil
 	}
 }
+
+func (ps *limitedSet) peerInfos(pstore peerstore.Peerstore) []peer.AddrInfo {
+	infos := make([]peer.AddrInfo, 0, len(ps.ps))
+	for p := range ps.ps {
+		infos = append(infos, peer.AddrInfo{
+			ID:    p,
+			Addrs: pstore.Addrs(p),
+		})
+	}
+	return infos
+}