<commit_before>use std::io::{IoError, MemReader};
use std::fmt;
use std::fmt::{Show, Formatter, FormatError};
use std::collections::HashMap;

use conduit::Response;
use curl::ErrCode;
use pg::error::{PostgresError, PostgresConnectError};
use serialize::json;
use git2;

use util::json_response;

pub trait CargoError: Send {
    fn description(&self) -> String;
    fn detail(&self) -> Option<String> { None }
    fn cause<'a>(&'a self) -> Option<&'a CargoError + Send> { None }

    fn to_error<E: FromError<Self>>(self) -> E {
        FromError::from_error(self)
    }

    fn box_error(self) -> Box<CargoError + Send> {
        box self as Box<CargoError + Send>
    }

    fn concrete(&self) -> ConcreteCargoError {
        ConcreteCargoError {
            description: self.description(),
            detail: self.detail(),
            cause: self.cause().map(|c| box c.concrete() as Box<CargoError + Send>),
            human: false,
        }
    }

    fn with_cause<E: CargoError + Send>(self, cause: E) -> Box<CargoError + Send> {
        let mut concrete = self.concrete();
        concrete.cause = Some(cause.box_error());
        box concrete as Box<CargoError + Send>
    }

    fn response(&self) -> Option<Response> {
        #[deriving(Encodable)] struct Error { detail: String }
        #[deriving(Encodable)] struct Bad { errors: Vec<Error> }

        if self.human() {
            Some(json_response(&Bad {
                errors: vec![Error { detail: self.description() }]
            }))
        } else {
            self.cause().and_then(|cause| cause.response())
        }
    }
    fn human(&self) -> bool { false }
}

pub trait FromError<E> {
    fn from_error(error: E) -> Self;
}

impl<E: CargoError + Send> FromError<E> for Box<CargoError + Send> {
    fn from_error(error: E) -> Box<CargoError + Send> {
        error.box_error()
    }
}

macro_rules! from_error (
    ($ty:ty) => {
        impl FromError<$ty> for $ty {
            fn from_error(error: $ty) -> $ty {
                error
            }
        }
    }
)

impl<'a> Show for &'a CargoError + Send {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        try!(write!(f, "{}", self.description()));

        match self.detail() {
            Some(s) => try!(write!(f, "\n  {}", s)),
            None => {}
        }

        match self.cause() {
            Some(cause) => try!(write!(f, "\nCaused by: {}", cause)),
            None => {}
        }

        Ok(())
    }
}

impl Show for Box<CargoError + Send> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let me: &CargoError + Send = &**self;
        me.fmt(f)
    }
}

impl CargoError for Box<CargoError + Send> {
    fn description(&self) -> String { (**self).description() }
    fn detail(&self) -> Option<String> { (**self).detail() }
    fn cause<'a>(&'a self) -> Option<&'a CargoError + Send> { (**self).cause() }
    fn box_error(self) -> Box<CargoError + Send> { self }
    fn human(&self) -> bool { (**self).human() }
    fn response(&self) -> Option<Response> { (**self).response() }
}

pub type CargoResult<T> = Result<T, Box<CargoError + Send>>;

pub trait BoxError<T> {
    fn box_error(self) -> CargoResult<T>;
}

pub trait ChainError<T> {
    fn chain_error<E: CargoError + Send>(self, callback: || -> E) -> CargoResult<T> ;
}

impl<'a, T> ChainError<T> for ||:'a -> CargoResult<T> {
    fn chain_error<E: CargoError + Send>(self, callback: || -> E) -> CargoResult<T> {
        self().map_err(|err| callback().with_cause(err))
    }
}

impl<T, E: CargoError + Send> BoxError<T> for Result<T, E> {
    fn box_error(self) -> CargoResult<T> {
        self.map_err(|err| err.box_error())
    }
}

impl<T, E: CargoError + Send> ChainError<T> for Result<T, E> {
    fn chain_error<E: CargoError + Send>(self, callback: || -> E) -> CargoResult<T>  {
        self.map_err(|err| callback().with_cause(err))
    }
}

impl CargoError for IoError {
    fn description(&self) -> String { self.to_string() }
}

from_error!(IoError)

impl CargoError for FormatError {
    fn description(&self) -> String {
        "formatting failed".to_string()
    }
}

from_error!(FormatError)

impl CargoError for PostgresError {
    fn description(&self) -> String { self.to_string() }
}

from_error!(PostgresError)

impl CargoError for PostgresConnectError {
    fn description(&self) -> String { self.to_string() }
}

from_error!(PostgresConnectError)

impl CargoError for ErrCode {
    fn description(&self) -> String { self.to_string() }
}

from_error!(ErrCode)

impl CargoError for json::DecoderError {
    fn description(&self) -> String { self.to_string() }
}

from_error!(json::DecoderError)

