From 0ed0cb58f89126b32c0ce8b9bc47a5d64029c562 Mon Sep 17 00:00:00 2001
From: Preston Van Loon <preston@prysmaticlabs.com>
Date: Sat, 18 Apr 2020 18:53:32 -0700
Subject: [PATCH] Configurable pubsub router (#5499)

* A p2p flag for configuring pubsub router
* Add randomsub
* lint
* Merge refs/heads/master into configurable-pubsub-router
* Default to gossip
* Merge branch 'configurable-pubsub-router' of github.com:prysmaticlabs/prysm into configurable-pubsub-router
* Merge refs/heads/master into configurable-pubsub-router
* Add flag to usage
* Merge branch 'configurable-pubsub-router' of github.com:prysmaticlabs/prysm into configurable-pubsub-router
* Fix build
* Merge refs/heads/master into configurable-pubsub-router
---
 beacon-chain/main.go        |  1 +
 beacon-chain/node/node.go   |  1 +
 beacon-chain/p2p/config.go  |  1 +
 beacon-chain/p2p/service.go | 22 +++++++++++++++++++++-
 beacon-chain/usage.go       |  1 +
 shared/cmd/flags.go         |  6 ++++++
 6 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/beacon-chain/main.go b/beacon-chain/main.go
index cf6f6a4b4bb..5765b015dc7 100644
--- a/beacon-chain/main.go
+++ b/beacon-chain/main.go
@@ -63,6 +63,7 @@ var appFlags = []cli.Flag{
 	cmd.P2PMetadata,
 	cmd.P2PWhitelist,
 	cmd.P2PEncoding,
+	cmd.P2PPubsub,
 	cmd.DataDirFlag,
 	cmd.VerbosityFlag,
 	cmd.EnableTracingFlag,
diff --git a/beacon-chain/node/node.go b/beacon-chain/node/node.go
index ee11f8dab4b..135b86687a9 100644
--- a/beacon-chain/node/node.go
+++ b/beacon-chain/node/node.go
@@ -312,6 +312,7 @@ func (b *BeaconNode) registerP2P(ctx *cli.Context) error {
 		DisableDiscv5:     ctx.Bool(flags.DisableDiscv5.Name),
 		Encoding:          ctx.String(cmd.P2PEncoding.Name),
 		StateNotifier:     b,
+		PubSub:            ctx.String(cmd.P2PPubsub.Name),
 	})
 	if err != nil {
 		return err
diff --git a/beacon-chain/p2p/config.go b/beacon-chain/p2p/config.go
index c325b6e68f5..ac5c3c8917c 100644
--- a/beacon-chain/p2p/config.go
+++ b/beacon-chain/p2p/config.go
@@ -27,4 +27,5 @@ type Config struct {
 	WhitelistCIDR         string
 	Encoding              string
 	StateNotifier         statefeed.Notifier
+	PubSub                string
 }
diff --git a/beacon-chain/p2p/service.go b/beacon-chain/p2p/service.go
index 77815f8d34b..00f33b89cd0 100644
--- a/beacon-chain/p2p/service.go
+++ b/beacon-chain/p2p/service.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/ecdsa"
+	"fmt"
 	"strconv"
 	"strings"
 	"time"
@@ -56,6 +57,12 @@ const prysmProtocolPrefix = "/prysm/0.0.0"
 // maxBadResponses is the maximum number of bad responses from a peer before we stop talking to it.
 const maxBadResponses = 3
 
+const (
+	pubsubFlood  = "flood"
+	pubsubGossip = "gossip"
+	pubsubRandom = "random"
+)
+
 // Service for managing peer to peer (p2p) networking.
 type Service struct {
 	started               bool
@@ -153,7 +160,20 @@ func NewService(cfg *Config) (*Service, error) {
 		pubsub.WithStrictSignatureVerification(false),
 		pubsub.WithMessageIdFn(msgIDFunction),
 	}
-	gs, err := pubsub.NewGossipSub(s.ctx, s.host, psOpts...)
+
+	var gs *pubsub.PubSub
+	if cfg.PubSub == "" {
+		cfg.PubSub = pubsubGossip
+	}
+	if cfg.PubSub == pubsubFlood {
+		gs, err = pubsub.NewFloodSub(s.ctx, s.host, psOpts...)
+	} else if cfg.PubSub == pubsubGossip {
+		gs, err = pubsub.NewGossipSub(s.ctx, s.host, psOpts...)
+	} else if cfg.PubSub == pubsubRandom {
+		gs, err = pubsub.NewRandomSub(s.ctx, s.host, psOpts...)
+	} else {
+		return nil, fmt.Errorf("unknown pubsub type %s", cfg.PubSub)
+	}
 	if err != nil {
 		log.WithError(err).Error("Failed to start pubsub")
 		return nil, err
diff --git a/beacon-chain/usage.go b/beacon-chain/usage.go
index 311bf07c338..685aa767509 100644
--- a/beacon-chain/usage.go
+++ b/beacon-chain/usage.go
@@ -108,6 +108,7 @@ var appHelpFlagGroups = []flagGroup{
 			cmd.StaticPeers,
 			cmd.EnableUPnPFlag,
 			cmd.P2PEncoding,
+			cmd.P2PPubsub,
 			flags.MinSyncPeers,
 		},
 	},
diff --git a/shared/cmd/flags.go b/shared/cmd/flags.go
index 6533587d8e0..35f955e6fe3 100644
--- a/shared/cmd/flags.go
+++ b/shared/cmd/flags.go
@@ -137,6 +137,12 @@ var (
 		Usage: "The encoding format of messages sent over the wire. The default is 0, which represents ssz",
 		Value: "ssz-snappy",
 	}
+	// P2PPubsub defines the pubsub router to use for p2p messages.
+	P2PPubsub = &cli.StringFlag{
+		Name:  "p2p-pubsub",
+		Usage: "The name of the pubsub router to use. Supported values are: gossip, flood, random",
+		Value: "gossip",
+	}
 	// ForceClearDB removes any previously stored data at the data directory.
 	ForceClearDB = &cli.BoolFlag{
 		Name:  "force-clear-db",