cleaned up connection registry

node/node.go

@@ -39,6 +39,8 @@ import (
 	"github.com/ava-labs/avalanchego/utils/constants"
 	"github.com/ava-labs/avalanchego/utils/hashing"
 	"github.com/ava-labs/avalanchego/utils/logging"
+	"github.com/ava-labs/avalanchego/utils/math"
+	"github.com/ava-labs/avalanchego/utils/timer"
 	"github.com/ava-labs/avalanchego/utils/wrappers"
 	"github.com/ava-labs/avalanchego/version"
 	"github.com/ava-labs/avalanchego/vms"
@@ -170,6 +172,26 @@ func (n *Node) initNetworking() error {
 		}
 	}
 
+	bootstrapWeight := n.beacons.Weight()
+	reqWeight := (3*bootstrapWeight + 3) / 4
+
+	if reqWeight > 0 {
+		timer := timer.NewTimer(func() {
+			n.Log.Fatal("Failed to connect to bootstrap nodes. Node shutting down...")
+			go n.Net.Close()
+		})
+
+		go timer.Dispatch()
+		timer.SetTimeoutIn(15 * time.Second)
+
+		consensusRouter = &beaconManager{
+			Router:         consensusRouter,
+			timer:          timer,
+			beacons:        n.beacons,
+			requiredWeight: reqWeight,
+		}
+	}
+
 	n.Net = network.NewDefaultNetwork(
 		n.Config.ConsensusParams.Metrics,
 		n.Log,
@@ -213,6 +235,46 @@ func (i *insecureValidatorManager) Disconnected(vdrID ids.ShortID) {
 	i.Router.Disconnected(vdrID)
 }
 
+type beaconManager struct {
+	router.Router
+	timer          *timer.Timer
+	beacons        validators.Set
+	requiredWeight uint64
+	weight         uint64
+}
+
+func (b *beaconManager) Connected(vdrID ids.ShortID) {
+	weight, ok := b.beacons.GetWeight(vdrID)
+	if !ok {
+		b.Router.Connected(vdrID)
+		return
+	}
+	weight, err := math.Add64(weight, b.weight)
+	if err != nil {
+		b.timer.Cancel()
+		b.Router.Connected(vdrID)
+		return
+	}
+	b.weight = weight
+	if b.weight >= b.requiredWeight {
+		b.timer.Cancel()
+	}
+	b.Router.Connected(vdrID)
+}
+
+func (b *beaconManager) Disconnected(vdrID ids.ShortID) {
+	if weight, ok := b.beacons.GetWeight(vdrID); ok {
+		// TODO: Account for weight changes in a more robust manner.
+
+		// Sub64 should rarely error since only validators that have added their
+		// weight can become disconnected. Because it is possible that there are
+		// changes to the validators set, we utilize that Sub64 returns 0 on
+		// error.
+		b.weight, _ = math.Sub64(b.weight, weight)
+	}
+	b.Router.Disconnected(vdrID)
+}
+
 // Dispatch starts the node's servers.
 // Returns when the node exits.
 func (n *Node) Dispatch() error {
@@ -366,30 +428,6 @@ func (n *Node) initChains(genesisBytes []byte, avaxAssetID ids.ID) error {
 		CustomBeacons: n.beacons,
 	})
 
-	// TODO: Introduce this back in
-
-	// bootstrapWeight := n.beacons.Weight()
-
-	// reqWeight := (3*bootstrapWeight + 3) / 4
-
-	// if reqWeight == 0 {
-	// 	return nil
-	// }
-
-	// connectToBootstrapsTimeout := timer.NewTimer(func() {
-	// 	n.Log.Fatal("Failed to connect to bootstrap nodes. Node shutting down...")
-	// 	go n.Net.Close()
-	// })
-
-	// awaiter := chains.NewAwaiter(n.beacons, reqWeight, func() {
-	// 	n.Log.Info("Connected to required bootstrap nodes. Starting Platform Chain...")
-	// 	connectToBootstrapsTimeout.Cancel()
-	// })
-
-	// go connectToBootstrapsTimeout.Dispatch()
-	// connectToBootstrapsTimeout.SetTimeoutIn(15 * time.Second)
-
-	// n.Net.RegisterConnector(awaiter)
 	return nil
 }
 

snow/engine/common/bootstrapper.go

@@ -176,7 +176,7 @@ func (b *Bootstrapper) Accepted(validatorID ids.ShortID, requestID uint32, conta
 
 // Connected implements the Engine interface.
 func (b *Bootstrapper) Connected(validatorID ids.ShortID) error {
-	weight, ok := b.Validators.GetWeight(validatorID)
+	weight, ok := b.Beacons.GetWeight(validatorID)
 	if !ok {
 		return nil
 	}
@@ -193,7 +193,7 @@ func (b *Bootstrapper) Connected(validatorID ids.ShortID) error {
 
 // Disconnected implements the Engine interface.
 func (b *Bootstrapper) Disconnected(validatorID ids.ShortID) error {
-	if weight, ok := b.Validators.GetWeight(validatorID); ok {
+	if weight, ok := b.Beacons.GetWeight(validatorID); ok {
 		// TODO: Account for weight changes in a more robust manner.
 
 		// Sub64 should rarely error since only validators that have added their

snow/networking/router/handler.go

@@ -441,14 +441,16 @@ func (h *Handler) QueryFailed(validatorID ids.ShortID, requestID uint32) {
 // Connected passes a new connection notification to the consensus engine
 func (h *Handler) Connected(validatorID ids.ShortID) {
 	h.sendReliableMsg(message{
-		messageType: gossipMsg,
+		messageType: connectedMsg,
+		validatorID: validatorID,
 	})
 }
 
 // Disconnected passes a new connection notification to the consensus engine
 func (h *Handler) Disconnected(validatorID ids.ShortID) {
 	h.sendReliableMsg(message{
-		messageType: gossipMsg,
+		messageType: disconnectedMsg,
+		validatorID: validatorID,
 	})
 }