fix(tinder): Update multi-driver, mocked driver & some tests

go/internal/tinder/driver_mocked.go

@@ -1,184 +1,148 @@
+// from https://github.com/libp2p/go-libp2p-discovery/blob/master/mocks_test.go
+
 package tinder
 
 import (
 	"context"
 	"sync"
 	"time"
 
-	libp2p_discovery "github.com/libp2p/go-libp2p-core/discovery"
-	libp2p_host "github.com/libp2p/go-libp2p-core/host"
-	libp2p_peer "github.com/libp2p/go-libp2p-core/peer"
+	p2p_discovery "github.com/libp2p/go-libp2p-core/discovery"
+	p2p_host "github.com/libp2p/go-libp2p-core/host"
+	p2p_peer "github.com/libp2p/go-libp2p-core/peer"
 )
 
-type MockedDriverFactory struct {
-	mockStore     *mockStore
-	mockedDrivers []Driver
+type MockDriverServer struct {
+	mx sync.RWMutex
+	db map[string]map[p2p_peer.ID]*discoveryRegistration
 }
 
-func NewMockedDriverFactory() *MockedDriverFactory {
-	return &MockedDriverFactory{
-		mockStore:     newMockStore(),
-		mockedDrivers: make([]Driver, 0),
-	}
+type discoveryRegistration struct {
+	info       p2p_peer.AddrInfo
+	expiration time.Time
 }
 
-func (mf *MockedDriverFactory) NewMockedDriver(h libp2p_host.Host) (d Driver) {
-	d = &mockedDriver{
-		Host:      h,
-		mockStore: mf.mockStore,
+func NewMockedDriverServer() *MockDriverServer {
+	return &MockDriverServer{
+		db: make(map[string]map[p2p_peer.ID]*discoveryRegistration),
 	}
-
-	mf.mockedDrivers = append(mf.mockedDrivers, d)
-	return d
 }
 
-type mockedDriver struct {
-	Host      libp2p_host.Host
-	mockStore *mockStore
-}
+func (s *MockDriverServer) Advertise(ns string, info p2p_peer.AddrInfo, ttl time.Duration) (time.Duration, error) {
+	s.mx.Lock()
+	defer s.mx.Unlock()
 
-func (md *mockedDriver) Advertise(ctx context.Context, ns string, opts ...libp2p_discovery.Option) (ttl time.Duration, err error) {
-	// Get options
-	var options libp2p_discovery.Options
-	err = options.Apply(opts...)
-	if err != nil {
-		return
+	peers, ok := s.db[ns]
+	if !ok {
+		peers = make(map[p2p_peer.ID]*discoveryRegistration)
+		s.db[ns] = peers
 	}
-
-	// get peer info
-	pi := md.Host.Peerstore().PeerInfo(md.Host.ID())
-	ttl = options.Ttl
-	md.mockStore.Put(ns, ttl, pi)
-	return
+	peers[info.ID] = &discoveryRegistration{info, time.Now().Add(ttl)}
+	return ttl, nil
 }
 
-func (md *mockedDriver) FindPeers(ctx context.Context, ns string, opts ...libp2p_discovery.Option) (<-chan libp2p_peer.AddrInfo, error) {
-	// Get options
-	var options libp2p_discovery.Options
-	err := options.Apply(opts...)
-	if err != nil {
-		return nil, err
+func (s *MockDriverServer) FindPeers(ns string, limit int) (<-chan p2p_peer.AddrInfo, error) {
+	s.mx.Lock()
+	defer s.mx.Unlock()
+
+	peers, ok := s.db[ns]
+	if !ok || len(peers) == 0 {
+		emptyCh := make(chan p2p_peer.AddrInfo)
+		close(emptyCh)
+		return emptyCh, nil
 	}
 
-	cpeers := make(chan libp2p_peer.AddrInfo)
-	go func() {
-		defer close(cpeers)
+	count := len(peers)
+	if limit != 0 && count > limit {
+		count = limit
+	}
 
-		pis := md.mockStore.Get(ns)
-		for _, pi := range pis {
-			cpeers <- pi
+	iterTime := time.Now()
+	ch := make(chan p2p_peer.AddrInfo, count)
+	numSent := 0
+	for p, reg := range peers {
+		if numSent == count {
+			break
+		}
+		if iterTime.After(reg.expiration) {
+			delete(peers, p)
+			continue
 		}
-	}()
 
-	return cpeers, nil
-}
+		numSent++
+		ch <- reg.info
+	}
 
-func (md *mockedDriver) Unregister(ctx context.Context, ns string) (err error) {
-	md.mockStore.Remove(ns, md.Host.ID())
-	return nil
-}
+	if len(peers) == 0 {
+		delete(s.db, ns)
+	}
 
-type mockStore struct {
-	mockRecords map[string]mockRecords
-	// dbRecords map[libp2p_peer.ID]*mockRecord
-	muStore sync.RWMutex
+	close(ch)
+
+	return ch, nil
 }
 
-func newMockStore() *mockStore {
-	return &mockStore{
-		mockRecords: make(map[string]mockRecords),
+func (s *MockDriverServer) Unregister(ns string, pid p2p_peer.ID) {
+	s.mx.Lock()
+	if peers, ok := s.db[ns]; ok {
+		if _, ok = peers[pid]; ok {
+			delete(peers, pid)
+		}
+
+		if len(peers) == 0 {
+			delete(s.db, ns)
+		}
 	}
+	s.mx.Unlock()
 }
 
-type mockRecords []*mockRecord
+func (s *MockDriverServer) HasPeerRecord(ns string, pid p2p_peer.ID) bool {
+	s.mx.RLock()
+	defer s.mx.RUnlock()
 
-func (rs mockRecords) AddrList() []libp2p_peer.AddrInfo {
-	pis := make([]libp2p_peer.AddrInfo, len(rs))
-	for i, r := range rs {
-		pis[i] = r.pi
+	if peers, ok := s.db[ns]; ok {
+		_, ok := peers[pid]
+		return ok
 	}
-	return pis
+	return false
 }
 
-type mockRecord struct {
-	pi      libp2p_peer.AddrInfo
-	ttl     time.Duration
-	timer   *time.Timer
-	deleted bool
+func (s *MockDriverServer) Reset() {
+	s.mx.Lock()
+	s.db = make(map[string]map[p2p_peer.ID]*discoveryRegistration)
+	s.mx.Unlock()
 }
 
-func newRecord(pi libp2p_peer.AddrInfo) *mockRecord {
-	return &mockRecord{pi, 0, time.NewTimer(0), false}
+type mockDriverClient struct {
+	host   p2p_host.Host
+	server *MockDriverServer
 }
 
-func (r *mockRecord) SetTimer(ttl time.Duration, f func()) {
-	if !r.timer.Stop() {
-		<-r.timer.C
-	}
-
-	r.timer = time.AfterFunc(ttl, f)
-	r.ttl = ttl
+func NewMockedDriverClient(host p2p_host.Host, server *MockDriverServer) Driver {
+	return &mockDriverClient{host, server}
 }
 
-func (s *mockStore) Get(key string) []libp2p_peer.AddrInfo {
-	s.muStore.RLock()
-	defer s.muStore.RUnlock()
-
-	recs, ok := s.mockRecords[key]
-	if !ok {
-		return []libp2p_peer.AddrInfo{}
+func (d *mockDriverClient) Advertise(ctx context.Context, ns string, opts ...p2p_discovery.Option) (time.Duration, error) {
+	var options p2p_discovery.Options
+	err := options.Apply(opts...)
+	if err != nil {
+		return 0, err
 	}
 
-	return recs.AddrList()
-
+	return d.server.Advertise(ns, *p2p_host.InfoFromHost(d.host), options.Ttl)
 }
 
-func (s *mockStore) Remove(key string, pid libp2p_peer.ID) {
-	s.muStore.Lock()
-	defer s.muStore.Unlock()
-
-	if recs, ok := s.mockRecords[key]; ok {
-		for i, rec := range recs {
-			if rec.pi.ID == pid {
-				rec.timer.Stop()
-				recs[len(recs)-1], recs[i] = recs[i], recs[len(recs)-1]
-				s.mockRecords[key] = recs[:len(recs)-1]
-				return
-			}
-		}
-
-		if len(s.mockRecords[key]) == 0 {
-			delete(s.mockRecords, key)
-		}
+func (d *mockDriverClient) FindPeers(ctx context.Context, ns string, opts ...p2p_discovery.Option) (<-chan p2p_peer.AddrInfo, error) {
+	var options p2p_discovery.Options
+	err := options.Apply(opts...)
+	if err != nil {
+		return nil, err
 	}
 
-	return
+	return d.server.FindPeers(ns, options.Limit)
 }
 
-func (s *mockStore) Put(key string, ttl time.Duration, pi libp2p_peer.AddrInfo) {
-	s.muStore.Lock()
-	defer s.muStore.Unlock()
-
-	recs, ok := s.mockRecords[key]
-	if !ok {
-		s.mockRecords[key] = make(mockRecords, 0)
-	}
-
-	var rec *mockRecord
-	for _, r := range recs {
-		if r.pi.ID == pi.ID {
-			rec = r
-		}
-	}
-
-	if rec == nil {
-		rec = newRecord(pi)
-		s.mockRecords[key] = append(s.mockRecords[key], rec)
-	}
-
-	if ttl > 0 {
-		rec.SetTimer(ttl, func() { s.Remove(key, pi.ID) })
-	}
-
-	return
+func (d *mockDriverClient) Unregister(ctx context.Context, ns string) error {
+	d.server.Unregister(ns, d.host.ID())
+	return nil
 }