impl CargoError for git2::Error {
    fn description(&self) -> String { self.to_string() }
}

from_error!(git2::Error)

impl<T: CargoError + Send> FromError<T> for Box<Show + 'static> {
    fn from_error(t: T) -> Box<Show + 'static> {
        box() (box t as Box<CargoError + Send>) as Box<Show + 'static>
    }
}

pub struct ConcreteCargoError {
    description: String,
    detail: Option<String>,
    cause: Option<Box<CargoError + Send>>,
    human: bool,
}

impl Show for ConcreteCargoError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{}", self.description)
    }
}

impl CargoError for ConcreteCargoError {
    fn description(&self) -> String {
        self.description.clone()
    }

    fn detail(&self) -> Option<String> {
        self.detail.clone()
    }

    fn cause<'a>(&'a self) -> Option<&'a CargoError + Send> {
        self.cause.as_ref().map(|c| { let err: &CargoError + Send = &**c; err })
    }

    fn with_cause<E: CargoError + Send>(mut self,
                                        err: E) -> Box<CargoError + Send> {
        self.cause = Some(err.box_error());
        box self as Box<CargoError + Send>
    }
    fn human(&self) -> bool { self.human }
}

pub struct NotFound;

impl CargoError for NotFound {
    fn description(&self) -> String { "not found".to_string() }

    fn response(&self) -> Option<Response> {
        Some(Response {
            status: (404, "Not Found"),
            headers: HashMap::new(),
            body: box MemReader::new(Vec::new()),
        })
    }
}

from_error!(NotFound)

pub struct Unauthorized;

impl CargoError for Unauthorized {
    fn description(&self) -> String { "unauthorized".to_string() }

    fn response(&self) -> Option<Response> {
        Some(Response {
            status: (403, "Forbidden"),
            headers: HashMap::new(),
            body: box MemReader::new(Vec::new()),
        })
    }
}

from_error!(Unauthorized)

pub fn internal_error<S1: Str, S2: Str>(error: S1,
                                        detail: S2) -> Box<CargoError + Send> {
    box ConcreteCargoError {
        description: error.as_slice().to_string(),
        detail: Some(detail.as_slice().to_string()),
        cause: None,
        human: false,
    } as Box<CargoError + Send>
}

pub fn internal<S: Show>(error: S) -> Box<CargoError + Send> {
    box ConcreteCargoError {
        description: error.to_string(),
        detail: None,
        cause: None,
        human: false,
    } as Box<CargoError + Send>
}

pub fn human<S: Show>(error: S) -> Box<CargoError + Send> {
    box ConcreteCargoError {
        description: error.to_string(),
        detail: None,
        cause: None,
        human: true,
    } as Box<CargoError + Send>
}
<commit_msg>Send some json back on a 403 unauthorized<commit_after>use std::io::{IoError, MemReader};
use std::fmt;
use std::fmt::{Show, Formatter, FormatError};
use std::collections::HashMap;

use conduit::Response;
use curl::ErrCode;
use pg::error::{PostgresError, PostgresConnectError};
use serialize::json;
use git2;

use util::json_response;

#[deriving(Encodable)] struct Error { detail: String }
#[deriving(Encodable)] struct Bad { errors: Vec<Error> }

pub trait CargoError: Send {
    fn description(&self) -> String;
    fn detail(&self) -> Option<String> { None }
    fn cause<'a>(&'a self) -> Option<&'a CargoError + Send> { None }

    fn to_error<E: FromError<Self>>(self) -> E {
        FromError::from_error(self)
    }

    fn box_error(self) -> Box<CargoError + Send> {
        box self as Box<CargoError + Send>
    }

    fn concrete(&self) -> ConcreteCargoError {
        ConcreteCargoError {
            description: self.description(),
            detail: self.detail(),
            cause: self.cause().map(|c| box c.concrete() as Box<CargoError + Send>),
            human: false,
        }
    }

    fn with_cause<E: CargoError + Send>(self, cause: E) -> Box<CargoError + Send> {
        let mut concrete = self.concrete();
        concrete.cause = Some(cause.box_error());
        box concrete as Box<CargoError + Send>
    }

    fn response(&self) -> Option<Response> {
        if self.human() {
            Some(json_response(&Bad {
                errors: vec![Error { detail: self.description() }]
            }))
        } else {
            self.cause().and_then(|cause| cause.response())
        }
    }
    fn human(&self) -> bool { false }
}

pub trait FromError<E> {
    fn from_error(error: E) -> Self;
}

impl<E: CargoError + Send> FromError<E> for Box<CargoError + Send> {
    fn from_error(error: E) -> Box<CargoError + Send> {
        error.box_error()
    }
}

macro_rules! from_error (
    ($ty:ty) => {
        impl FromError<$ty> for $ty {
            fn from_error(error: $ty) -> $ty {
                error
            }
        }
    }
)

