Add Discovery

comm.go

@@ -19,7 +19,7 @@ import (
 // get the initial RPC containing all of our subscriptions to send to new peers
 func (p *PubSub) getHelloPacket() *RPC {
 	var rpc RPC
-	for t := range p.myTopics {
+	for t := range p.mySubs {
 		as := &pb.RPC_SubOpts{
 			Topicid:   proto.String(t),
 			Subscribe: proto.Bool(true),

discovery.go

@@ -0,0 +1,336 @@
+package pubsub
+
+import (
+	"context"
+	"math/rand"
+	"time"
+
+	"github.com/libp2p/go-libp2p-core/discovery"
+	"github.com/libp2p/go-libp2p-core/host"
+	"github.com/libp2p/go-libp2p-core/peer"
+	discimpl "github.com/libp2p/go-libp2p-discovery"
+)
+
+var (
+	// poll interval
+
+	// DiscoveryPollInitialDelay is how long the discovery system waits after it first starts before polling
+	DiscoveryPollInitialDelay = 0 * time.Millisecond
+	// DiscoveryPollInterval is approximately how long the discovery system waits in between checks for whether the
+	// more peers are needed for any topic
+	DiscoveryPollInterval = 1 * time.Second
+)
+
+type DiscoverOpt func(*discoverOptions) error
+
+type discoverOptions struct {
+	connFactory BackoffConnectorFactory
+	opts        []discovery.Option
+}
+
+func defaultDiscoverOptions() *discoverOptions {
+	rng := rand.New(rand.NewSource(rand.Int63()))
+	minBackoff, maxBackoff := time.Second*10, time.Hour
+	cacheSize := 100
+	dialTimeout := time.Minute * 2
+	discoverOpts := &discoverOptions{
+		connFactory: func(host host.Host) (*discimpl.BackoffConnector, error) {
+			backoff := discimpl.NewExponentialBackoff(minBackoff, maxBackoff, discimpl.FullJitter, time.Second, 5.0, 0, rng)
+			return discimpl.NewBackoffConnector(host, cacheSize, dialTimeout, backoff)
+		},
+	}
+
+	return discoverOpts
+}
+
+// discover represents the discovery pipeline.
+// The discovery pipeline handles advertising and discovery of peers
+type discover struct {
+	p *PubSub
+
+	// discovery assists in discovering and advertising peers for a topic
+	discovery discovery.Discovery
+
+	// advertising tracks which topics are being advertised
+	advertising map[string]context.CancelFunc
+
+	// discoverQ handles continuing peer discovery
+	discoverQ chan *discoverReq
+
+	// ongoing tracks ongoing discovery requests
+	ongoing map[string]struct{}
+
+	// done handles completion of a discovery request
+	done chan string
+
+	// connector handles connecting to new peers found via discovery
+	connector *discimpl.BackoffConnector
+
+	// options are the set of options to be used to complete struct construction in Start
+	options *discoverOptions
+}
+
+// MinTopicSize returns a function that checks if a router is ready for publishing based on the topic size.
+// The router ultimately decides the whether it is ready or not, the given size is just a suggestion.
+func MinTopicSize(size int) RouterReady {
+	return func(rt PubSubRouter, topic string) (bool, error) {
+		return rt.EnoughPeers(topic, size), nil
+	}
+}
+
+// Start attaches the discovery pipeline to a pubsub instance, initializes discovery and starts event loop
+func (d *discover) Start(p *PubSub, opts ...DiscoverOpt) error {
+	if d.discovery == nil || p == nil {
+		return nil
+	}
+
+	d.p = p
+	d.advertising = make(map[string]context.CancelFunc)
+	d.discoverQ = make(chan *discoverReq, 32)
+	d.ongoing = make(map[string]struct{})
+	d.done = make(chan string)
+
+	conn, err := d.options.connFactory(p.host)
+	if err != nil {
+		return err
+	}
+	d.connector = conn
+
+	go d.discoverLoop()
+	go d.pollTimer()
+
+	return nil
+}
+
+func (d *discover) pollTimer() {
+	select {
+	case <-time.After(DiscoveryPollInitialDelay):
+	case <-d.p.ctx.Done():
+		return
+	}
+
+	select {
+	case d.p.eval <- d.requestDiscovery:
+	case <-d.p.ctx.Done():
+		return
+	}
+
+	ticker := time.NewTicker(DiscoveryPollInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			select {
+			case d.p.eval <- d.requestDiscovery:
+			case <-d.p.ctx.Done():
+				return
+			}
+		case <-d.p.ctx.Done():
+			return
+		}
+	}
+}
+
+func (d *discover) requestDiscovery() {
+	for t := range d.p.myTopics {
+		if !d.p.rt.EnoughPeers(t, 0) {
+			d.discoverQ <- &discoverReq{topic: t, done: make(chan struct{}, 1)}
+		}
+	}
+}
+
+func (d *discover) discoverLoop() {
+	for {
+		select {
+		case discover := <-d.discoverQ:
+			topic := discover.topic
+
+			if _, ok := d.ongoing[topic]; ok {
+				discover.done <- struct{}{}
+				continue
+			}
+
+			d.ongoing[topic] = struct{}{}
+
+			go func() {
+				d.handleDiscovery(d.p.ctx, topic, discover.opts)
+				select {
+				case d.done <- topic:
+				case <-d.p.ctx.Done():
+				}
+				discover.done <- struct{}{}
+			}()
+		case topic := <-d.done:
+			delete(d.ongoing, topic)
+		case <-d.p.ctx.Done():
+			return
+		}
+	}
+}
+
+// Advertise advertises this node's interest in a topic to a discovery service. Advertise is not thread-safe.
+func (d *discover) Advertise(topic string) {
+	if d.discovery == nil {
+		return
+	}
+
+	advertisingCtx, cancel := context.WithCancel(d.p.ctx)
+
+	if _, ok := d.advertising[topic]; ok {
+		cancel()
+		return
+	}
+	d.advertising[topic] = cancel
+
+	go func() {
+		next, err := d.discovery.Advertise(advertisingCtx, topic)
+		if err != nil {
+			log.Warningf("bootstrap: error providing rendezvous for %s: %s", topic, err.Error())
+		}
+
+		t := time.NewTimer(next)
+		for {
+			select {
+			case <-t.C:
+				next, err = d.discovery.Advertise(advertisingCtx, topic)
+				if err != nil {
+					log.Warningf("bootstrap: error providing rendezvous for %s: %s", topic, err.Error())
+				}
+				t.Reset(next)
+			case <-advertisingCtx.Done():
+				t.Stop()
+				return
+			}
+		}
+	}()
+}
+
+// StopAdvertise stops advertising this node's interest in a topic. StopAdvertise is not thread-safe.
+func (d *discover) StopAdvertise(topic string) {
+	if d.discovery == nil {
+		return
+	}
+
+	if advertiseCancel, ok := d.advertising[topic]; ok {
+		advertiseCancel()
+		delete(d.advertising, topic)
+	}
+}
+
+// Discover searches for additional peers interested in a given topic
+func (d *discover) Discover(topic string, opts ...discovery.Option) {
+	if d.discovery == nil {
+		return
+	}
+
+	d.discoverQ <- &discoverReq{topic, opts, make(chan struct{}, 1)}
+}
+
+// Bootstrap attempts to bootstrap to a given topic. Returns true if bootstrapped successfully, false otherwise.
+func (d *discover) Bootstrap(ctx context.Context, topic string, ready RouterReady, opts ...discovery.Option) bool {
+	if d.discovery == nil {
+		return true
+	}
+
+	t := time.NewTimer(time.Hour)
+	if !t.Stop() {
+		<-t.C
+	}
+
+	for {
+		// Check if ready for publishing
+		bootstrapped := make(chan bool, 1)
+		select {
+		case d.p.eval <- func() {
+			done, _ := ready(d.p.rt, topic)
+			bootstrapped <- done
+		}:
+			if <-bootstrapped {
+				return true
+			}
+		case <-d.p.ctx.Done():
+			return false
+		case <-ctx.Done():
+			return false
+		}
+
+		// If not ready discover more peers
+		disc := &discoverReq{topic, opts, make(chan struct{}, 1)}
+		select {
+		case d.discoverQ <- disc:
+		case <-d.p.ctx.Done():
+			return false
+		case <-ctx.Done():
+			return false
+		}
+
+		select {
+		case <-disc.done:
+		case <-d.p.ctx.Done():
+			return false
+		case <-ctx.Done():
+			return false
+		}
+
+		t.Reset(time.Millisecond * 100)
+		select {
+		case <-t.C:
+		case <-d.p.ctx.Done():
+			return false
+		case <-ctx.Done():
+			return false
+		}
+	}
+}
+
+func (d *discover) handleDiscovery(ctx context.Context, topic string, opts []discovery.Option) {
+	discoverCtx, cancel := context.WithTimeout(ctx, time.Second*10)
+	defer cancel()
+
+	peerCh, err := d.discovery.FindPeers(discoverCtx, topic, opts...)
+	if err != nil {
+		log.Debugf("error finding peers for topic %s: %v", topic, err)
+		return
+	}
+
+	d.connector.Connect(ctx, peerCh)
+}
+
+type discoverReq struct {
+	topic string
+	opts  []discovery.Option
+	done  chan struct{}
+}
+
+type pubSubDiscovery struct {
+	discovery.Discovery
+	opts []discovery.Option
+}
+
+func (d *pubSubDiscovery) Advertise(ctx context.Context, ns string, opts ...discovery.Option) (time.Duration, error) {
+	return d.Discovery.Advertise(ctx, "floodsub:"+ns, append(opts, d.opts...)...)
+}
+
+func (d *pubSubDiscovery) FindPeers(ctx context.Context, ns string, opts ...discovery.Option) (<-chan peer.AddrInfo, error) {
+	return d.Discovery.FindPeers(ctx, "floodsub:"+ns, append(opts, d.opts...)...)
+}
+
+// WithDiscoveryOpts passes libp2p Discovery options into the PubSub discovery subsystem
+func WithDiscoveryOpts(opts ...discovery.Option) DiscoverOpt {
+	return func(d *discoverOptions) error {
+		d.opts = opts
+		return nil
+	}
+}
+
+// BackoffConnectorFactory creates a BackoffConnector that is attached to a given host
+type BackoffConnectorFactory func(host host.Host) (*discimpl.BackoffConnector, error)
+
+// WithDiscoverConnector adds a custom connector that deals with how the discovery subsystem connects to peers
+func WithDiscoverConnector(connFactory BackoffConnectorFactory) DiscoverOpt {
+	return func(d *discoverOptions) error {
+		d.connFactory = connFactory
+		return nil
+	}
+}