ceremonyclient/go-libp2p/p2p/discovery/backoff/backoffcache.go

package backoff

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p/core/discovery"
	"github.com/libp2p/go-libp2p/core/peer"

	ma "github.com/multiformats/go-multiaddr"
)

// BackoffDiscovery is an implementation of discovery that caches peer data and attenuates repeated queries
type BackoffDiscovery struct {
	disc         discovery.Discovery
	stratFactory BackoffFactory
	peerCache    map[string]*backoffCache
	peerCacheMux sync.RWMutex

	parallelBufSz int
	returnedBufSz int

	clock clock
}

type BackoffDiscoveryOption func(*BackoffDiscovery) error

func NewBackoffDiscovery(disc discovery.Discovery, stratFactory BackoffFactory, opts ...BackoffDiscoveryOption) (discovery.Discovery, error) {
	b := &BackoffDiscovery{
		disc:         disc,
		stratFactory: stratFactory,
		peerCache:    make(map[string]*backoffCache),

		parallelBufSz: 32,
		returnedBufSz: 32,

		clock: realClock{},
	}

	for _, opt := range opts {
		if err := opt(b); err != nil {
			return nil, err
		}
	}

	return b, nil
}

// WithBackoffDiscoverySimultaneousQueryBufferSize sets the buffer size for the channels between the main FindPeers query
// for a given namespace and all simultaneous FindPeers queries for the namespace
func WithBackoffDiscoverySimultaneousQueryBufferSize(size int) BackoffDiscoveryOption {
	return func(b *BackoffDiscovery) error {
		if size < 0 {
			return fmt.Errorf("cannot set size to be smaller than 0")
		}
		b.parallelBufSz = size
		return nil
	}
}

// WithBackoffDiscoveryReturnedChannelSize sets the size of the buffer to be used during a FindPeer query.
// Note: This does not apply if the query occurs during the backoff time
func WithBackoffDiscoveryReturnedChannelSize(size int) BackoffDiscoveryOption {
	return func(b *BackoffDiscovery) error {
		if size < 0 {
			return fmt.Errorf("cannot set size to be smaller than 0")
		}
		b.returnedBufSz = size
		return nil
	}
}

type clock interface {
	Now() time.Time
}

type realClock struct{}

func (c realClock) Now() time.Time {
	return time.Now()
}

// withClock lets you override the default time.Now() call. Useful for tests.
func withClock(c clock) BackoffDiscoveryOption {
	return func(b *BackoffDiscovery) error {
		b.clock = c
		return nil
	}
}

type backoffCache struct {
	// strat is assigned on creation and not written to
	strat BackoffStrategy

	mux          sync.Mutex // guards writes to all following fields
	nextDiscover time.Time
	prevPeers    map[peer.ID]peer.AddrInfo
	peers        map[peer.ID]peer.AddrInfo
	sendingChs   map[chan peer.AddrInfo]int
	ongoing      bool

	clock clock
}

func (d *BackoffDiscovery) Advertise(ctx context.Context, ns string, opts ...discovery.Option) (time.Duration, error) {
	return d.disc.Advertise(ctx, ns, opts...)
}

func (d *BackoffDiscovery) FindPeers(ctx context.Context, ns string, opts ...discovery.Option) (<-chan peer.AddrInfo, error) {
	// Get options
	var options discovery.Options
	err := options.Apply(opts...)
	if err != nil {
		return nil, err
	}

	// Get cached peers
	d.peerCacheMux.RLock()
	c, ok := d.peerCache[ns]
	d.peerCacheMux.RUnlock()

	/*
		Overall plan:
		If it's time to look for peers, look for peers, then return them
		If it's not time then return cache
		If it's time to look for peers, but we have already started looking. Get up to speed with ongoing request
	*/

	// Setup cache if we don't have one yet
	if !ok {
		pc := &backoffCache{
			nextDiscover: time.Time{},
			prevPeers:    make(map[peer.ID]peer.AddrInfo),
			peers:        make(map[peer.ID]peer.AddrInfo),
			sendingChs:   make(map[chan peer.AddrInfo]int),
			strat:        d.stratFactory(),
			clock:        d.clock,
		}

		d.peerCacheMux.Lock()
		c, ok = d.peerCache[ns]

		if !ok {
			d.peerCache[ns] = pc
			c = pc
		}

		d.peerCacheMux.Unlock()
	}

	c.mux.Lock()
	defer c.mux.Unlock()

	timeExpired := d.clock.Now().After(c.nextDiscover)

	// If it's not yet time to search again and no searches are in progress then return cached peers
	if !(timeExpired || c.ongoing) {
		chLen := options.Limit

		if chLen == 0 {
			chLen = len(c.prevPeers)
		} else if chLen > len(c.prevPeers) {
			chLen = len(c.prevPeers)
		}
		pch := make(chan peer.AddrInfo, chLen)
		for _, ai := range c.prevPeers {
			select {
			case pch <- ai:
			default:
				// skip if we have asked for a lower limit than the number of peers known
			}
		}
		close(pch)
		return pch, nil
	}

	// If a request is not already in progress setup a dispatcher channel for dispatching incoming peers
	if !c.ongoing {
		pch, err := d.disc.FindPeers(ctx, ns, opts...)
		if err != nil {
			return nil, err
		}

		c.ongoing = true
		go findPeerDispatcher(ctx, c, pch)
	}

	// Setup receiver channel for receiving peers from ongoing requests
	evtCh := make(chan peer.AddrInfo, d.parallelBufSz)
	pch := make(chan peer.AddrInfo, d.returnedBufSz)
	rcvPeers := make([]peer.AddrInfo, 0, 32)
	for _, ai := range c.peers {
		rcvPeers = append(rcvPeers, ai)
	}
	c.sendingChs[evtCh] = options.Limit

	go findPeerReceiver(ctx, pch, evtCh, rcvPeers)

	return pch, nil
}