impl<'a> Show for &'a CargoError + Send {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        try!(write!(f, "{}", self.description()));

        match self.detail() {
            Some(s) => try!(write!(f, "\n  {}", s)),
            None => {}
        }

        match self.cause() {
            Some(cause) => try!(write!(f, "\nCaused by: {}", cause)),
            None => {}
        }

        Ok(())
    }
}

impl Show for Box<CargoError + Send> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let me: &CargoError + Send = &**self;
        me.fmt(f)
    }
}

impl CargoError for Box<CargoError + Send> {
    fn description(&self) -> String { (**self).description() }
    fn detail(&self) -> Option<String> { (**self).detail() }
    fn cause<'a>(&'a self) -> Option<&'a CargoError + Send> { (**self).cause() }
    fn box_error(self) -> Box<CargoError + Send> { self }
    fn human(&self) -> bool { (**self).human() }
    fn response(&self) -> Option<Response> { (**self).response() }
}

pub type CargoResult<T> = Result<T, Box<CargoError + Send>>;

pub trait BoxError<T> {
    fn box_error(self) -> CargoResult<T>;
}

pub trait ChainError<T> {
    fn chain_error<E: CargoError + Send>(self, callback: || -> E) -> CargoResult<T> ;
}

impl<'a, T> ChainError<T> for ||:'a -> CargoResult<T> {
    fn chain_error<E: CargoError + Send>(self, callback: || -> E) -> CargoResult<T> {
        self().map_err(|err| callback().with_cause(err))
    }
}

impl<T, E: CargoError + Send> BoxError<T> for Result<T, E> {
    fn box_error(self) -> CargoResult<T> {
        self.map_err(|err| err.box_error())
    }
}

impl<T, E: CargoError + Send> ChainError<T> for Result<T, E> {
    fn chain_error<E: CargoError + Send>(self, callback: || -> E) -> CargoResult<T>  {
        self.map_err(|err| callback().with_cause(err))
    }
}

impl CargoError for IoError {
    fn description(&self) -> String { self.to_string() }
}

from_error!(IoError)

impl CargoError for FormatError {
    fn description(&self) -> String {
        "formatting failed".to_string()
    }
}

from_error!(FormatError)

impl CargoError for PostgresError {
    fn description(&self) -> String { self.to_string() }
}

from_error!(PostgresError)

impl CargoError for PostgresConnectError {
    fn description(&self) -> String { self.to_string() }
}

from_error!(PostgresConnectError)

impl CargoError for ErrCode {
    fn description(&self) -> String { self.to_string() }
}

from_error!(ErrCode)

impl CargoError for json::DecoderError {
    fn description(&self) -> String { self.to_string() }
}

from_error!(json::DecoderError)

impl CargoError for git2::Error {
    fn description(&self) -> String { self.to_string() }
}

from_error!(git2::Error)

impl<T: CargoError + Send> FromError<T> for Box<Show + 'static> {
    fn from_error(t: T) -> Box<Show + 'static> {
        box() (box t as Box<CargoError + Send>) as Box<Show + 'static>
    }
}

pub struct ConcreteCargoError {
    description: String,
    detail: Option<String>,
    cause: Option<Box<CargoError + Send>>,
    human: bool,
}

impl Show for ConcreteCargoError {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "{}", self.description)
    }
}

impl CargoError for ConcreteCargoError {
    fn description(&self) -> String {
        self.description.clone()
    }

    fn detail(&self) -> Option<String> {
        self.detail.clone()
    }

    fn cause<'a>(&'a self) -> Option<&'a CargoError + Send> {
        self.cause.as_ref().map(|c| { let err: &CargoError + Send = &**c; err })
    }

    fn with_cause<E: CargoError + Send>(mut self,
                                        err: E) -> Box<CargoError + Send> {
        self.cause = Some(err.box_error());
        box self as Box<CargoError + Send>
    }
    fn human(&self) -> bool { self.human }
}

pub struct NotFound;

impl CargoError for NotFound {
    fn description(&self) -> String { "not found".to_string() }

    fn response(&self) -> Option<Response> {
        Some(Response {
            status: (404, "Not Found"),
            headers: HashMap::new(),
            body: box MemReader::new(Vec::new()),
        })
    }
}

from_error!(NotFound)

pub struct Unauthorized;

impl CargoError for Unauthorized {
    fn description(&self) -> String { "unauthorized".to_string() }

    fn response(&self) -> Option<Response> {
        let mut response = json_response(&Bad {
            errors: vec![Error {
                detail: "must be logged in to perform that action".to_string(),
            }],
        });
        response.status = (403, "Forbidden");
        return Some(response);
    }
}

from_error!(Unauthorized)

