network/hive: move message handlers from Peer to Hive (#1801)

Author: acud

network/hive.go

@@ -17,6 +17,7 @@
 package network
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"time"
@@ -167,7 +168,7 @@ func (h *Hive) connect() {
 func (h *Hive) tickHive() {
 	addr, depth, changed := h.SuggestPeer()
 	if h.Discovery && changed {
-		NotifyDepth(uint8(depth), h.Kademlia)
+		h.NotifyDepth(uint8(depth))
 	}
 	if addr != nil {
 		log.Trace(fmt.Sprintf("%08x hive connect() suggested %08x", h.BaseAddr()[:4], addr.Address()[:4]))
@@ -194,15 +195,15 @@ func (h *Hive) Run(p *BzzPeer) error {
 	if h.Discovery {
 		if changed {
 			// if depth changed, send to all peers
-			NotifyDepth(depth, h.Kademlia)
+			h.NotifyDepth(depth)
 		} else {
 			// otherwise just send depth to new peer
 			dp.NotifyDepth(depth)
 		}
-		NotifyPeer(p.BzzAddr, h.Kademlia)
+		h.NotifyPeer(p.BzzAddr)
 	}
 	defer h.Off(dp)
-	return dp.Run(dp.HandleMsg)
+	return dp.Run(h.handleMsg(dp))
 }
 
 func (h *Hive) trackPeer(p *BzzPeer) {
@@ -329,3 +330,84 @@ func (h *Hive) savePeers() error {
 	}
 	return nil
 }
+
+var sortPeers = noSortPeers
+
+// handleMsg is the message handler that delegates incoming messages
+func (h *Hive) handleMsg(p *Peer) func(context.Context, interface{}) error {
+	return func(ctx context.Context, msg interface{}) error {
+		switch msg := msg.(type) {
+		case *peersMsg:
+			return h.handlePeersMsg(p, msg)
+		case *subPeersMsg:
+			return h.handleSubPeersMsg(ctx, p, msg)
+		}
+
+		return fmt.Errorf("unknown message type: %T", msg)
+	}
+}
+
+// NotifyDepth sends a message to all connections if depth of saturation is changed
+func (h *Hive) NotifyDepth(depth uint8) {
+	f := func(val *Peer, po int) bool {
+		val.NotifyDepth(depth)
+		return true
+	}
+	h.EachConn(nil, 255, f)
+}
+
+// NotifyPeer informs all peers about a newly added node
+func (h *Hive) NotifyPeer(p *BzzAddr) {
+	f := func(val *Peer, po int) bool {
+		val.NotifyPeer(p, uint8(po))
+		return true
+	}
+	h.EachConn(p.Address(), 255, f)
+}
+
+// handlePeersMsg called by the protocol when receiving peerset (for target address)
+// list of nodes ([]PeerAddr in peersMsg) is added to the overlay db using the
+// Register interface method
+func (h *Hive) handlePeersMsg(d *Peer, msg *peersMsg) error {
+	// register all addresses
+	if len(msg.Peers) == 0 {
+		return nil
+	}
+	for _, a := range msg.Peers {
+		d.seen(a)
+		h.NotifyPeer(a)
+	}
+	return h.Register(msg.Peers...)
+}
+
+// handleSubPeersMsg handles incoming subPeersMsg
+// this message represents the saturation depth of the remote peer
+// saturation depth is the radius within which the peer subscribes to peers
+// the first time this is received we send peer info on all
+// our connected peers that fall within peers saturation depth
+// otherwise this depth is just recorded on the peer, so that
+// subsequent new connections are sent iff they fall within the radius
+func (h *Hive) handleSubPeersMsg(ctx context.Context, d *Peer, msg *subPeersMsg) error {
+	d.setDepth(msg.Depth)
+	// only send peers after the initial subPeersMsg
+	if !d.sentPeers {
+		var peers []*BzzAddr
+		// iterate connection in ascending order of disctance from the remote address
+		h.EachConn(d.Over(), 255, func(p *Peer, po int) bool {
+			// terminate if we are beyond the radius
+			if uint8(po) < msg.Depth {
+				return false
+			}
+			if !d.seen(p.BzzAddr) { // here just records the peer sent
+				peers = append(peers, p.BzzAddr)
+			}
+			return true
+		})
+		// if useful  peers are found, send them over
+		if len(peers) > 0 {
+			go d.Send(ctx, &peersMsg{Peers: sortPeers(peers)})
+		}
+	}
+	d.sentPeers = true
+	return nil
+}

network/discovery.go network/peer.go

@@ -25,16 +25,12 @@ import (
 	"github.com/ethersphere/swarm/pot"
 )
 
-// discovery bzz extension for requesting and relaying node address records
-
-var sortPeers = noSortPeers
-
 // Peer wraps BzzPeer and embeds Kademlia overlay connectivity driver
 type Peer struct {
 	*BzzPeer
 	kad       *Kademlia
 	sentPeers bool            // whether we already sent peer closer to this address
-	mtx       sync.RWMutex    //
+	mtx       sync.RWMutex    // protect peers map
 	peers     map[string]bool // tracks node records sent to the peer
 	depth     uint8           // the proximity order advertised by remote as depth of saturation
 }
@@ -51,39 +47,6 @@ func NewPeer(p *BzzPeer, kad *Kademlia) *Peer {
 	return d
 }
 
-// HandleMsg is the message handler that delegates incoming messages
-func (d *Peer) HandleMsg(ctx context.Context, msg interface{}) error {
-	switch msg := msg.(type) {
-
-	case *peersMsg:
-		return d.handlePeersMsg(msg)
-
-	case *subPeersMsg:
-		return d.handleSubPeersMsg(msg)
-
-	default:
-		return fmt.Errorf("unknown message type: %T", msg)
-	}
-}
-
-// NotifyDepth sends a message to all connections if depth of saturation is changed
-func NotifyDepth(depth uint8, kad *Kademlia) {
-	f := func(val *Peer, po int) bool {
-		val.NotifyDepth(depth)
-		return true
-	}
-	kad.EachConn(nil, 255, f)
-}
-
-// NotifyPeer informs all peers about a newly added node
-func NotifyPeer(p *BzzAddr, k *Kademlia) {
-	f := func(val *Peer, po int) bool {
-		val.NotifyPeer(p, uint8(po))
-		return true
-	}
-	k.EachConn(p.Address(), 255, f)
-}
-
 // NotifyPeer notifies the remote node (recipient) about a peer if
 // the peer's PO is within the recipients advertised depth
 // OR the peer is closer to the recipient than self
@@ -154,21 +117,6 @@ func (msg peersMsg) String() string {
 	return fmt.Sprintf("%T: %v", msg, msg.Peers)
 }
 
-// handlePeersMsg called by the protocol when receiving peerset (for target address)
-// list of nodes ([]PeerAddr in peersMsg) is added to the overlay db using the
-// Register interface method
-func (d *Peer) handlePeersMsg(msg *peersMsg) error {
-	// register all addresses
-	if len(msg.Peers) == 0 {
-		return nil
-	}
-	for _, a := range msg.Peers {
-		d.seen(a)
-		NotifyPeer(a, d.kad)
-	}
-	return d.kad.Register(msg.Peers...)
-}
-
 // subPeers msg is communicating the depth of the overlay table of a peer
 type subPeersMsg struct {
 	Depth uint8
@@ -179,38 +127,6 @@ func (msg subPeersMsg) String() string {
 	return fmt.Sprintf("%T: request peers > PO%02d. ", msg, msg.Depth)
 }
 
-// handleSubPeersMsg handles incoming subPeersMsg
-// this message represents the saturation depth of the remote peer
-// saturation depth is the radius within which the peer subscribes to peers
-// the first time this is received we send peer info on all
-// our connected peers that fall within peers saturation depth
-// otherwise this depth is just recorded on the peer, so that
-// subsequent new connections are sent iff they fall within the radius
-func (d *Peer) handleSubPeersMsg(msg *subPeersMsg) error {
-	d.setDepth(msg.Depth)
-	// only send peers after the initial subPeersMsg
-	if !d.sentPeers {
-		var peers []*BzzAddr
-		// iterate connection in ascending order of disctance from the remote address
-		d.kad.EachConn(d.Over(), 255, func(p *Peer, po int) bool {
-			// terminate if we are beyond the radius
-			if uint8(po) < msg.Depth {
-				return false
-			}
-			if !d.seen(p.BzzAddr) { // here just records the peer sent
-				peers = append(peers, p.BzzAddr)
-			}
-			return true
-		})
-		// if useful  peers are found, send them over
-		if len(peers) > 0 {
-			go d.Send(context.TODO(), &peersMsg{Peers: sortPeers(peers)})
-		}
-	}
-	d.sentPeers = true
-	return nil
-}
-
 // seen takes a peer address and checks if it was sent to a peer already
 // if not, marks the peer as sent
 func (d *Peer) seen(p *BzzAddr) bool {