Skip to content

Commit

Permalink
fix: keep peer store address unique
Browse files Browse the repository at this point in the history
  • Loading branch information
driftluo committed Sep 5, 2023
1 parent e585192 commit 9c26137
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 8 deletions.
15 changes: 9 additions & 6 deletions network/src/peer_store/addr_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,14 +18,17 @@ impl AddrManager {
/// Add an address information to address manager
pub fn add(&mut self, mut addr_info: AddrInfo) {
if let Some(key) = multiaddr_to_socketaddr(&addr_info.addr) {
if let Some(exists_last_connected_at_ms) = self
.get(&addr_info.addr)
.map(|addr| addr.last_connected_at_ms)
{
if let Some(&id) = self.addr_to_id.get(&key) {
let (exist_last_connected_at_ms, random_id_pos) = {
let info = self.id_to_info.get(&id).expect("must exists");
(info.last_connected_at_ms, info.random_id_pos)
};
// Get time earlier than record time, return directly
if addr_info.last_connected_at_ms < exists_last_connected_at_ms {
return;
if addr_info.last_connected_at_ms >= exist_last_connected_at_ms {
addr_info.random_id_pos = random_id_pos;
self.id_to_info.insert(id, addr_info);
}
return;
}

let id = self.next_id;
Expand Down
2 changes: 1 addition & 1 deletion network/src/peer_store/peer_store_impl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ impl PeerStore {
.entry(extract_peer_id(&addr).expect("connected addr should have peer id"))
{
Entry::Occupied(mut entry) => {
let mut peer = entry.get_mut();
let peer = entry.get_mut();
peer.connected_addr = addr;
peer.last_connected_at_ms = now_ms;
peer.session_type = session_type;
Expand Down
3 changes: 2 additions & 1 deletion network/src/tests/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,8 @@ mod peer_store_db;

fn random_addr() -> crate::multiaddr::Multiaddr {
format!(
"/ip4/127.0.0.1/tcp/42/p2p/{}",
"/ip4/127.0.0.1/tcp/{}/p2p/{}",
rand::random::<u16>(),
crate::PeerId::random().to_base58()
)
.parse()
Expand Down
19 changes: 19 additions & 0 deletions network/src/tests/peer_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -533,3 +533,22 @@ fn test_eviction() {
.unwrap();
assert!(peer_store.mut_addr_manager().get(&new_peer_addr).is_some());
}

#[test]
fn test_addr_unique() {
let mut peer_store = PeerStore::default();
let addr = random_addr();
let addr_1 = random_addr();

peer_store
.add_addr(addr.clone(), Flags::COMPATIBILITY)
.unwrap();
peer_store.add_addr(addr_1, Flags::COMPATIBILITY).unwrap();
assert_eq!(peer_store.addr_manager().addrs_iter().count(), 2);
assert_eq!(peer_store.fetch_addrs_to_feeler(2).len(), 2);

peer_store.add_addr(addr, Flags::COMPATIBILITY).unwrap();
assert_eq!(peer_store.fetch_addrs_to_feeler(2).len(), 2);

assert_eq!(peer_store.addr_manager().addrs_iter().count(), 2);
}

0 comments on commit 9c26137

Please sign in to comment.