pub fn internal_error<S1: Str, S2: Str>(error: S1,
                                        detail: S2) -> Box<CargoError + Send> {
    box ConcreteCargoError {
        description: error.as_slice().to_string(),
        detail: Some(detail.as_slice().to_string()),
        cause: None,
        human: false,
    } as Box<CargoError + Send>
}

pub fn internal<S: Show>(error: S) -> Box<CargoError + Send> {
    box ConcreteCargoError {
        description: error.to_string(),
        detail: None,
        cause: None,
        human: false,
    } as Box<CargoError + Send>
}

pub fn human<S: Show>(error: S) -> Box<CargoError + Send> {
    box ConcreteCargoError {
        description: error.to_string(),
        detail: None,
        cause: None,
        human: true,
    } as Box<CargoError + Send>
}
<|endoftext|>


<commit_before>use std::cmp::max;
use std::any::Any;

use ncurses;

use color;
use ::{Cursive};
use event::EventResult;
use view::{View,SizeRequest,DimensionRequest,Selector};
use view::{Button,SizedView};
use vec::{Vec2,Vec4,ToVec4};
use printer::Printer;

#[derive(PartialEq)]
enum Focus {
    Content,
    Button(usize),
}

/// Popup-like view with a main content, and optional buttons under it.
///
/// # Examples
///
/// ```
/// let dialog = Dialog::new(TextView::new("Hello!")).button("Ok", |s,_| s.quit());
/// ```
pub struct Dialog {
    title: String,
    content: Box<View>,

    buttons: Vec<SizedView<Button>>,

    padding: Vec4,
    borders: Vec4,

    focus: Focus,
}

impl Dialog {
    /// Creates a new Dialog with the given content.
    pub fn new<V: View + 'static>(view: V) -> Self {
        Dialog {
            content: Box::new(view),
            buttons: Vec::new(),
            title: String::new(),
            focus: Focus::Content,
            padding: Vec4::new(1,1,0,0),
            borders: Vec4::new(1,1,1,1),
        }
    }

    /// Adds a button to the dialog with the given label and callback.
    ///
    /// Consumes and returns self for easy chaining.
    pub fn button<'a, F>(mut self, label: &'a str, cb: F) -> Self
        where F: Fn(&mut Cursive) + 'static
    {
        self.buttons.push(SizedView::new(Button::new(label, cb)));

        self
    }

    /// Shortcut method to add a button that will dismiss the dialog.
    pub fn dismiss_button<'a>(self, label: &'a str) -> Self {
        self.button(label, |s| s.screen_mut().pop_layer())
    }

    /// Sets the title of the dialog.
    /// If not empty, it will be visible at the top.
    pub fn title(mut self, label: &str) -> Self {
        self.title = label.to_string();
        self
    }

    /// Sets the padding in the dialog (around content and buttons).
    pub fn padding<T: ToVec4>(mut self, padding: T) -> Self {
        self.padding = padding.to_vec4();

        self
    }

}

impl View for Dialog {
    fn draw(&mut self, printer: &Printer, focused: bool) {

        // This will be the height used by the buttons.
        let mut height = 0;
        // Current horizontal position of the next button we'll draw.
        let mut x = 0;
        for (i,button) in self.buttons.iter_mut().enumerate().rev() {
            let size = button.size;
            let offset = printer.size - self.borders.bot_right() - self.padding.bot_right() - size - Vec2::new(x, 0);
            // Add some special effect to the focused button
            button.draw(&printer.sub_printer(offset, size), focused && (self.focus == Focus::Button(i)));
            // Keep 1 blank between two buttons
            x += size.x + 1;
            // Also keep 1 blank above the buttons
            height = max(height, size.y+1);
        }

        // What do we have left?
        let inner_size = printer.size
            - Vec2::new(0, height)
            - self.borders.combined()
            - self.padding.combined();

        self.content.draw(&printer.sub_printer(self.borders.top_left() + self.padding.top_left(), inner_size), focused && self.focus == Focus::Content);

        printer.print_box(Vec2::new(0,0), printer.size);

        if self.title.len() > 0 {
            let x = (printer.size.x - self.title.len()) / 2;
            printer.print((x-2,0), "┤ ");
            printer.print((x+self.title.len(),0), " ├");

            printer.with_style(color::TITLE_PRIMARY, |p| p.print((x,0), &self.title));
        }

    }

    fn get_min_size(&self, req: SizeRequest) -> Vec2 {
        // Padding and borders are not available for kids.
        let content_req = req.reduced(self.padding.combined() + self.borders.combined());
        let content_size = self.content.get_min_size(content_req);

        let mut buttons_size = Vec2::new(0,0);
        for button in self.buttons.iter() {
            let s = button.view.get_min_size(req);
            buttons_size.x += s.x + 1;
            buttons_size.y = max(buttons_size.y, s.y + 1);
        }

        // On the Y axis, we add buttons and content.
        // On the X axis, we take the max.
        let mut inner_size = Vec2::new(max(content_size.x, buttons_size.x),
                                   content_size.y + buttons_size.y)
                        + self.padding.combined() + self.borders.combined();

        if self.title.len() > 0 {
            // If we have a title, we have to fit it too!
            inner_size.x = max(inner_size.x, self.title.len() + 6);
        }

        inner_size
    }

