use std::fmt::{Debug, Display};
use std::time::Instant;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use futures::channel::mpsc;
use futures::prelude::*;
use crate::common::*;
use crate::store::Datastore;
use crate::table::NodeTable;
use super::{Dht, Operation, OperationKind, OperationState, RequestState};
pub struct ConnectFuture<Id, Info> {
done: bool,
rx: mpsc::Receiver<Result<Vec<Entry<Id, Info>>, Error>>,
}
impl<Id, Info> Future for ConnectFuture<Id, Info> {
type Output = Result<Vec<Entry<Id, Info>>, Error>;
fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
if self.done {
return Poll::Pending;
}
match self.rx.poll_next_unpin(ctx) {
Poll::Ready(Some(r)) => {
self.done = true;
Poll::Ready(r)
}
_ => Poll::Pending,
}
}
}
impl<Id, Info, Data, ReqId, Table, Store> Dht<Id, Info, Data, ReqId, Table, Store>
where
Id: DatabaseId + Clone + Sized + Send + 'static,
Info: PartialEq + Clone + Sized + Debug + Send + 'static,
Data: PartialEq + Clone + Sized + Debug + Send + 'static,
ReqId: RequestId + Clone + Sized + Display + Debug + Send + 'static,
Table: NodeTable<Id, Info> + Send + 'static,
Store: Datastore<Id, Data> + Send + 'static,
{
pub fn connect(
&mut self,
peers: &[Entry<Id, Info>],
) -> Result<(ConnectFuture<Id, Info>, ReqId), Error> {
let req_id = ReqId::generate();
let (done_tx, done_rx) = mpsc::channel(1);
let mut op = Operation::new(
req_id.clone(),
self.id.clone(),
OperationKind::Connect(done_tx),
);
for e in peers {
op.nodes
.insert(e.id().clone(), (e.clone(), RequestState::Active));
}
self.operations.insert(req_id.clone(), op);
Ok((
ConnectFuture {
done: false,
rx: done_rx,
},
req_id,
))
}
pub fn connect_start(
&mut self,
) -> Result<(ConnectFuture<Id, Info>, ReqId, Request<Id, Data>), Error> {
let req_id = ReqId::generate();
let (done_tx, done_rx) = mpsc::channel(1);
let mut op = Operation::new(
req_id.clone(),
self.id.clone(),
OperationKind::Connect(done_tx),
);
let req = Request::FindNode(self.id.clone());
op.state = OperationState::Connecting;
op.last_update = Instant::now();
self.operations.insert(req_id.clone(), op);
Ok((
ConnectFuture {
done: false,
rx: done_rx,
},
req_id,
req,
))
}
}
#[cfg(test)]
mod tests {
use log::*;
use simplelog::{Config as LogConfig, LevelFilter, SimpleLogger};
use crate::{Config, Dht};
use super::*;
#[async_std::test]
async fn test_connect() {
let _ = SimpleLogger::init(LevelFilter::Debug, LogConfig::default());
let n1 = Entry::new([0b1000], 100);
let n2 = Entry::new([0b0011], 200);
let n3 = Entry::new([0b0010], 300);
let n4 = Entry::new([0b1001], 400);
let n5 = Entry::new([0b1010], 400);
let mut config = Config::default();
config.concurrency = 2;
config.k = 2;
let (tx, mut rx) = mpsc::channel(10);
let mut dht: Dht<_, u32, u32, u16> = Dht::standard(n1.id().clone(), config, tx);
info!("Start connect");
let (connect, req_id) = dht
.connect(&[n2.clone(), n3.clone()])
.expect("Error starting lookup");
info!("Search round 0");
dht.update().unwrap();
assert_eq!(
rx.try_next().unwrap(),
Some((req_id, n3.clone(), Request::FindNode(n1.id().clone())))
);
assert_eq!(
rx.try_next().unwrap(),
Some((req_id, n2.clone(), Request::FindNode(n1.id().clone())))
);
dht.handle_resp(
req_id,
&n2,
&Response::NodesFound(n1.id().clone(), vec![n4.clone()]),
)
.unwrap();
dht.handle_resp(
req_id,
&n3,
&Response::NodesFound(n1.id().clone(), vec![n5.clone()]),
)
.unwrap();
info!("Search round 1");
dht.update().unwrap();
dht.update().unwrap();
assert_eq!(
rx.try_next().unwrap(),
Some((req_id, n4.clone(), Request::FindNode(n1.id().clone())))
);
assert_eq!(
rx.try_next().unwrap(),
Some((req_id, n5.clone(), Request::FindNode(n1.id().clone())))
);
dht.handle_resp(
req_id,
&n4,
&Response::NodesFound(n1.id().clone(), vec![n4.clone()]),
)
.unwrap();
dht.handle_resp(
req_id,
&n5,
&Response::NodesFound(n1.id().clone(), vec![n5.clone()]),
)
.unwrap();
dht.update().unwrap();
dht.update().unwrap();
dht.update().unwrap();
info!("Expecting completion");
assert_eq!(connect.await, Ok(vec![n4, n5, n3, n2]));
}
}