func findPeerDispatcher(ctx context.Context, c *backoffCache, pch <-chan peer.AddrInfo) {
	defer func() {
		c.mux.Lock()

		// If the peer addresses have changed reset the backoff
		if checkUpdates(c.prevPeers, c.peers) {
			c.strat.Reset()
			c.prevPeers = c.peers
		}
		c.nextDiscover = c.clock.Now().Add(c.strat.Delay())

		c.ongoing = false
		c.peers = make(map[peer.ID]peer.AddrInfo)

		for ch := range c.sendingChs {
			close(ch)
		}
		c.sendingChs = make(map[chan peer.AddrInfo]int)
		c.mux.Unlock()
	}()

	for {
		select {
		case ai, ok := <-pch:
			if !ok {
				return
			}
			c.mux.Lock()

			// If we receive the same peer multiple times return the address union
			var sendAi peer.AddrInfo
			if prevAi, ok := c.peers[ai.ID]; ok {
				if combinedAi := mergeAddrInfos(prevAi, ai); combinedAi != nil {
					sendAi = *combinedAi
				} else {
					c.mux.Unlock()
					continue
				}
			} else {
				sendAi = ai
			}

			c.peers[ai.ID] = sendAi

			for ch, rem := range c.sendingChs {
				if rem > 0 {
					ch <- sendAi
					c.sendingChs[ch] = rem - 1
				}
			}

			c.mux.Unlock()
		case <-ctx.Done():
			return
		}
	}
}

func findPeerReceiver(ctx context.Context, pch, evtCh chan peer.AddrInfo, rcvPeers []peer.AddrInfo) {
	defer close(pch)

	for {
		select {
		case ai, ok := <-evtCh:
			if ok {
				rcvPeers = append(rcvPeers, ai)

				sentAll := true
			sendPeers:
				for i, p := range rcvPeers {
					select {
					case pch <- p:
					default:
						rcvPeers = rcvPeers[i:]
						sentAll = false
						break sendPeers
					}
				}
				if sentAll {
					rcvPeers = []peer.AddrInfo{}
				}
			} else {
				for _, p := range rcvPeers {
					select {
					case pch <- p:
					case <-ctx.Done():
						return
					}
				}
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

func mergeAddrInfos(prevAi, newAi peer.AddrInfo) *peer.AddrInfo {
	seen := make(map[string]struct{}, len(prevAi.Addrs))
	combinedAddrs := make([]ma.Multiaddr, 0, len(prevAi.Addrs))
	addAddrs := func(addrs []ma.Multiaddr) {
		for _, addr := range addrs {
			if _, ok := seen[addr.String()]; ok {
				continue
			}
			seen[addr.String()] = struct{}{}
			combinedAddrs = append(combinedAddrs, addr)
		}
	}
	addAddrs(prevAi.Addrs)
	addAddrs(newAi.Addrs)

	if len(combinedAddrs) > len(prevAi.Addrs) {
		combinedAi := &peer.AddrInfo{ID: prevAi.ID, Addrs: combinedAddrs}
		return combinedAi
	}
	return nil
}

func checkUpdates(orig, update map[peer.ID]peer.AddrInfo) bool {
	if len(orig) != len(update) {
		return true
	}
	for p, ai := range update {
		if prevAi, ok := orig[p]; ok {
			if combinedAi := mergeAddrInfos(prevAi, ai); combinedAi != nil {
				return true
			}
		} else {
			return true
		}
	}
	return false
}