    fn layout(&mut self, mut size: Vec2) {
        // Padding and borders are taken, sorry.
        size = size - (self.borders.combined() + self.padding.combined());
        let req = SizeRequest {
            w: DimensionRequest::AtMost(size.x),
            h: DimensionRequest::AtMost(size.y),
        };

        // Buttons are kings, we give them everything they want.
        let mut buttons_height = 0;
        for button in self.buttons.iter_mut().rev() {
            let size = button.get_min_size(req);
            buttons_height = max(buttons_height, size.y+1);
            button.layout(size);
        }

        // Poor content will have to make do with what's left.
        self.content.layout(size - Vec2::new(0, buttons_height));
    }

    fn on_key_event(&mut self, ch: i32) -> EventResult {
        match self.focus {
            // If we are on the content, we can only go down.
            Focus::Content => match self.content.on_key_event(ch) {
                EventResult::Ignored if !self.buttons.is_empty() => match ch {
                    ncurses::KEY_DOWN => {
                        // Default to leftmost button when going down.
                        self.focus = Focus::Button(0);
                        EventResult::Consumed(None)
                    },
                    _ => EventResult::Ignored,
                },
                res => res,
            },
            // If we are on a button, we have more choice
            Focus::Button(i) => match self.buttons[i].on_key_event(ch) {
                EventResult::Ignored => match ch {
                    // Up goes back to the content
                    ncurses::KEY_UP => {
                        if self.content.take_focus() {
                            self.focus = Focus::Content;
                            EventResult::Consumed(None)
                        } else {
                            EventResult::Ignored
                        }
                    },
                    // Left and Right move to other buttons
                    ncurses::KEY_RIGHT if i+1 < self.buttons.len() => {
                        self.focus = Focus::Button(i+1);
                        EventResult::Consumed(None)
                    },
                    ncurses::KEY_LEFT if i > 0 => {
                        self.focus = Focus::Button(i-1);
                        EventResult::Consumed(None)
                    },
                    _ => EventResult::Ignored,
                },
                res => res,
            },
        }
    }

    fn take_focus(&mut self) -> bool {
        // TODO: add a direction to the focus. Meanwhile, takes button first.
        if !self.buttons.is_empty() {
            self.focus = Focus::Button(0);
            true
        } else {
            self.focus = Focus::Content;
            self.content.take_focus()
        }
    }

    fn find(&mut self, selector: &Selector) -> Option<&mut Any> {
        self.content.find(selector)
    }
}
<commit_msg>Dialog gives focus to its content first<commit_after>use std::cmp::max;
use std::any::Any;

use ncurses;

use color;
use ::{Cursive};
use event::EventResult;
use view::{View,SizeRequest,DimensionRequest,Selector};
use view::{Button,SizedView};
use vec::{Vec2,Vec4,ToVec4};
use printer::Printer;

#[derive(PartialEq)]
enum Focus {
    Content,
    Button(usize),
}

/// Popup-like view with a main content, and optional buttons under it.
///
/// # Examples
///
/// ```
/// let dialog = Dialog::new(TextView::new("Hello!")).button("Ok", |s,_| s.quit());
/// ```
pub struct Dialog {
    title: String,
    content: Box<View>,

    buttons: Vec<SizedView<Button>>,

    padding: Vec4,
    borders: Vec4,

    focus: Focus,
}

impl Dialog {
    /// Creates a new Dialog with the given content.
    pub fn new<V: View + 'static>(view: V) -> Self {
        Dialog {
            content: Box::new(view),
            buttons: Vec::new(),
            title: String::new(),
            focus: Focus::Content,
            padding: Vec4::new(1,1,0,0),
            borders: Vec4::new(1,1,1,1),
        }
    }

    /// Adds a button to the dialog with the given label and callback.
    ///
    /// Consumes and returns self for easy chaining.
    pub fn button<'a, F>(mut self, label: &'a str, cb: F) -> Self
        where F: Fn(&mut Cursive) + 'static
    {
        self.buttons.push(SizedView::new(Button::new(label, cb)));

        self
    }

    /// Shortcut method to add a button that will dismiss the dialog.
    pub fn dismiss_button<'a>(self, label: &'a str) -> Self {
        self.button(label, |s| s.screen_mut().pop_layer())
    }

    /// Sets the title of the dialog.
    /// If not empty, it will be visible at the top.
    pub fn title(mut self, label: &str) -> Self {
        self.title = label.to_string();
        self
    }

    /// Sets the padding in the dialog (around content and buttons).
    pub fn padding<T: ToVec4>(mut self, padding: T) -> Self {
        self.padding = padding.to_vec4();

        self
    }

}

