Add a Configurable Bootstrap Timeout for DHT Discovery

cmd/config/default.go

@@ -49,6 +49,7 @@ var defaultConfig = harmonyconfig.HarmonyConfig{
 		DialTimeout:                     nodeconfig.DefaultDialTimeout,
 		Muxer:                           nodeconfig.DefaultMuxer,
 		NoRelay:                         nodeconfig.DefaultNoRelay,
+		DiscBootstrapTimeout:            nodeconfig.DefaultDiscBootstrapTimeout,
 	},
 	HTTP: harmonyconfig.HttpConfig{
 		Enabled:        true,

cmd/config/flags.go

@@ -60,6 +60,7 @@ var (
 		p2pKeyFileFlag,
 		p2pDHTDataStoreFlag,
 		p2pDiscoveryConcurrencyFlag,
+		p2pDiscBootstrapTimeoutFlag,
 		legacyKeyFileFlag,
 		p2pDisablePrivateIPScanFlag,
 		maxConnPerIPFlag,
@@ -631,6 +632,11 @@ var (
 		Usage:    "the pubsub's DHT discovery concurrency num (default with raw libp2p dht option)",
 		DefValue: defaultConfig.P2P.DiscConcurrency,
 	}
+	p2pDiscBootstrapTimeoutFlag = cli.DurationFlag{
+		Name:     "p2p.disc.bootstrap-timeout",
+		Usage:    "timeout for the DHT bootstrap process",
+		DefValue: defaultConfig.P2P.DiscBootstrapTimeout,
+	}
 	p2pDisablePrivateIPScanFlag = cli.BoolFlag{
 		Name:     "p2p.no-private-ip-scan",
 		Usage:    "disable scanning of private ip4/6 addresses by DHT",
@@ -729,6 +735,10 @@ func applyP2PFlags(cmd *cobra.Command, config *harmonyconfig.HarmonyConfig) {
 		config.P2P.DiscConcurrency = cli.GetIntFlagValue(cmd, p2pDiscoveryConcurrencyFlag)
 	}
 
+	if cli.IsFlagChanged(cmd, p2pDiscBootstrapTimeoutFlag) {
+		config.P2P.DiscBootstrapTimeout = cli.GetDurationFlagValue(cmd, p2pDiscBootstrapTimeoutFlag)
+	}
+
 	if cli.IsFlagChanged(cmd, maxConnPerIPFlag) {
 		config.P2P.MaxConnsPerIP = cli.GetIntFlagValue(cmd, maxConnPerIPFlag)
 	}

internal/cli/flag.go

@@ -4,6 +4,8 @@
 package cli
 
 import (
+	"time"
+
 	"github.com/spf13/pflag"
 )
 
@@ -130,6 +132,23 @@ func (f IntSliceFlag) RegisterTo(fs *pflag.FlagSet) error {
 	return markHiddenOrDeprecated(fs, f.Name, f.Deprecated, f.Hidden)
 }
 
+// DurationFlag is the flag with duration value
+type DurationFlag struct {
+	Name       string
+	Shorthand  string
+	Usage      string
+	Deprecated string
+	Hidden     bool
+
+	DefValue time.Duration
+}
+
+// RegisterTo register the duration flag to FlagSet
+func (f DurationFlag) RegisterTo(fs *pflag.FlagSet) error {
+	fs.DurationP(f.Name, f.Shorthand, f.DefValue, f.Usage)
+	return markHiddenOrDeprecated(fs, f.Name, f.Deprecated, f.Hidden)
+}
+
 func markHiddenOrDeprecated(fs *pflag.FlagSet, name string, deprecated string, hidden bool) error {
 	if len(deprecated) != 0 {
 		// TODO: after totally removed node.sh, change MarkHidden to MarkDeprecated
@@ -160,6 +179,8 @@ func getFlagName(flag Flag) string {
 		return f.Name
 	case Uint64Flag:
 		return f.Name
+	case DurationFlag:
+		return f.Name
 	}
 	return ""
 }

internal/cli/parse.go

@@ -1,6 +1,8 @@
 package cli
 
 import (
+	"time"
+
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 )
@@ -82,6 +84,11 @@ func GetUint64FlagValue(cmd *cobra.Command, flag Uint64Flag) uint64 {
 	return getUint64FlagValue(cmd.Flags(), flag)
 }
 
+// GetDurationFlagValue get the duration value for the given DurationFlag from the local flags of the cobra command.
+func GetDurationFlagValue(cmd *cobra.Command, flag DurationFlag) time.Duration {
+	return getDurationFlagValue(cmd.Flags(), flag)
+}
+
 // GetIntPersistentFlagValue get the int value for the given IntFlag from the persistent
 // flags of the cobra command.
 func GetIntPersistentFlagValue(cmd *cobra.Command, flag IntFlag) int {
@@ -157,6 +164,15 @@ func getIntSliceFlagValue(fs *pflag.FlagSet, flag IntSliceFlag) []int {
 	return val
 }
 
+func getDurationFlagValue(fs *pflag.FlagSet, flag DurationFlag) time.Duration {
+	val, err := fs.GetDuration(flag.Name)
+	if err != nil {
+		handleParseError(err)
+		return 0
+	}
+	return val
+}
+
 // IsFlagChanged returns whether the flag has been changed in command
 func IsFlagChanged(cmd *cobra.Command, flag Flag) bool {
 	name := getFlagName(flag)

internal/configs/bootnode/bootnode.go

@@ -95,6 +95,8 @@ type P2pConfig struct {
 	Muxer string
 	// No relay services, direct connections between peers only
 	NoRelay bool
+	// Timeout for the initial DHT bootstrap process
+	DiscBootstrapTimeout time.Duration
 }
 
 type GeneralConfig struct {

internal/configs/harmony/harmony.go

@@ -146,6 +146,8 @@ type P2pConfig struct {
 	Muxer string
 	// No relay services, direct connections between peers only
 	NoRelay bool
+	// Timeout for the initial DHT bootstrap process
+	DiscBootstrapTimeout time.Duration
 }
 
 type GeneralConfig struct {

internal/configs/node/network.go

@@ -93,6 +93,8 @@ const (
 	DefaultUserAgent = ""
 	// DefaultDialTimeout dial timeout
 	DefaultDialTimeout = time.Minute
+	// DefaultDiscBootstrapTimeout is the timeout for DHT bootstrap (0=unlimited)
+	DefaultDiscBootstrapTimeout = 0
 	// DefaultMuxerType P2P multiplexer type
 	DefaultMuxer = "mplex, yamux"
 	// DefaultNoRelay disables p2p host relay

p2p/discovery/discovery.go

@@ -54,7 +54,18 @@ func NewDHTDiscovery(ctx context.Context, cancel context.CancelFunc, host libp2p
 
 // Start bootstrap the dht discovery service.
 func (d *dhtDiscovery) Start() error {
-	return d.dht.Bootstrap(d.ctx)
+	ctx := d.ctx
+	if d.opt.BootstrapTimeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, d.opt.BootstrapTimeout)
+		defer cancel()
+	}
+	start := time.Now()
+	err := d.dht.Bootstrap(ctx)
+	if err == nil {
+		d.logger.Debug().Dur("duration", time.Since(start)).Msg("dht bootstrap completed")
+	}
+	return err
 }
 
 // Stop stop the dhtDiscovery service

p2p/discovery/option.go

@@ -1,6 +1,8 @@
 package discovery
 
 import (
+	"time"
+
 	"github.com/pkg/errors"
 
 	p2ptypes "github.com/harmony-one/harmony/p2p/types"
@@ -15,7 +17,10 @@ type DHTConfig struct {
 	BootNodes       []string
 	DataStoreFile   *string // File path to store DHT data. Shall be only used for bootstrap nodes.
 	DiscConcurrency int
-	DHT             *dht.IpfsDHT
+	// BootstrapTimeout defines how long the DHT bootstrap should run. Zero
+	// means using the passed in context without additional timeout.
+	BootstrapTimeout time.Duration
+	DHT              *dht.IpfsDHT
 }
 
 // GetLibp2pRawOptions get the raw libp2p options as a slice.

p2p/host.go

@@ -107,6 +107,7 @@ type HostConfig struct {
 	TrustedNodes                    []string
 	DataStoreFile                   *string
 	DiscConcurrency                 int
+	DiscBootstrapTimeout            time.Duration
 	MaxConnPerIP                    int
 	DisablePrivateIPScan            bool
 	MaxPeers                        int64
@@ -315,9 +316,10 @@ func NewHost(cfg HostConfig) (Host, error) {
 
 	// DHT
 	opt := discovery.DHTConfig{
-		BootNodes:       cfg.BootNodes,
-		DataStoreFile:   cfg.DataStoreFile,
-		DiscConcurrency: cfg.DiscConcurrency,
+		BootNodes:        cfg.BootNodes,
+		DataStoreFile:    cfg.DataStoreFile,
+		DiscConcurrency:  cfg.DiscConcurrency,
+		BootstrapTimeout: cfg.DiscBootstrapTimeout,
 	}
 	opts, err := opt.GetLibp2pRawOptions()
 	if err != nil {