diff --git a/src/avalanche/peermanager.h b/src/avalanche/peermanager.h --- a/src/avalanche/peermanager.h +++ b/src/avalanche/peermanager.h @@ -121,6 +121,8 @@ * Node API. */ bool addNodeToPeer(PeerId peerid, NodeId nodeid, CPubKey pubkey); + bool removeNode(NodeId nodeid); + NodeId getSuitableNodeToQuery(); /** diff --git a/src/avalanche/peermanager.cpp b/src/avalanche/peermanager.cpp --- a/src/avalanche/peermanager.cpp +++ b/src/avalanche/peermanager.cpp @@ -35,6 +35,14 @@ slots[i] = slots[i].withPeerId(NO_PEER); } + // Remove nodes associated with this peer, unless their timeout is still + // active. This ensure that we don't overquery them in case their are + // subsequently added to another peer. + auto &nview = nodes.get(); + nview.erase(nview.lower_bound(boost::make_tuple(p, TimePoint())), + nview.upper_bound( + boost::make_tuple(p, std::chrono::steady_clock::now()))); + peers.erase(it); return true; } @@ -100,6 +108,10 @@ }); } +bool PeerManager::removeNode(NodeId nodeid) { + return nodes.erase(nodeid) > 0; +} + NodeId PeerManager::getSuitableNodeToQuery() { for (int retry = 0; retry < SELECT_NODE_MAX_RETRY; retry++) { const PeerId p = selectPeer(); @@ -112,9 +124,9 @@ } // See if that peer has an available node. - auto it = nodes.get().lower_bound( - boost::make_tuple(p, TimePoint())); - if (it != nodes.get().end() && it->peerid == p && + auto &nview = nodes.get(); + auto it = nview.lower_bound(boost::make_tuple(p, TimePoint())); + if (it != nview.end() && it->peerid == p && it->nextRequestTime <= std::chrono::steady_clock::now()) { return it->nodeid; } diff --git a/src/avalanche/test/peermanager_tests.cpp b/src/avalanche/test/peermanager_tests.cpp --- a/src/avalanche/test/peermanager_tests.cpp +++ b/src/avalanche/test/peermanager_tests.cpp @@ -350,7 +350,7 @@ BOOST_CHECK_EQUAL(pm.getFragmentation(), 0); } -BOOST_AUTO_TEST_CASE(add_and_get_node) { +BOOST_AUTO_TEST_CASE(node_crud) { PeerManager pm; // Create one peer. @@ -364,7 +364,15 @@ for (int i = 0; i < 100; i++) { NodeId n = pm.getSuitableNodeToQuery(); - BOOST_CHECK((n >= 0 && n < 4) || n == NO_PEER); + BOOST_CHECK(n >= 0 && n < 4); + } + + // Remove a node, check that it doesn't show up. + BOOST_CHECK(pm.removeNode(2)); + + for (int i = 0; i < 100; i++) { + NodeId n = pm.getSuitableNodeToQuery(); + BOOST_CHECK(n == 0 || n == 1 || n == 3); } }