impl View for Dialog {
    fn draw(&mut self, printer: &Printer, focused: bool) {

        // This will be the height used by the buttons.
        let mut height = 0;
        // Current horizontal position of the next button we'll draw.
        let mut x = 0;
        for (i,button) in self.buttons.iter_mut().enumerate().rev() {
            let size = button.size;
            let offset = printer.size - self.borders.bot_right() - self.padding.bot_right() - size - Vec2::new(x, 0);
            // Add some special effect to the focused button
            button.draw(&printer.sub_printer(offset, size), focused && (self.focus == Focus::Button(i)));
            // Keep 1 blank between two buttons
            x += size.x + 1;
            // Also keep 1 blank above the buttons
            height = max(height, size.y+1);
        }

        // What do we have left?
        let inner_size = printer.size
            - Vec2::new(0, height)
            - self.borders.combined()
            - self.padding.combined();

        self.content.draw(&printer.sub_printer(self.borders.top_left() + self.padding.top_left(), inner_size), focused && self.focus == Focus::Content);

        printer.print_box(Vec2::new(0,0), printer.size);

        if self.title.len() > 0 {
            let x = (printer.size.x - self.title.len()) / 2;
            printer.print((x-2,0), "┤ ");
            printer.print((x+self.title.len(),0), " ├");

            printer.with_style(color::TITLE_PRIMARY, |p| p.print((x,0), &self.title));
        }

    }

    fn get_min_size(&self, req: SizeRequest) -> Vec2 {
        // Padding and borders are not available for kids.
        let content_req = req.reduced(self.padding.combined() + self.borders.combined());
        let content_size = self.content.get_min_size(content_req);

        let mut buttons_size = Vec2::new(0,0);
        for button in self.buttons.iter() {
            let s = button.view.get_min_size(req);
            buttons_size.x += s.x + 1;
            buttons_size.y = max(buttons_size.y, s.y + 1);
        }

        // On the Y axis, we add buttons and content.
        // On the X axis, we take the max.
        let mut inner_size = Vec2::new(max(content_size.x, buttons_size.x),
                                   content_size.y + buttons_size.y)
                        + self.padding.combined() + self.borders.combined();

        if self.title.len() > 0 {
            // If we have a title, we have to fit it too!
            inner_size.x = max(inner_size.x, self.title.len() + 6);
        }

        inner_size
    }

    fn layout(&mut self, mut size: Vec2) {
        // Padding and borders are taken, sorry.
        size = size - (self.borders.combined() + self.padding.combined());
        let req = SizeRequest {
            w: DimensionRequest::AtMost(size.x),
            h: DimensionRequest::AtMost(size.y),
        };

        // Buttons are kings, we give them everything they want.
        let mut buttons_height = 0;
        for button in self.buttons.iter_mut().rev() {
            let size = button.get_min_size(req);
            buttons_height = max(buttons_height, size.y+1);
            button.layout(size);
        }

        // Poor content will have to make do with what's left.
        self.content.layout(size - Vec2::new(0, buttons_height));
    }

    fn on_key_event(&mut self, ch: i32) -> EventResult {
        match self.focus {
            // If we are on the content, we can only go down.
            Focus::Content => match self.content.on_key_event(ch) {
                EventResult::Ignored if !self.buttons.is_empty() => match ch {
                    ncurses::KEY_DOWN => {
                        // Default to leftmost button when going down.
                        self.focus = Focus::Button(0);
                        EventResult::Consumed(None)
                    },
                    _ => EventResult::Ignored,
                },
                res => res,
            },
            // If we are on a button, we have more choice
            Focus::Button(i) => match self.buttons[i].on_key_event(ch) {
                EventResult::Ignored => match ch {
                    // Up goes back to the content
                    ncurses::KEY_UP => {
                        if self.content.take_focus() {
                            self.focus = Focus::Content;
                            EventResult::Consumed(None)
                        } else {
                            EventResult::Ignored
                        }
                    },
                    // Left and Right move to other buttons
                    ncurses::KEY_RIGHT if i+1 < self.buttons.len() => {
                        self.focus = Focus::Button(i+1);
                        EventResult::Consumed(None)
                    },
                    ncurses::KEY_LEFT if i > 0 => {
                        self.focus = Focus::Button(i-1);
                        EventResult::Consumed(None)
                    },
                    _ => EventResult::Ignored,
                },
                res => res,
            },
        }
    }

