diff --git a/p2p/src/protocol.rs b/p2p/src/protocol.rs index 0b1bf2a1..81faa67f 100644 --- a/p2p/src/protocol.rs +++ b/p2p/src/protocol.rs @@ -19,6 +19,9 @@ mod syncmgr; #[cfg(test)] mod tests; +// Futures executor. +mod executor; + use addrmgr::AddressManager; use cbfmgr::FilterManager; use invmgr::InventoryManager; diff --git a/p2p/src/protocol/cbfmgr.rs b/p2p/src/protocol/cbfmgr.rs index 89ad3ff1..d2361b01 100644 --- a/p2p/src/protocol/cbfmgr.rs +++ b/p2p/src/protocol/cbfmgr.rs @@ -20,7 +20,7 @@ use nakamoto_common::collections::{AddressBook, HashMap, HashSet}; use nakamoto_common::source; use super::filter_cache::FilterCache; -use super::output::{Disconnect, Wakeup}; +use super::output::{Blocks, Disconnect, Wakeup}; use super::{DisconnectReason, Link, PeerId, Socket}; /// Idle timeout. @@ -361,7 +361,7 @@ pub struct FilterManager { inflight: HashMap, } -impl FilterManager { +impl FilterManager { /// Create a new filter manager. pub fn new(config: Config, rng: fastrand::Rng, filters: F, upstream: U) -> Self { let peers = AddressBook::new(rng.clone()); diff --git a/p2p/src/protocol/executor.rs b/p2p/src/protocol/executor.rs new file mode 100644 index 00000000..fe0aa1c9 --- /dev/null +++ b/p2p/src/protocol/executor.rs @@ -0,0 +1,146 @@ +#![allow(dead_code)] +use std::cell::RefCell; +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::rc::Rc; +use std::sync::Arc; +use std::task::{Context, Poll}; + +struct Waker; + +impl std::task::Wake for Waker { + fn wake(self: Arc) {} + fn wake_by_ref(self: &Arc) {} +} + +type BoxFuture<'a, T> = Pin + 'a>>; + +struct Task { + future: Option>, +} + +impl fmt::Debug for Task { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Task").finish() + } +} + +#[derive(Debug, Clone)] +pub struct Request { + result: Rc>>>, +} + +impl Request { + pub fn new() -> Self { + Self { + result: Rc::new(RefCell::new(None)), + } + } + + pub fn complete(&mut self, result: Result) { + *self.result.borrow_mut() = Some(result); + } +} + +impl Future for Request { + type Output = Result; + + fn poll( + self: std::pin::Pin<&mut Self>, + _ctx: &mut std::task::Context<'_>, + ) -> std::task::Poll { + // TODO: Use `take()` instead of cloning, once you figure it out. + // For now we have to clone, as multiple futures may share the same + // refcell. + if let Some(result) = self.get_mut().result.borrow().clone() { + Poll::Ready(result) + } else { + Poll::Pending + } + } +} + +#[derive(Clone, Debug)] +pub struct Executor { + tasks: Rc>>, +} + +impl Executor { + pub fn new() -> Self { + Self { + tasks: Default::default(), + } + } + + /// Spawn a future to be executed. + pub fn spawn(&mut self, future: impl Future + 'static) { + self.tasks.borrow_mut().push(Task { + future: Some(Box::pin(future)), + }); + } + + /// Poll all tasks for completion. + pub fn poll(&mut self) -> Poll<()> { + let mut tasks = self.tasks.borrow_mut(); + let waker = Arc::new(Waker).into(); + let mut cx = Context::from_waker(&waker); + + for task in tasks.iter_mut() { + if let Some(mut fut) = task.future.take() { + if fut.as_mut().poll(&mut cx).is_pending() { + task.future = Some(fut); + } + } + } + // Clear out all completed futures. + tasks.retain(|t| t.future.is_some()); + + if tasks.is_empty() { + return Poll::Ready(()); + } + Poll::Pending + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct Random { + val: T, + } + + impl Future for Random { + type Output = T; + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + if fastrand::bool() { + Poll::Ready(self.val.clone()) + } else { + Poll::Pending + } + } + } + + #[test] + fn test_executor() { + let mut exe = Executor::new(); + + exe.spawn(async { + Random { val: 1 }.await; + }); + exe.spawn(async { + Random { val: 2 }.await; + }); + exe.spawn(async { + Random { val: 3 }.await; + }); + + loop { + if let Poll::Ready(()) = exe.poll() { + break; + } + } + } +} diff --git a/p2p/src/protocol/output.rs b/p2p/src/protocol/output.rs index c910e793..5f35e063 100644 --- a/p2p/src/protocol/output.rs +++ b/p2p/src/protocol/output.rs @@ -7,6 +7,7 @@ //! communicate with the network. use log::*; use std::cell::RefCell; +use std::collections::hash_map::Entry; use std::collections::{HashMap, VecDeque}; use std::rc::Rc; use std::sync::Arc; @@ -23,13 +24,16 @@ use nakamoto_common::bitcoin::network::message_filter::{ CFHeaders, CFilter, GetCFHeaders, GetCFilters, }; use nakamoto_common::bitcoin::network::message_network::VersionMessage; -use nakamoto_common::bitcoin::Transaction; +use nakamoto_common::bitcoin::{Block, Transaction}; use nakamoto_common::block::time::LocalDuration; use nakamoto_common::block::{BlockHash, BlockHeader, BlockTime, Height}; use crate::protocol::{Event, PeerId}; +use super::executor::Executor; +use super::executor::Request; +use super::invmgr::Inventories; use super::network::Network; use super::{addrmgr, cbfmgr, invmgr, peermgr, pingmgr, syncmgr, Locators}; @@ -157,10 +161,14 @@ pub(crate) mod message { pub struct Outbox { /// Protocol version. version: u32, + /// Futures executor. + executor: Executor, /// Output queue. outbound: Rc>>, /// Message outbox. outbox: Rc>>>, + /// Block requests to the network. + block_requests: Rc>>>, /// Network message builder. builder: message::Builder, /// Log target. @@ -172,8 +180,10 @@ impl Outbox { pub fn new(network: Network, version: u32, target: &'static str) -> Self { Self { version, + executor: Executor::new(), outbound: Rc::new(RefCell::new(VecDeque::new())), outbox: Rc::new(RefCell::new(HashMap::new())), + block_requests: Rc::new(RefCell::new(HashMap::new())), builder: message::Builder::new(network), target, } @@ -236,6 +246,17 @@ impl Outbox { } } +impl Outbox { + /// A block was received from the network. + pub fn block_received(&mut self, blk: Block) { + let block_hash = blk.block_hash(); + + if let Some(mut req) = self.block_requests.borrow_mut().remove(&block_hash) { + req.complete(Ok(blk.clone())); + } + } +} + /// Draining iterator over outbound channel queue. pub struct Drain { items: Rc>>, @@ -266,6 +287,29 @@ impl Disconnect for () { fn disconnect(&self, _addr: net::SocketAddr, _reason: DisconnectReason) {} } +pub trait Blocks { + fn get_block(&mut self, hash: BlockHash, peer: &PeerId) -> Request; +} + +impl Blocks for Outbox { + /// Fetch a block from a peer. + fn get_block(&mut self, hash: BlockHash, addr: &PeerId) -> Request { + let request = Request::new(); + + match self.block_requests.borrow_mut().entry(hash) { + Entry::Vacant(e) => { + e.insert(request.clone()); + } + Entry::Occupied(e) => { + return e.get().clone(); + } + } + self.getdata(*addr, vec![Inventory::Block(hash)]); + + request + } +} + /// The ability to be woken up in the future. pub trait Wakeup { /// Ask to be woken up in a predefined amount of time. @@ -494,6 +538,7 @@ impl cbfmgr::Events for Outbox { pub mod test { use super::*; use nakamoto_common::bitcoin::network::message::{NetworkMessage, RawNetworkMessage}; + use nakamoto_test::assert_matches; pub fn messages( channel: &mut Outbox, @@ -511,4 +556,41 @@ pub mod test { } msgs.into_iter() } + + #[test] + fn test_get_block() { + let network = Network::Mainnet; + let remote: net::SocketAddr = ([88, 88, 88, 88], 8333).into(); + let mut outbox = Outbox::new(network, crate::protocol::PROTOCOL_VERSION, "test"); + + outbox.executor.spawn({ + let mut outbox = outbox.clone(); + + async move { + outbox + .get_block(network.genesis_hash(), &remote) + .await + .unwrap(); + } + }); + outbox.executor.spawn({ + let mut outbox = outbox.clone(); + + async move { + outbox + .get_block(network.genesis_hash(), &remote) + .await + .unwrap(); + } + }); + assert!(outbox.executor.poll().is_pending()); + + assert_matches!( + messages(&mut outbox, &remote).next(), + Some(NetworkMessage::GetData(_)) + ); + + outbox.block_received(network.genesis_block()); + assert!(outbox.executor.poll().is_ready()); + } }