feat: add qgb dht wait for peers (#141)

* feat: add qgb dht wait for peers

* fix: close dhts in test

p2p/dht.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"time"
 
+	tmlog "github.com/tendermint/tendermint/libs/log"
+
 	"github.com/celestiaorg/orchestrator-relayer/types"
 
 	ds "github.com/ipfs/go-datastore"
@@ -21,10 +23,11 @@ const (
 // Used to add helper methods to easily handle the DHT.
 type QgbDHT struct {
 	*dht.IpfsDHT
+	logger tmlog.Logger
 }
 
 // NewQgbDHT create a new IPFS DHT using a suitable configuration for the QGB.
-func NewQgbDHT(ctx context.Context, h host.Host, store ds.Batching) (*QgbDHT, error) {
+func NewQgbDHT(ctx context.Context, h host.Host, store ds.Batching, logger tmlog.Logger) (*QgbDHT, error) {
 	router, err := dht.New(
 		ctx,
 		h,
@@ -39,7 +42,42 @@ func NewQgbDHT(ctx context.Context, h host.Host, store ds.Batching) (*QgbDHT, er
 		return nil, err
 	}
 
-	return &QgbDHT{router}, nil
+	return &QgbDHT{
+		IpfsDHT: router,
+		logger:  logger,
+	}, nil
+}
+
+// WaitForPeers waits for peers to be connected to the DHT.
+// Returns nil if the context is done or the peers list has more peers than the specified peersThreshold.
+// Returns error if it times out.
+func (q QgbDHT) WaitForPeers(ctx context.Context, timeout time.Duration, rate time.Duration, peersThreshold int) error {
+	if peersThreshold < 1 {
+		return ErrPeersThresholdCannotBeNegative
+	}
+
+	t := time.After(timeout)
+	ticker := time.NewTicker(rate)
+	for {
+		select {
+		case <-ctx.Done():
+			return nil
+		case <-t:
+			return ErrPeersTimeout
+		case <-ticker.C:
+			peersLen := len(q.RoutingTable().ListPeers())
+			if peersLen >= peersThreshold {
+				return nil
+			}
+			q.logger.Info(
+				"waiting for routing table to populate",
+				"target number of peers",
+				peersThreshold,
+				"current number",
+				peersLen,
+			)
+		}
+	}
 }
 
 // Note: The Get and Put methods do not run any validations on the data commitment confirms

p2p/dht_test.go

@@ -3,6 +3,10 @@ package p2p_test
 import (
 	"context"
 	"testing"
+	"time"
+
+	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/stretchr/testify/require"
 
 	"github.com/celestiaorg/orchestrator-relayer/p2p"
 	qgbtesting "github.com/celestiaorg/orchestrator-relayer/testing"
@@ -137,3 +141,29 @@ func TestNetworkGetNonExistentValsetConfirm(t *testing.T) {
 	assert.Error(t, err)
 	assert.True(t, types.IsEmptyValsetConfirm(actualConfirm))
 }
+
+func TestWaitForPeers(t *testing.T) {
+	ctx := context.Background()
+	// create first dht
+	h1, _, dht1 := qgbtesting.NewTestDHT(ctx)
+	defer dht1.Close()
+
+	// wait for peers
+	err := dht1.WaitForPeers(ctx, 10*time.Millisecond, time.Millisecond, 1)
+	// should error because no peer is connected to this dht
+	assert.Error(t, err)
+
+	// create second dht
+	h2, _, dht2 := qgbtesting.NewTestDHT(ctx)
+	defer dht2.Close()
+	// connect to first dht
+	err = h2.Connect(ctx, peer.AddrInfo{
+		ID:    h1.ID(),
+		Addrs: h1.Addrs(),
+	})
+	require.NoError(t, err)
+
+	// wait for peers
+	err = dht1.WaitForPeers(ctx, 10*time.Millisecond, time.Millisecond, 1)
+	assert.NoError(t, err)
+}

p2p/errors.go

@@ -3,6 +3,8 @@ package p2p
 import "errors"
 
 var (
+	ErrPeersTimeout                    = errors.New("timeout while waiting for peers")
+	ErrPeersThresholdCannotBeNegative  = errors.New("peers threshold cannot be negative")
 	ErrNilPrivateKey                   = errors.New("private key cannot be nil")
 	ErrNotEnoughValsetConfirms         = errors.New("couldn't find enough valset confirms")
 	ErrNotEnoughDataCommitmentConfirms = errors.New("couldn't find enough data commitment confirms")

testing/dht_network.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"time"
 
+	tmlog "github.com/tendermint/tendermint/libs/log"
+
 	"github.com/celestiaorg/orchestrator-relayer/p2p"
 	ds "github.com/ipfs/go-datastore"
 	dssync "github.com/ipfs/go-datastore/sync"
@@ -68,7 +70,7 @@ func NewTestDHT(ctx context.Context) (host.Host, ds.Batching, *p2p.QgbDHT) {
 		panic(err)
 	}
 	dataStore := dssync.MutexWrap(ds.NewMapDatastore())
-	dht, err := p2p.NewQgbDHT(ctx, h, dataStore)
+	dht, err := p2p.NewQgbDHT(ctx, h, dataStore, tmlog.NewNopLogger())
 	if err != nil {
 		panic(err)
 	}