leyoda/rqbit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Nothing

Browse source
This commit is contained in:
Igor Katson 2023-11-28 09:23:05 +00:00
parent e012cd94a3
commit 0478577a72
No known key found for this signature in database
GPG key ID: B4EC22B66D61A3F5
3 changed files with 31 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

40
crates/dht/src/dht.rs
View file

@ -11,12 +11,13 @@ use std::{
use crate::{
bprotocol::{
self, CompactNodeInfo, CompactPeerInfo, FindNodeRequest, GetPeersRequest, Message,
MessageKind, Node, PingRequest,
MessageKind, Node, PingRequest, Response,
},
routing_table::{InsertResult, RoutingTable},
RESPONSE_TIMEOUT,
};
use anyhow::Context;
use bencode::ByteString;
use bencode::{ByteBuf, ByteString};
use dashmap::DashMap;
use futures::{stream::FuturesUnordered, Stream, StreamExt};
use indexmap::IndexSet;
@ -53,8 +54,6 @@ pub struct DhtState {
// Created requests: (transaction_id, addr) => Requests.
// If we get a response, it gets removed from here.
//
// TODO: clean up old entries
outstanding_requests_by_transaction_id: DashMap<(u16, SocketAddr), OutstandingRequest>,
// TODO: clean up old entries
@ -112,7 +111,7 @@ impl DhtState {
spawn(
debug_span!("dht_request", tid = tid, addr = addr.to_string()),
async move {
match tokio::time::timeout(Duration::from_secs(60), rx).await {
match tokio::time::timeout(RESPONSE_TIMEOUT, rx).await {
Ok(Ok(_)) => {}
Ok(Err(e)) => {
this.outstanding_requests_by_transaction_id
@ -165,6 +164,24 @@ impl DhtState {
(transaction_id, message)
}
fn on_response(
self: &Arc<Self>,
addr: SocketAddr,
request: Request,
response: Response<ByteString>,
) -> anyhow::Result<()> {
match request {
Request::FindNode(id) => {
let nodes = response
.nodes
.ok_or_else(|| anyhow::anyhow!("expected nodes for find_node requests"))?;
self.on_found_nodes(response.id, addr, id, nodes)
}
Request::GetPeers(id) => self.on_found_peers_or_nodes(response.id, addr, id, response),
Request::Ping => Ok(()),
}
}
fn on_incoming_from_remote(
self: &Arc<Self>,
msg: Message<ByteString>,
@ -220,18 +237,7 @@ impl DhtState {
_ => unreachable!(),
};
self.routing_table.write().mark_response(&response.id);
match request {
Request::FindNode(id) => {
let nodes = response.nodes.ok_or_else(|| {
anyhow::anyhow!("expected nodes for find_node requests")
})?;
self.on_found_nodes(response.id, addr, id, nodes)
}
Request::GetPeers(id) => {
self.on_found_peers_or_nodes(response.id, addr, id, response)
}
Request::Ping => Ok(()),
}
self.on_response(addr, request, response)
}
MessageKind::PingRequest(_) => {
let message = Message {