diff --git a/network/default.go b/network/default.go index ab8905c7..2da925eb 100644 --- a/network/default.go +++ b/network/default.go @@ -276,7 +276,7 @@ func (n *network) processNetChan(client transport.Client, listener tunnel.Listen } n.Unlock() // get all the node peers down to MaxDepth encoded in protobuf - msg := n.node.getProtoTopology(MaxDepth) + msg := PeersToProto(n.node, n.Peers(), MaxDepth) // advertise yourself to the network if err := n.sendMsg("peer", msg, NetworkChannel); err != nil { log.Debugf("Network failed to advertise peers: %v", err) @@ -393,7 +393,7 @@ func (n *network) announce(client transport.Client) { case <-n.closed: return case <-announce.C: - msg := n.node.getProtoTopology(MaxDepth) + msg := PeersToProto(n.node, n.Peers(), MaxDepth) // advertise yourself to the network if err := n.sendMsg("peer", msg, NetworkChannel); err != nil { log.Debugf("Network failed to advertise peers: %v", err) diff --git a/network/handler/handler.go b/network/handler/handler.go index 486b3fa5..759e21a4 100644 --- a/network/handler/handler.go +++ b/network/handler/handler.go @@ -15,35 +15,6 @@ type Network struct { Network network.Network } -// toplogyToProto recursively traverses node topology and returns it -func peerTopology(peer network.Node, depth uint) *pbNet.Peer { - node := &pbNet.Node{ - Id: peer.Id(), - Address: peer.Address(), - } - - pbPeers := &pbNet.Peer{ - Node: node, - Peers: make([]*pbNet.Peer, 0), - } - - // return if we reached the end of topology or depth - if depth == 0 || len(peer.Peers()) == 0 { - return pbPeers - } - - // decrement the depth - depth-- - - // iterate through peers of peers aka pops - for _, pop := range peer.Peers() { - peer := peerTopology(pop, depth) - pbPeers.Peers = append(pbPeers.Peers, peer) - } - - return pbPeers -} - // ListPeers returns a list of all the nodes the node has a direct link with func (n *Network) ListPeers(ctx context.Context, req *pbNet.PeerRequest, resp *pbNet.PeerResponse) error { depth := uint(req.Depth) @@ -54,21 +25,8 @@ func (n *Network) ListPeers(ctx context.Context, req *pbNet.PeerRequest, resp *p // get node peers nodePeers := n.Network.Peers() - // network node aka root node - node := &pbNet.Node{ - Id: n.Network.Id(), - Address: n.Network.Address(), - } - // we will build proto topology into this - peers := &pbNet.Peer{ - Node: node, - Peers: make([]*pbNet.Peer, 0), - } - - for _, nodePeer := range nodePeers { - peer := peerTopology(nodePeer, depth) - peers.Peers = append(peers.Peers, peer) - } + // get peers encoded into protobuf + peers := network.PeersToProto(n.Network, nodePeers, depth) resp.Peers = peers diff --git a/network/node.go b/network/node.go index dad609e5..b5a318ee 100644 --- a/network/node.go +++ b/network/node.go @@ -128,42 +128,22 @@ func (n *node) Peers() []Node { return peers } -// getProtoTopology returns node topology down to the given depth encoded in protobuf -func (n *node) getProtoTopology(depth uint) *pb.Peer { - n.RLock() - defer n.RUnlock() +// updateTopology updates node peer topology down to given depth +func (n *node) updatePeerTopology(pbPeer *pb.Peer, depth uint) error { + n.Lock() + defer n.Unlock() - node := &pb.Node{ - Id: n.id, - Address: n.address, + if pbPeer == nil { + return errors.New("peer not initialized") } - pbPeers := &pb.Peer{ - Node: node, - Peers: make([]*pb.Peer, 0), - } + // unpack Peer topology into *node + peer := unpackPeer(pbPeer, depth) - // return if have either reached the depth or have no more peers - if depth == 0 || len(n.peers) == 0 { - return pbPeers - } + // update node peers with new topology + n.peers[pbPeer.Node.Id] = peer - // decrement the depth - depth-- - - var peers []*pb.Peer - for _, peer := range n.peers { - // get peers of the node peers - // NOTE: this is [not] a recursive call - pbPeerPeer := peer.getProtoTopology(depth) - // add current peer to explored peers - peers = append(peers, pbPeerPeer) - } - - // add peers to the parent topology - pbPeers.Peers = peers - - return pbPeers + return nil } // unpackPeer unpacks pb.Peer into node topology of given depth @@ -194,20 +174,51 @@ func unpackPeer(pbPeer *pb.Peer, depth uint) *node { return peerNode } -// updateTopology updates node peer topology down to given depth -func (n *node) updatePeerTopology(pbPeer *pb.Peer, depth uint) error { - n.Lock() - defer n.Unlock() - - if pbPeer == nil { - return errors.New("peer not initialized") +func peerTopology(peer Node, depth uint) *pb.Peer { + node := &pb.Node{ + Id: peer.Id(), + Address: peer.Address(), } - // unpack Peer topology into *node - peer := unpackPeer(pbPeer, depth) + pbPeers := &pb.Peer{ + Node: node, + Peers: make([]*pb.Peer, 0), + } - // update node peers with new topology - n.peers[pbPeer.Node.Id] = peer + // return if we reached the end of topology or depth + if depth == 0 || len(peer.Peers()) == 0 { + return pbPeers + } - return nil + // decrement the depth + depth-- + + // iterate through peers of peers aka pops + for _, pop := range peer.Peers() { + peer := peerTopology(pop, depth) + pbPeers.Peers = append(pbPeers.Peers, peer) + } + + return pbPeers +} + +// PeersToProto returns node peers graph encoded into protobuf +func PeersToProto(root Node, peers []Node, depth uint) *pb.Peer { + // network node aka root node + node := &pb.Node{ + Id: root.Id(), + Address: root.Address(), + } + // we will build proto topology into this + pbPeers := &pb.Peer{ + Node: node, + Peers: make([]*pb.Peer, 0), + } + + for _, peer := range peers { + pbPeer := peerTopology(peer, depth) + pbPeers.Peers = append(pbPeers.Peers, pbPeer) + } + + return pbPeers } diff --git a/network/node_test.go b/network/node_test.go index 2a49be11..f5941c1a 100644 --- a/network/node_test.go +++ b/network/node_test.go @@ -256,7 +256,7 @@ func TestUpdatePeerTopology(t *testing.T) { } } -func TestGetProtoTopology(t *testing.T) { +func TestPeersToProto(t *testing.T) { // single node single := &node{ id: testNodeId, @@ -266,10 +266,10 @@ func TestGetProtoTopology(t *testing.T) { } topCount := 0 - protoTop := single.getProtoTopology(10) + protoPeers := PeersToProto(single, single.Peers(), 0) - if len(protoTop.Peers) != topCount { - t.Errorf("Expected to find %d nodes, found: %d", topCount, len(protoTop.Peers)) + if len(protoPeers.Peers) != topCount { + t.Errorf("Expected to find %d nodes, found: %d", topCount, len(protoPeers.Peers)) } // complicated node graph @@ -282,9 +282,9 @@ func TestGetProtoTopology(t *testing.T) { peerIds[id] = true } // depth 1 should give us immmediate neighbours only - protoTop = node.getProtoTopology(1) + protoPeers = PeersToProto(node, node.Peers(), 1) - if len(protoTop.Peers) != topCount { - t.Errorf("Expected to find %d nodes, found: %d", topCount, len(protoTop.Peers)) + if len(protoPeers.Peers) != topCount { + t.Errorf("Expected to find %d nodes, found: %d", topCount, len(protoPeers.Peers)) } }