Restructure proto
The existing code has been moved out and is being copied back piece / by piece while restructuring the code to (hopefully) be more manageable.
This commit is contained in:
Carl Lerche
GitHub
@@ -1,362 +1,60 @@
|
||||
use ConnectionError;
|
||||
use frame::{Frame, Ping, SettingSet};
|
||||
use frame::Ping;
|
||||
use proto::*;
|
||||
|
||||
/// Acknowledges ping requests from the remote.
|
||||
#[derive(Debug)]
|
||||
pub struct PingPong<T, U> {
|
||||
inner: T,
|
||||
sending_pong: Option<Frame<U>>,
|
||||
pub struct PingPong<B> {
|
||||
// TODO: this doesn't need to save the entire frame
|
||||
sending_pong: Option<Frame<B>>,
|
||||
received_pong: Option<PingPayload>,
|
||||
// TODO: factor this out
|
||||
blocked_ping: Option<task::Task>,
|
||||
expecting_pong: bool,
|
||||
}
|
||||
|
||||
impl<T, U> PingPong<T, U>
|
||||
where T: Stream<Item = Frame, Error = ConnectionError>,
|
||||
T: Sink<SinkItem = Frame<U>, SinkError = ConnectionError>,
|
||||
impl<B> PingPong<B>
|
||||
where B: Buf,
|
||||
{
|
||||
pub fn new(inner: T) -> Self {
|
||||
pub fn new() -> Self {
|
||||
PingPong {
|
||||
inner,
|
||||
sending_pong: None,
|
||||
received_pong: None,
|
||||
expecting_pong: false,
|
||||
blocked_ping: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U> ControlPing for PingPong<T, U>
|
||||
where T: Sink<SinkItem = Frame<U>, SinkError = ConnectionError>,
|
||||
T: ReadySink,
|
||||
{
|
||||
fn start_ping(&mut self, body: PingPayload) -> StartSend<PingPayload, ConnectionError> {
|
||||
if self.inner.poll_ready()?.is_not_ready() {
|
||||
return Ok(AsyncSink::NotReady(body));
|
||||
}
|
||||
/// Process a ping
|
||||
pub fn recv_ping(&mut self, ping: Ping) {
|
||||
// The caller should always check that `send_pongs` returns ready before
|
||||
// calling `recv_ping`.
|
||||
assert!(self.sending_pong.is_none());
|
||||
|
||||
// Only allow one in-flight ping.
|
||||
if self.expecting_pong || self.received_pong.is_some() {
|
||||
self.blocked_ping = Some(task::current());
|
||||
return Ok(AsyncSink::NotReady(body))
|
||||
}
|
||||
if ping.is_ack() {
|
||||
// Save acknowledgements to be returned from take_pong().
|
||||
self.received_pong = Some(ping.into_payload());
|
||||
|
||||
match self.inner.start_send(Ping::ping(body).into())? {
|
||||
AsyncSink::NotReady(_) => {
|
||||
// By virtual of calling inner.poll_ready(), this must not happen.
|
||||
unreachable!()
|
||||
}
|
||||
AsyncSink::Ready => {
|
||||
self.expecting_pong = true;
|
||||
Ok(AsyncSink::Ready)
|
||||
if let Some(task) = self.blocked_ping.take() {
|
||||
task.notify();
|
||||
}
|
||||
} else {
|
||||
// Save the ping's payload to be sent as an acknowledgement.
|
||||
let pong = Ping::pong(ping.into_payload());
|
||||
self.sending_pong = Some(pong.into());
|
||||
}
|
||||
}
|
||||
|
||||
fn take_pong(&mut self) -> Option<PingPayload> {
|
||||
match self.received_pong.take() {
|
||||
None => None,
|
||||
Some(p) => {
|
||||
self.expecting_pong = false;
|
||||
if let Some(task) = self.blocked_ping.take() {
|
||||
task.notify();
|
||||
}
|
||||
Some(p)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U> PingPong<T, U>
|
||||
where T: Sink<SinkItem = Frame<U>, SinkError = ConnectionError>,
|
||||
{
|
||||
fn try_send_pong(&mut self) -> Poll<(), ConnectionError> {
|
||||
/// Send any pending pongs.
|
||||
pub fn send_pending_pong<T>(&mut self, dst: &mut Codec<T, B>) -> Poll<(), ConnectionError>
|
||||
where T: AsyncWrite,
|
||||
{
|
||||
if let Some(pong) = self.sending_pong.take() {
|
||||
if let AsyncSink::NotReady(pong) = self.inner.start_send(pong)? {
|
||||
if let AsyncSink::NotReady(pong) = dst.start_send(pong)? {
|
||||
// If the pong can't be sent, save it.
|
||||
self.sending_pong = Some(pong);
|
||||
return Ok(Async::NotReady);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Async::Ready(()))
|
||||
}
|
||||
}
|
||||
|
||||
/// > Receivers of a PING frame that does not include an ACK flag MUST send
|
||||
/// > a PING frame with the ACK flag set in response, with an identical
|
||||
/// > payload. PING responses SHOULD be given higher priority than any
|
||||
/// > other frame.
|
||||
impl<T, U> Stream for PingPong<T, U>
|
||||
where T: Stream<Item = Frame, Error = ConnectionError>,
|
||||
T: Sink<SinkItem = Frame<U>, SinkError = ConnectionError>,
|
||||
{
|
||||
type Item = Frame;
|
||||
type Error = ConnectionError;
|
||||
|
||||
/// Reads the next frame from the underlying socket, eliding ping requests.
|
||||
///
|
||||
/// If a PING is received without the ACK flag, the frame is sent to the remote with
|
||||
/// its ACK flag set.
|
||||
fn poll(&mut self) -> Poll<Option<Frame>, ConnectionError> {
|
||||
loop {
|
||||
// Don't read any frames until `inner` accepts any pending pong.
|
||||
try_ready!(self.try_send_pong());
|
||||
|
||||
match self.inner.poll()? {
|
||||
Async::Ready(Some(Frame::Ping(ping))) => {
|
||||
if ping.is_ack() {
|
||||
// Save acknowledgements to be returned from take_pong().
|
||||
self.received_pong = Some(ping.into_payload());
|
||||
if let Some(task) = self.blocked_ping.take() {
|
||||
task.notify();
|
||||
}
|
||||
} else {
|
||||
// Save the ping's payload to be sent as an acknowledgement.
|
||||
let pong = Ping::pong(ping.into_payload());
|
||||
self.sending_pong = Some(pong.into());
|
||||
}
|
||||
}
|
||||
|
||||
// Everything other than ping gets passed through.
|
||||
f => return Ok(f),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U> Sink for PingPong<T, U>
|
||||
where T: Sink<SinkItem = Frame<U>, SinkError = ConnectionError>,
|
||||
{
|
||||
type SinkItem = Frame<U>;
|
||||
type SinkError = ConnectionError;
|
||||
|
||||
fn start_send(&mut self, item: Self::SinkItem)
|
||||
-> StartSend<Self::SinkItem, Self::SinkError>
|
||||
{
|
||||
// Pings _SHOULD_ have priority over other messages, so attempt to send pending
|
||||
// ping frames before attempting to send `item`.
|
||||
if self.try_send_pong()?.is_not_ready() {
|
||||
return Ok(AsyncSink::NotReady(item));
|
||||
}
|
||||
|
||||
self.inner.start_send(item)
|
||||
}
|
||||
|
||||
/// Polls the underlying sink and tries to flush pending pong frames.
|
||||
fn poll_complete(&mut self) -> Poll<(), ConnectionError> {
|
||||
try_ready!(self.try_send_pong());
|
||||
self.inner.poll_complete()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, U> ReadySink for PingPong<T, U>
|
||||
where T: Sink<SinkItem = Frame<U>, SinkError = ConnectionError>,
|
||||
T: ReadySink,
|
||||
{
|
||||
fn poll_ready(&mut self) -> Poll<(), ConnectionError> {
|
||||
try_ready!(self.try_send_pong());
|
||||
self.inner.poll_ready()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: ApplySettings, U> ApplySettings for PingPong<T, U> {
|
||||
fn apply_local_settings(&mut self, set: &SettingSet) -> Result<(), ConnectionError> {
|
||||
self.inner.apply_local_settings(set)
|
||||
}
|
||||
|
||||
fn apply_remote_settings(&mut self, set: &SettingSet) -> Result<(), ConnectionError> {
|
||||
self.inner.apply_remote_settings(set)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
use proto::ControlPing;
|
||||
use std::cell::RefCell;
|
||||
use std::collections::VecDeque;
|
||||
use std::rc::Rc;
|
||||
|
||||
#[test]
|
||||
fn responds_to_ping_with_pong() {
|
||||
let trans = Transport::default();
|
||||
let mut ping_pong = PingPong::new(trans.clone());
|
||||
|
||||
{
|
||||
let mut trans = trans.0.borrow_mut();
|
||||
let ping = Ping::ping(*b"buoyant_");
|
||||
trans.from_socket.push_back(ping.into());
|
||||
}
|
||||
|
||||
match ping_pong.poll() {
|
||||
Ok(Async::NotReady) => {} // cool
|
||||
rsp => panic!("unexpected poll result: {:?}", rsp),
|
||||
}
|
||||
|
||||
{
|
||||
let mut trans = trans.0.borrow_mut();
|
||||
assert_eq!(trans.to_socket.len(), 1);
|
||||
match trans.to_socket.pop_front().unwrap() {
|
||||
Frame::Ping(pong) => {
|
||||
assert!(pong.is_ack());
|
||||
assert_eq!(&pong.into_payload(), b"buoyant_");
|
||||
}
|
||||
f => panic!("unexpected frame: {:?}", f),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responds_to_ping_even_when_blocked() {
|
||||
let trans = Transport::default();
|
||||
let mut ping_pong = PingPong::new(trans.clone());
|
||||
|
||||
// Configure the transport so that writes can't proceed.
|
||||
{
|
||||
let mut trans = trans.0.borrow_mut();
|
||||
trans.start_send_blocked = true;
|
||||
}
|
||||
|
||||
// The transport receives a ping but can't send it immediately.
|
||||
{
|
||||
let mut trans = trans.0.borrow_mut();
|
||||
let ping = Ping::ping(*b"buoyant?");
|
||||
trans.from_socket.push_back(ping.into());
|
||||
}
|
||||
assert!(ping_pong.poll().unwrap().is_not_ready());
|
||||
|
||||
// The transport receives another ping but can't send it immediately.
|
||||
{
|
||||
let mut trans = trans.0.borrow_mut();
|
||||
let ping = Ping::ping(*b"buoyant!");
|
||||
trans.from_socket.push_back(ping.into());
|
||||
}
|
||||
assert!(ping_pong.poll().unwrap().is_not_ready());
|
||||
|
||||
// At this point, ping_pong is holding two pongs that it cannot send.
|
||||
{
|
||||
let mut trans = trans.0.borrow_mut();
|
||||
assert!(trans.to_socket.is_empty());
|
||||
|
||||
trans.start_send_blocked = false;
|
||||
}
|
||||
|
||||
// Now that start_send_blocked is disabled, the next poll will successfully send
|
||||
// the pongs on the transport.
|
||||
assert!(ping_pong.poll().unwrap().is_not_ready());
|
||||
{
|
||||
let mut trans = trans.0.borrow_mut();
|
||||
assert_eq!(trans.to_socket.len(), 2);
|
||||
match trans.to_socket.pop_front().unwrap() {
|
||||
Frame::Ping(pong) => {
|
||||
assert!(pong.is_ack());
|
||||
assert_eq!(&pong.into_payload(), b"buoyant?");
|
||||
}
|
||||
f => panic!("unexpected frame: {:?}", f),
|
||||
}
|
||||
match trans.to_socket.pop_front().unwrap() {
|
||||
Frame::Ping(pong) => {
|
||||
assert!(pong.is_ack());
|
||||
assert_eq!(&pong.into_payload(), b"buoyant!");
|
||||
}
|
||||
f => panic!("unexpected frame: {:?}", f),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pong_passes_through() {
|
||||
let trans = Transport::default();
|
||||
let mut ping_pong = PingPong::new(trans.clone());
|
||||
|
||||
{
|
||||
let mut trans = trans.0.borrow_mut();
|
||||
let pong = Ping::pong(*b"buoyant!");
|
||||
trans.from_socket.push_back(pong.into());
|
||||
}
|
||||
|
||||
assert!(ping_pong.poll().unwrap().is_not_ready());
|
||||
match ping_pong.take_pong() {
|
||||
Some(pong) => assert_eq!(&pong, b"buoyant!"),
|
||||
None => panic!("no pong received"),
|
||||
}
|
||||
|
||||
{
|
||||
let trans = trans.0.borrow();
|
||||
assert_eq!(trans.to_socket.len(), 0);
|
||||
}
|
||||
}
|
||||
|
||||
/// A stubbed transport for tests.a
|
||||
///
|
||||
/// We probably want/have something generic for this?
|
||||
#[derive(Clone, Default)]
|
||||
struct Transport(Rc<RefCell<Inner>>);
|
||||
|
||||
#[derive(Default)]
|
||||
struct Inner {
|
||||
from_socket: VecDeque<Frame>,
|
||||
to_socket: VecDeque<Frame>,
|
||||
read_blocked: bool,
|
||||
start_send_blocked: bool,
|
||||
closing: bool,
|
||||
}
|
||||
|
||||
impl Stream for Transport {
|
||||
type Item = Frame;
|
||||
type Error = ConnectionError;
|
||||
|
||||
fn poll(&mut self) -> Poll<Option<Frame>, ConnectionError> {
|
||||
let mut trans = self.0.borrow_mut();
|
||||
if trans.read_blocked || (!trans.closing && trans.from_socket.is_empty()) {
|
||||
Ok(Async::NotReady)
|
||||
} else {
|
||||
Ok(trans.from_socket.pop_front().into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Sink for Transport {
|
||||
type SinkItem = Frame;
|
||||
type SinkError = ConnectionError;
|
||||
|
||||
fn start_send(&mut self, item: Frame) -> StartSend<Frame, ConnectionError> {
|
||||
let mut trans = self.0.borrow_mut();
|
||||
if trans.closing || trans.start_send_blocked {
|
||||
Ok(AsyncSink::NotReady(item))
|
||||
} else {
|
||||
trans.to_socket.push_back(item);
|
||||
Ok(AsyncSink::Ready)
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_complete(&mut self) -> Poll<(), ConnectionError> {
|
||||
let trans = self.0.borrow();
|
||||
if !trans.to_socket.is_empty() {
|
||||
Ok(Async::NotReady)
|
||||
} else {
|
||||
Ok(Async::Ready(()))
|
||||
}
|
||||
}
|
||||
|
||||
fn close(&mut self) -> Poll<(), ConnectionError> {
|
||||
{
|
||||
let mut trans = self.0.borrow_mut();
|
||||
trans.closing = true;
|
||||
}
|
||||
self.poll_complete()
|
||||
}
|
||||
}
|
||||
|
||||
impl ReadySink for Transport {
|
||||
fn poll_ready(&mut self) -> Poll<(), ConnectionError> {
|
||||
let trans = self.0.borrow();
|
||||
if trans.closing || trans.start_send_blocked {
|
||||
Ok(Async::NotReady)
|
||||
} else {
|
||||
Ok(Async::Ready(()))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user