    fn take_focus(&mut self) -> bool {
        // TODO: add a direction to the focus. Meanwhile, takes button first.
        if self.content.take_focus() {
            self.focus = Focus::Content;
            true
        } else if !self.buttons.is_empty() {
            self.focus = Focus::Button(0);
            true
        } else {
            false
        }
    }

    fn find(&mut self, selector: &Selector) -> Option<&mut Any> {
        self.content.find(selector)
    }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    protocol::{self, AMQPError, AMQPHardError},
    reactor::{ReactorBuilder, ReactorHandle, Slot},
    socket_state::SocketState,
    tcp::HandshakeResult,
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result, TcpStream,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{debug, error, trace};
use std::{
    collections::VecDeque,
    convert::TryFrom,
    io::{self, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Connected,
    Stop,
}

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: Box<dyn ReactorHandle + Send>,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    stream: TcpStream,
    slot: Slot,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

impl IoLoop {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: HandshakeResult,
        reactor_builder: &dyn ReactorBuilder,
    ) -> Result<Self> {
        let mut stream = TcpStream::try_from(stream)?;
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = reactor_builder.build(heartbeat.clone())?;
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE as usize,
            configuration.frame_max() as usize,
        );
        let slot = reactor.register(stream.inner_mut(), socket_state.handle())?;

        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            stream,
            slot,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
        })
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
                self.reactor.start_heartbeat();
            }
            let peer = self.stream.inner().peer_addr()?;
            debug!("Connected to {}", peer);
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        match self.status {
            Status::Initial => self.finish_setup(),
            Status::Connected => Ok(true),
            Status::Stop => Ok(false),
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable() && self.receive_buffer.available_space() > 0
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status != Status::Connected
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        let io_loop_handle = self.connection_io_loop_handle.clone();
        io_loop_handle.register(
            ThreadBuilder::new()
                .name("lapin-io-loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.heartbeat.cancel();
                    self.reactor.shutdown()?;
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events()?;
        if !self.ensure_setup()? {
            return Ok(());
        }
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        if self.stream.is_handshaking() {
            self.stream.handshake()?;
            if self.stream.is_handshaking() {
                // We hit WOULDBLOCK while handshaking, wait for the next socket event
                return Ok(());
            }
        }
        self.attempt_flush()?;
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        self.reactor.shutdown()?;
        Err(error)
    }

    fn attempt_flush(&mut self) -> Result<()> {
        let res = self.flush();
        if let Err(e) = self.socket_state.handle_flush_result(res) {
            error!("error flushing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_read_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_read_result(result, &*self.reactor, self.slot)
        {
            error!("error reading: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_write_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_write_result(result, &*self.reactor, self.slot)
        {
            error!("error writing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn flush(&mut self) -> Result<()> {
        self.stream.flush()?;
        self.poll_internal_rpc()
    }

    fn write(&mut self) -> Result<()> {
        while self.can_write() {
            let res = self.write_to_stream();
            self.handle_write_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            self.handle_read_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.flush()?;
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.stream)?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                let sz = self.receive_buffer.read_from(&mut self.stream)?;

                if sz > 0 {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                } else {
                    error!("Socket was readable but we read 0, marking as wouldblock");
                    self.handle_read_result(
                        Err(io::Error::from(io::ErrorKind::WouldBlock).into()),
                    )?;
                }
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                let frame_max = self.configuration.frame_max() as usize;
                if frame_max > 0 && consumed > frame_max {
                    error!("received large ({} bytes) frame", consumed);
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("frame too large: {} bytes", consumed).into(),
                    );
                    self.internal_rpc.handle().close_connection(
                        error.get_id(),
                        error.get_message().to_string(),
                        0,
                        0,
                    );
                    self.poll_internal_rpc()?;
                    self.critical_error(Error::ProtocolError(error))?;
                }
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<commit_msg>consider the connection to unresponsive servers as broken<commit_after>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    protocol::{self, AMQPError, AMQPHardError},
    reactor::{ReactorBuilder, ReactorHandle, Slot},
    socket_state::SocketState,
    tcp::HandshakeResult,
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result, TcpStream,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{debug, error, trace};
use std::{
    collections::VecDeque,
    convert::TryFrom,
    io::{self, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

const FRAMES_STORAGE: usize = 32;

// The amount of Read(0) we tolerate before onsidering the connection as broken.
// This might seem high but better safe than sorry.
const BAD_READS_TOLERANCE: u8 = 16;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Connected,
    Stop,
}

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: Box<dyn ReactorHandle + Send>,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    stream: TcpStream,
    slot: Slot,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    bad_reads: u8,
}

impl IoLoop {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: HandshakeResult,
        reactor_builder: &dyn ReactorBuilder,
    ) -> Result<Self> {
        let mut stream = TcpStream::try_from(stream)?;
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = reactor_builder.build(heartbeat.clone())?;
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE as usize,
            configuration.frame_max() as usize,
        );
        let slot = reactor.register(stream.inner_mut(), socket_state.handle())?;

        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            stream,
            slot,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
            bad_reads: 0,
        })
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
                self.reactor.start_heartbeat();
            }
            let peer = self.stream.inner().peer_addr()?;
            debug!("Connected to {}", peer);
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        match self.status {
            Status::Initial => self.finish_setup(),
            Status::Connected => Ok(true),
            Status::Stop => Ok(false),
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable() && self.receive_buffer.available_space() > 0
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status != Status::Connected
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        let io_loop_handle = self.connection_io_loop_handle.clone();
        io_loop_handle.register(
            ThreadBuilder::new()
                .name("lapin-io-loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.heartbeat.cancel();
                    self.reactor.shutdown()?;
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events()?;
        if !self.ensure_setup()? {
            return Ok(());
        }
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        if self.stream.is_handshaking() {
            self.stream.handshake()?;
            if self.stream.is_handshaking() {
                // We hit WOULDBLOCK while handshaking, wait for the next socket event
                return Ok(());
            }
        }
        self.attempt_flush()?;
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        self.reactor.shutdown()?;
        Err(error)
    }

    fn attempt_flush(&mut self) -> Result<()> {
        let res = self.flush();
        if let Err(e) = self.socket_state.handle_flush_result(res) {
            error!("error flushing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_read_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_read_result(result, &*self.reactor, self.slot)
        {
            error!("error reading: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_write_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_write_result(result, &*self.reactor, self.slot)
        {
            error!("error writing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn flush(&mut self) -> Result<()> {
        self.stream.flush()?;
        self.poll_internal_rpc()
    }

    fn write(&mut self) -> Result<()> {
        while self.can_write() {
            let res = self.write_to_stream();
            self.handle_write_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            self.handle_read_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.flush()?;
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.stream)?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                let sz = self.receive_buffer.read_from(&mut self.stream)?;

                if sz > 0 {
                    self.bad_reads = 0;
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                } else {
                    self.bad_reads += 1;
                    let io_error_kind = if self.bad_reads > BAD_READS_TOLERANCE {
                        error!("Socket was readable but we read 0 {} times in a row, mark the connection as broken", BAD_READS_TOLERANCE);
                        io::ErrorKind::ConnectionAborted
                    } else {
                        error!("Socket was readable but we read 0, marking as wouldblock");
                        io::ErrorKind::WouldBlock
                    };
                    self.handle_read_result(Err(io::Error::from(io_error_kind).into()))?;
                }
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                let frame_max = self.configuration.frame_max() as usize;
                if frame_max > 0 && consumed > frame_max {
                    error!("received large ({} bytes) frame", consumed);
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("frame too large: {} bytes", consumed).into(),
                    );
                    self.internal_rpc.handle().close_connection(
                        error.get_id(),
                        error.get_message().to_string(),
                        0,
                        0,
                    );
                    self.poll_internal_rpc()?;
                    self.critical_error(Error::ProtocolError(error))?;
                }
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat * 500)); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.connection.send_heartbeat()?;
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
            self.status = Status::Stop;
        }
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            self.last_write = Instant::now();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                self.send_continue()?;
            }
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>IoLoop: always send the first heartbeat<commit_after>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    io::{Read, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
    last_write: Instant,
    sent_initial_heartbeat: bool,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
            last_write: Instant::now(),
            sent_initial_heartbeat: false,
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat * 500)); // * 1000 (ms) / 2 (half the negociated timeout)
                self.poll_timeout = Some(heartbeat);
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&self) -> bool {
        self.can_write && self.has_data() && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn poll_timeout(&self) -> Option<Duration> {
        self.poll_timeout.map(|timeout| {
            timeout
                .checked_sub(self.last_write.elapsed())
                .unwrap_or_else(|| Duration::from_secs(0))
        })
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout())?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.should_heartbeat() {
                self.sent_initial_heartbeat = true;
                self.connection.send_heartbeat()?;
            }
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn should_heartbeat(&self) -> bool {
        if let Some(heartbeat_timeout) = self.poll_timeout {
            self.last_write.elapsed() > heartbeat_timeout
        } else {
            false
        }
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
            self.status = Status::Stop;
        }
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.send_buffer.write_to(&mut self.socket)?;

        if sz > 0 {
            // If we didn't send the first heartbeat, don't update last_write to trick us into sending it
            if self.sent_initial_heartbeat {
                self.last_write = Instant::now();
            }

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz, true);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                self.send_continue()?;
            }
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.receive_buffer
                    .read_from(&mut self.socket, false)
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz, false);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((next_msg, resolver)) = self.connection.next_frame() {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => {
                    self.serialized_frames.push_back((sz, resolver));
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((next_msg, resolver))?;
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                self.receive_buffer.consume(consumed, false);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<|endoftext|>