<commit_before>use std::path::Path;

use regex::bytes::Regex;
use term::{Attr, Terminal};
use term::color;

use pathutil::strip_prefix;
use types::FileTypeDef;

/// Printer encapsulates all output logic for searching.
///
/// Note that we currently ignore all write errors. It's probably worthwhile
/// to fix this, but printers are only ever used for writes to stdout or
/// writes to memory, neither of which commonly fail.
pub struct Printer<W> {
    /// The underlying writer.
    wtr: W,
    /// Whether anything has been printed to wtr yet.
    has_printed: bool,
    /// Whether to show column numbers for the first match or not.
    column: bool,
    /// The string to use to separate non-contiguous runs of context lines.
    context_separator: Vec<u8>,
    /// The end-of-line terminator used by the printer. In general, eols are
    /// printed via the match directly, but occasionally we need to insert them
    /// ourselves (for example, to print a context separator).
    eol: u8,
    /// Whether to show file name as a heading or not.
    ///
    /// N.B. If with_filename is false, then this setting has no effect.
    heading: bool,
    /// Whether to show every match on its own line.
    line_per_match: bool,
    /// Whether to suppress all output.
    quiet: bool,
    /// A string to use as a replacement of each match in a matching line.
    replace: Option<Vec<u8>>,
    /// Whether to prefix each match with the corresponding file name.
    with_filename: bool,
}

impl<W: Terminal + Send> Printer<W> {
    /// Create a new printer that writes to wtr.
    pub fn new(wtr: W) -> Printer<W> {
        Printer {
            wtr: wtr,
            has_printed: false,
            column: false,
            context_separator: "--".to_string().into_bytes(),
            eol: b'\n',
            heading: false,
            line_per_match: false,
            quiet: false,
            replace: None,
            with_filename: false,
        }
    }

    /// When set, column numbers will be printed for the first match on each
    /// line.
    pub fn column(mut self, yes: bool) -> Printer<W> {
        self.column = yes;
        self
    }

    /// Set the context separator. The default is `--`.
    pub fn context_separator(mut self, sep: Vec<u8>) -> Printer<W> {
        self.context_separator = sep;
        self
    }

    /// Set the end-of-line terminator. The default is `\n`.
    pub fn eol(mut self, eol: u8) -> Printer<W> {
        self.eol = eol;
        self
    }

    /// Whether to show file name as a heading or not.
    ///
    /// N.B. If with_filename is false, then this setting has no effect.
    pub fn heading(mut self, yes: bool) -> Printer<W> {
        self.heading = yes;
        self
    }

    /// Whether to show every match on its own line.
    pub fn line_per_match(mut self, yes: bool) -> Printer<W> {
        self.line_per_match = yes;
        self
    }

    /// When set, all output is suppressed.
    pub fn quiet(mut self, yes: bool) -> Printer<W> {
        self.quiet = yes;
        self
    }

    /// Replace every match in each matching line with the replacement string
    /// given.
    ///
    /// The replacement string syntax is documented here:
    /// https://doc.rust-lang.org/regex/regex/bytes/struct.Captures.html#method.expand
    pub fn replace(mut self, replacement: Vec<u8>) -> Printer<W> {
        self.replace = Some(replacement);
        self
    }

    /// When set, each match is prefixed with the file name that it came from.
    pub fn with_filename(mut self, yes: bool) -> Printer<W> {
        self.with_filename = yes;
        self
    }

    /// Returns true if and only if something has been printed.
    pub fn has_printed(&self) -> bool {
        self.has_printed
    }

    /// Returns true if the printer has been configured to be quiet.
    pub fn is_quiet(&self) -> bool {
        self.quiet
    }

    /// Flushes the underlying writer and returns it.
    pub fn into_inner(mut self) -> W {
        let _ = self.wtr.flush();
        self.wtr
    }

    /// Prints a type definition.
    pub fn type_def(&mut self, def: &FileTypeDef) {
        self.write(def.name().as_bytes());
        self.write(b": ");
        let mut first = true;
        for pat in def.patterns() {
            if !first {
                self.write(b", ");
            }
            self.write(pat.as_bytes());
            first = false;
        }
        self.write_eol();
    }

    /// Prints the given path.
    pub fn path<P: AsRef<Path>>(&mut self, path: P) {
        let path = strip_prefix("./", path.as_ref()).unwrap_or(path.as_ref());
        self.write_path(path);
        self.write_eol();
    }

    /// Prints the given path and a count of the number of matches found.
    pub fn path_count<P: AsRef<Path>>(&mut self, path: P, count: u64) {
        if self.with_filename {
            self.write_path(path);
            self.write(b":");
        }
        self.write(count.to_string().as_bytes());
        self.write_eol();
    }

    /// Prints the context separator.
    pub fn context_separate(&mut self) {
        // N.B. We can't use `write` here because of borrowing restrictions.
        if self.quiet {
            return;
        }
        if self.context_separator.is_empty() {
            return;
        }
        self.has_printed = true;
        let _ = self.wtr.write_all(&self.context_separator);
        let _ = self.wtr.write_all(&[self.eol]);
    }

    pub fn matched<P: AsRef<Path>>(
        &mut self,
        re: &Regex,
        path: P,
        buf: &[u8],
        start: usize,
        end: usize,
        line_number: Option<u64>,
    ) {
        if !self.line_per_match {
            let column =
                if self.column {
                    Some(re.find(&buf[start..end])
                           .map(|(s, _)| s + 1).unwrap_or(0) as u64)
                } else {
                    None
                };
            return self.write_match(
                re, path, buf, start, end, line_number, column);
        }
        for (s, _) in re.find_iter(&buf[start..end]) {
            let column = if self.column { Some(s as u64) } else { None };
            self.write_match(
                re, path.as_ref(), buf, start, end, line_number, column);
        }
    }

    fn write_match<P: AsRef<Path>>(
        &mut self,
        re: &Regex,
        path: P,
        buf: &[u8],
        start: usize,
        end: usize,
        line_number: Option<u64>,
        column: Option<u64>,
    ) {
        if self.heading && self.with_filename && !self.has_printed {
            self.write_heading(path.as_ref());
        } else if !self.heading && self.with_filename {
            self.write_path(path.as_ref());
            self.write(b":");
        }
        if let Some(line_number) = line_number {
            self.line_number(line_number, b':');
        }
        if let Some(c) = column {
            self.write((c + 1).to_string().as_bytes());
            self.write(b":");
        }
        if self.replace.is_some() {
            let line = re.replace_all(
                &buf[start..end], &**self.replace.as_ref().unwrap());
            self.write(&line);
        } else {
            self.write_matched_line(re, &buf[start..end]);
        }
        if buf[start..end].last() != Some(&self.eol) {
            self.write_eol();
        }
    }

    fn write_matched_line(&mut self, re: &Regex, buf: &[u8]) {
        if !self.wtr.supports_color() {
            self.write(buf);
            return;
        }
        let mut last_written = 0;
        for (s, e) in re.find_iter(buf) {
            self.write(&buf[last_written..s]);
            let _ = self.wtr.fg(color::BRIGHT_RED);
            let _ = self.wtr.attr(Attr::Bold);
            self.write(&buf[s..e]);
            let _ = self.wtr.reset();
            last_written = e;
        }
        self.write(&buf[last_written..]);
    }

    pub fn context<P: AsRef<Path>>(
        &mut self,
        path: P,
        buf: &[u8],
        start: usize,
        end: usize,
        line_number: Option<u64>,
    ) {
        if self.heading && self.with_filename && !self.has_printed {
            self.write_heading(path.as_ref());
        } else if !self.heading && self.with_filename {
            self.write_path(path.as_ref());
            self.write(b"-");
        }
        if let Some(line_number) = line_number {
            self.line_number(line_number, b'-');
        }
        self.write(&buf[start..end]);
        if buf[start..end].last() != Some(&self.eol) {
            self.write_eol();
        }
    }

    fn write_heading<P: AsRef<Path>>(&mut self, path: P) {
        if self.wtr.supports_color() {
            let _ = self.wtr.fg(color::BRIGHT_GREEN);
            let _ = self.wtr.attr(Attr::Bold);
        }
        self.write_path(path.as_ref());
        self.write_eol();
        if self.wtr.supports_color() {
            let _ = self.wtr.reset();
        }
    }

    fn line_number(&mut self, n: u64, sep: u8) {
        if self.wtr.supports_color() {
            let _ = self.wtr.fg(color::BRIGHT_BLUE);
            let _ = self.wtr.attr(Attr::Bold);
        }
        self.write(n.to_string().as_bytes());
        if self.wtr.supports_color() {
            let _ = self.wtr.reset();
        }
        self.write(&[sep]);
    }

    #[cfg(unix)]
    fn write_path<P: AsRef<Path>>(&mut self, path: P) {
        use std::os::unix::ffi::OsStrExt;

        let path = path.as_ref().as_os_str().as_bytes();
        self.write(path);
    }

    #[cfg(not(unix))]
    fn write_path<P: AsRef<Path>>(&mut self, path: P) {
        self.write(path.as_ref().to_string_lossy().as_bytes());
    }

    fn write(&mut self, buf: &[u8]) {
        if self.quiet {
            return;
        }
        self.has_printed = true;
        let _ = self.wtr.write_all(buf);
    }

    fn write_eol(&mut self) {
        let eol = self.eol;
        self.write(&[eol]);
    }
}
<commit_msg>Add colour choice<commit_after>use std::path::Path;

use regex::bytes::Regex;
use term::{Attr, Terminal};
use term::color;

use pathutil::strip_prefix;
use types::FileTypeDef;

/// Printer encapsulates all output logic for searching.
///
/// Note that we currently ignore all write errors. It's probably worthwhile
/// to fix this, but printers are only ever used for writes to stdout or
/// writes to memory, neither of which commonly fail.
pub struct Printer<W> {
    /// The underlying writer.
    wtr: W,
    /// Whether anything has been printed to wtr yet.
    has_printed: bool,
    /// Whether to show column numbers for the first match or not.
    column: bool,
    /// The string to use to separate non-contiguous runs of context lines.
    context_separator: Vec<u8>,
    /// The end-of-line terminator used by the printer. In general, eols are
    /// printed via the match directly, but occasionally we need to insert them
    /// ourselves (for example, to print a context separator).
    eol: u8,
    /// Whether to show file name as a heading or not.
    ///
    /// N.B. If with_filename is false, then this setting has no effect.
    heading: bool,
    /// Whether to show every match on its own line.
    line_per_match: bool,
    /// Whether to suppress all output.
    quiet: bool,
    /// A string to use as a replacement of each match in a matching line.
    replace: Option<Vec<u8>>,
    /// Whether to prefix each match with the corresponding file name.
    with_filename: bool,

    /// The choice of Colours
    color_choice: ColorChoice
}

struct ColorChoice {
    matched_line: color::Color,
    heading: color::Color,
    line_number: color::Color
}

impl ColorChoice {

    #[cfg(unix)]
    pub fn new() -> ColorChoice {
        ColorChoice {
            matched_line: color::RED,
            heading: color::GREEN,
            line_number: color::BLUE
        }
    }

    #[cfg(not(unix))]
    pub fn new() -> ColorChoice {
        ColorChoice {
            matched_line: color::BRIGHT_RED,
            heading: color::BRIGHT_GREEN,
            line_number: color::BRIGHT_BLUE
        }
    }
}

impl<W: Terminal + Send> Printer<W> {
    /// Create a new printer that writes to wtr.
    pub fn new(wtr: W) -> Printer<W> {
        Printer {
            wtr: wtr,
            has_printed: false,
            column: false,
            context_separator: "--".to_string().into_bytes(),
            eol: b'\n',
            heading: false,
            line_per_match: false,
            quiet: false,
            replace: None,
            with_filename: false,
            color_choice: ColorChoice::new()
        }
    }

    /// When set, column numbers will be printed for the first match on each
    /// line.
    pub fn column(mut self, yes: bool) -> Printer<W> {
        self.column = yes;
        self
    }

    /// Set the context separator. The default is `--`.
    pub fn context_separator(mut self, sep: Vec<u8>) -> Printer<W> {
        self.context_separator = sep;
        self
    }

    /// Set the end-of-line terminator. The default is `\n`.
    pub fn eol(mut self, eol: u8) -> Printer<W> {
        self.eol = eol;
        self
    }

    /// Whether to show file name as a heading or not.
    ///
    /// N.B. If with_filename is false, then this setting has no effect.
    pub fn heading(mut self, yes: bool) -> Printer<W> {
        self.heading = yes;
        self
    }

    /// Whether to show every match on its own line.
    pub fn line_per_match(mut self, yes: bool) -> Printer<W> {
        self.line_per_match = yes;
        self
    }

    /// When set, all output is suppressed.
    pub fn quiet(mut self, yes: bool) -> Printer<W> {
        self.quiet = yes;
        self
    }

    /// Replace every match in each matching line with the replacement string
    /// given.
    ///
    /// The replacement string syntax is documented here:
    /// https://doc.rust-lang.org/regex/regex/bytes/struct.Captures.html#method.expand
    pub fn replace(mut self, replacement: Vec<u8>) -> Printer<W> {
        self.replace = Some(replacement);
        self
    }

    /// When set, each match is prefixed with the file name that it came from.
    pub fn with_filename(mut self, yes: bool) -> Printer<W> {
        self.with_filename = yes;
        self
    }

    /// Returns true if and only if something has been printed.
    pub fn has_printed(&self) -> bool {
        self.has_printed
    }

    /// Returns true if the printer has been configured to be quiet.
    pub fn is_quiet(&self) -> bool {
        self.quiet
    }

    /// Flushes the underlying writer and returns it.
    pub fn into_inner(mut self) -> W {
        let _ = self.wtr.flush();
        self.wtr
    }

    /// Prints a type definition.
    pub fn type_def(&mut self, def: &FileTypeDef) {
        self.write(def.name().as_bytes());
        self.write(b": ");
        let mut first = true;
        for pat in def.patterns() {
            if !first {
                self.write(b", ");
            }
            self.write(pat.as_bytes());
            first = false;
        }
        self.write_eol();
    }

    /// Prints the given path.
    pub fn path<P: AsRef<Path>>(&mut self, path: P) {
        let path = strip_prefix("./", path.as_ref()).unwrap_or(path.as_ref());
        self.write_path(path);
        self.write_eol();
    }

    /// Prints the given path and a count of the number of matches found.
    pub fn path_count<P: AsRef<Path>>(&mut self, path: P, count: u64) {
        if self.with_filename {
            self.write_path(path);
            self.write(b":");
        }
        self.write(count.to_string().as_bytes());
        self.write_eol();
    }

    /// Prints the context separator.
    pub fn context_separate(&mut self) {
        // N.B. We can't use `write` here because of borrowing restrictions.
        if self.quiet {
            return;
        }
        if self.context_separator.is_empty() {
            return;
        }
        self.has_printed = true;
        let _ = self.wtr.write_all(&self.context_separator);
        let _ = self.wtr.write_all(&[self.eol]);
    }

    pub fn matched<P: AsRef<Path>>(
        &mut self,
        re: &Regex,
        path: P,
        buf: &[u8],
        start: usize,
        end: usize,
        line_number: Option<u64>,
    ) {
        if !self.line_per_match {
            let column =
                if self.column {
                    Some(re.find(&buf[start..end])
                           .map(|(s, _)| s + 1).unwrap_or(0) as u64)
                } else {
                    None
                };
            return self.write_match(
                re, path, buf, start, end, line_number, column);
        }
        for (s, _) in re.find_iter(&buf[start..end]) {
            let column = if self.column { Some(s as u64) } else { None };
            self.write_match(
                re, path.as_ref(), buf, start, end, line_number, column);
        }
    }

    fn write_match<P: AsRef<Path>>(
        &mut self,
        re: &Regex,
        path: P,
        buf: &[u8],
        start: usize,
        end: usize,
        line_number: Option<u64>,
        column: Option<u64>,
    ) {
        if self.heading && self.with_filename && !self.has_printed {
            self.write_heading(path.as_ref());
        } else if !self.heading && self.with_filename {
            self.write_path(path.as_ref());
            self.write(b":");
        }
        if let Some(line_number) = line_number {
            self.line_number(line_number, b':');
        }
        if let Some(c) = column {
            self.write((c + 1).to_string().as_bytes());
            self.write(b":");
        }
        if self.replace.is_some() {
            let line = re.replace_all(
                &buf[start..end], &**self.replace.as_ref().unwrap());
            self.write(&line);
        } else {
            self.write_matched_line(re, &buf[start..end]);
        }
        if buf[start..end].last() != Some(&self.eol) {
            self.write_eol();
        }
    }

    fn write_matched_line(&mut self, re: &Regex, buf: &[u8]) {
        if !self.wtr.supports_color() {
            self.write(buf);
            return;
        }
        let mut last_written = 0;
        for (s, e) in re.find_iter(buf) {
            self.write(&buf[last_written..s]);
            let _ = self.wtr.fg(self.color_choice.matched_line);
            let _ = self.wtr.attr(Attr::Bold);
            self.write(&buf[s..e]);
            let _ = self.wtr.reset();
            last_written = e;
        }
        self.write(&buf[last_written..]);
    }

    pub fn context<P: AsRef<Path>>(
        &mut self,
        path: P,
        buf: &[u8],
        start: usize,
        end: usize,
        line_number: Option<u64>,
    ) {
        if self.heading && self.with_filename && !self.has_printed {
            self.write_heading(path.as_ref());
        } else if !self.heading && self.with_filename {
            self.write_path(path.as_ref());
            self.write(b"-");
        }
        if let Some(line_number) = line_number {
            self.line_number(line_number, b'-');
        }
        self.write(&buf[start..end]);
        if buf[start..end].last() != Some(&self.eol) {
            self.write_eol();
        }
    }

    fn write_heading<P: AsRef<Path>>(&mut self, path: P) {
        if self.wtr.supports_color() {
            let _ = self.wtr.fg(self.color_choice.heading);
            let _ = self.wtr.attr(Attr::Bold);
        }
        self.write_path(path.as_ref());
        self.write_eol();
        if self.wtr.supports_color() {
            let _ = self.wtr.reset();
        }
    }

    fn line_number(&mut self, n: u64, sep: u8) {
        if self.wtr.supports_color() {
            let _ = self.wtr.fg(self.color_choice.line_number);
            let _ = self.wtr.attr(Attr::Bold);
        }
        self.write(n.to_string().as_bytes());
        if self.wtr.supports_color() {
            let _ = self.wtr.reset();
        }
        self.write(&[sep]);
    }

    #[cfg(unix)]
    fn write_path<P: AsRef<Path>>(&mut self, path: P) {
        use std::os::unix::ffi::OsStrExt;

        let path = path.as_ref().as_os_str().as_bytes();
        self.write(path);
    }

    #[cfg(not(unix))]
    fn write_path<P: AsRef<Path>>(&mut self, path: P) {
        self.write(path.as_ref().to_string_lossy().as_bytes());
    }

    fn write(&mut self, buf: &[u8]) {
        if self.quiet {
            return;
        }
        self.has_printed = true;
        let _ = self.wtr.write_all(buf);
    }

    fn write_eol(&mut self) {
        let eol = self.eol;
        self.write(&[eol]);
    }
}
<|endoftext|>


<commit_before>//
// process.rs
// Copyright (C) 2017 Szymon Urbaś <szymon.urbas@aol.com>
// Distributed under terms of the BSD (2-clause) license.
//
// Created on: 15 Mar 2017 18:06:31 +0100 (CET)
//

#[derive(Debug)]
pub struct Process {
  pub id: usize,
  pub name: String,
  pub arrival_time: usize,
  pub burst_time: usize,
  pub execution_time: usize,
}

impl Process {
  pub fn with_burst_time(&mut self, burst_time: usize) -> Process {
    Process {
      id: self.id,
      name: self.name.clone(),
      arrival_time: self.arrival_time,
      burst_time: burst_time,
      execution_time: self.execution_time,
    }
  }

  pub fn done_executing(&self) -> bool {
    self.execution_time >= self.burst_time
  }

  pub fn record_execution(&mut self) {
    self.execution_time += 1;
  }
}

pub struct ProcessSpawner {
  current_id: usize,
}

pub fn new_spawner() -> ProcessSpawner {
  ProcessSpawner {
    current_id: 0,
  }
}

impl ProcessSpawner {
  pub fn spawn(&mut self, burst_time: usize) -> Process {
    self.current_id += 1;

    Process {
      id: self.current_id - 1,
      name: format!("proc_{:?}", self.current_id - 1),
      arrival_time: 0,
      execution_time: 0,
      burst_time: burst_time,
    }
  }
}

/*
 * vi: ts=2 sw=2 expandtab
 */

<commit_msg>Remove unused function<commit_after>//
// process.rs
// Copyright (C) 2017 Szymon Urbaś <szymon.urbas@aol.com>
// Distributed under terms of the BSD (2-clause) license.
//
// Created on: 15 Mar 2017 18:06:31 +0100 (CET)
//

#[derive(Debug)]
pub struct Process {
  pub id: usize,
  pub name: String,
  pub arrival_time: usize,
  pub burst_time: usize,
  pub execution_time: usize,
}

impl Process {
  pub fn done_executing(&self) -> bool {
    self.execution_time >= self.burst_time
  }

  pub fn record_execution(&mut self) {
    self.execution_time += 1;
  }
}

pub struct ProcessSpawner {
  current_id: usize,
}

pub fn new_spawner() -> ProcessSpawner {
  ProcessSpawner {
    current_id: 0,
  }
}

impl ProcessSpawner {
  pub fn spawn(&mut self, burst_time: usize) -> Process {
    self.current_id += 1;

    Process {
      id: self.current_id - 1,
      name: format!("proc_{:?}", self.current_id - 1),
      arrival_time: 0,
      execution_time: 0,
      burst_time: burst_time,
    }
  }
}

/*
 * vi: ts=2 sw=2 expandtab
 */

<|endoftext|>


<commit_before>use std::mem;
use std::ptr;
use std::ffi::CString;
use super::gl_lib as gl;
use super::gl_lib::types::*;
use super::types::GLError;
use super::context::Context;
use super::shader::Shader;

pub struct Program {
    gl_id: GLuint
}

impl Program {
    pub fn gl_id(&self) -> GLuint {
        self.gl_id
    }
}

impl Drop for Program {
    fn drop(&mut self) {
        unsafe {
            gl::DeleteProgram(self.gl_id);
        }
    }
}

impl Context {
    pub fn create_program(&self) -> Result<Program, ()> {
        unsafe {
            let id = gl::CreateProgram();
            if id > 0 {
                Ok(Program { gl_id: id })
            }
            else {
                Err(())
            }
        }
    }

    pub fn attach_shader(&self, program: &mut Program, shader: &Shader) {
        unsafe {
            gl::AttachShader(program.gl_id(), shader.gl_id());
        }
    }

    pub fn link_program(&self, program: &mut Program) -> Result<(), GLError> {
        let success = unsafe {
            gl::LinkProgram(program.gl_id());
            let mut link_status : GLint = 0;
            gl::GetProgramiv(program.gl_id(),
                             gl::LINK_STATUS,
                             &mut link_status as *mut GLint);

            link_status == gl::TRUE as GLint
        };

        if success {
            Ok(())
        }
        else {
            unsafe {
                let mut info_length : GLint = mem::uninitialized();
                gl::GetProgramiv(program.gl_id(),
                                 gl::INFO_LOG_LENGTH,
                                 &mut info_length as *mut GLint);

                let mut bytes = Vec::<u8>::with_capacity(info_length as usize);

                gl::GetProgramInfoLog(program.gl_id(),
                                      info_length,
                                      ptr::null_mut(),
                                      bytes.as_mut_ptr() as *mut GLchar);
                bytes.set_len((info_length - 1) as usize);

                let msg = String::from_utf8(bytes)
                                 .unwrap_or(String::from("<Unknown error>"));

                Err(GLError { message: msg })
            }
        }
    }

    pub fn get_attrib_location(&self, program: &Program, name: &str)
        -> Result<ProgramAttrib, ()>
    {
        let c_str = try!(CString::new(name).or(Err(())));
        let str_ptr = c_str.as_ptr() as *const GLchar;
        unsafe {
            let index = gl::GetAttribLocation(program.gl_id(), str_ptr);
            if index >= 0 {
                Ok(ProgramAttrib { gl_index: index as GLuint })
            }
            else {
                Err(())
            }
        }
    }

    pub fn use_program(&self, program: &Program) {
        unsafe {
            gl::UseProgram(program.gl_id())
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct ProgramAttrib {
    pub gl_index: GLuint
}

#[derive(Debug, Clone, Copy)]
pub struct ProgramUniform {
    pub gl_index: GLuint
}
<commit_msg>Add `Context.get_uniform_location()`<commit_after>use std::mem;
use std::ptr;
use std::ffi::CString;
use super::gl_lib as gl;
use super::gl_lib::types::*;
use super::types::GLError;
use super::context::Context;
use super::shader::Shader;

pub struct Program {
    gl_id: GLuint
}

impl Program {
    pub fn gl_id(&self) -> GLuint {
        self.gl_id
    }
}

impl Drop for Program {
    fn drop(&mut self) {
        unsafe {
            gl::DeleteProgram(self.gl_id);
        }
    }
}

impl Context {
    pub fn create_program(&self) -> Result<Program, ()> {
        unsafe {
            let id = gl::CreateProgram();
            if id > 0 {
                Ok(Program { gl_id: id })
            }
            else {
                Err(())
            }
        }
    }

    pub fn attach_shader(&self, program: &mut Program, shader: &Shader) {
        unsafe {
            gl::AttachShader(program.gl_id(), shader.gl_id());
        }
    }

    pub fn link_program(&self, program: &mut Program) -> Result<(), GLError> {
        let success = unsafe {
            gl::LinkProgram(program.gl_id());
            let mut link_status : GLint = 0;
            gl::GetProgramiv(program.gl_id(),
                             gl::LINK_STATUS,
                             &mut link_status as *mut GLint);

            link_status == gl::TRUE as GLint
        };

        if success {
            Ok(())
        }
        else {
            unsafe {
                let mut info_length : GLint = mem::uninitialized();
                gl::GetProgramiv(program.gl_id(),
                                 gl::INFO_LOG_LENGTH,
                                 &mut info_length as *mut GLint);

                let mut bytes = Vec::<u8>::with_capacity(info_length as usize);

                gl::GetProgramInfoLog(program.gl_id(),
                                      info_length,
                                      ptr::null_mut(),
                                      bytes.as_mut_ptr() as *mut GLchar);
                bytes.set_len((info_length - 1) as usize);

                let msg = String::from_utf8(bytes)
                                 .unwrap_or(String::from("<Unknown error>"));

                Err(GLError { message: msg })
            }
        }
    }

    pub fn get_attrib_location(&self, program: &Program, name: &str)
        -> Result<ProgramAttrib, ()>
    {
        let c_str = try!(CString::new(name).or(Err(())));
        let str_ptr = c_str.as_ptr() as *const GLchar;
        unsafe {
            let index = gl::GetAttribLocation(program.gl_id(), str_ptr);
            if index >= 0 {
                Ok(ProgramAttrib { gl_index: index as GLuint })
            }
            else {
                Err(())
            }
        }
    }

    pub fn get_uniform_location(&self, program: &Program, name: &str)
        -> Result<ProgramUniform, ()>
    {
        let c_str = try!(CString::new(name).or(Err(())));
        let str_ptr = c_str.as_ptr() as *const GLchar;
        unsafe {
            let index = gl::GetUniformLocation(program.gl_id(), str_ptr);
            if index >= 0 {
                Ok(ProgramUniform { gl_index: index as GLuint })
            }
            else {
                Err(())
            }
        }
    }

    pub fn use_program(&self, program: &Program) {
        unsafe {
            gl::UseProgram(program.gl_id())
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct ProgramAttrib {
    pub gl_index: GLuint
}

#[derive(Debug, Clone, Copy)]
pub struct ProgramUniform {
    pub gl_index: GLuint
}
<|endoftext|>


<commit_before>use crate::{buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error};
use amq_protocol::frame::{gen_frame, parse_frame, GenError, Offset};
use log::{error, trace};
use mio::{Evented, Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use parking_lot::Mutex;
use std::{
    io::{Read, Write},
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    thread::{self, Builder as ThreadBuilder, JoinHandle},
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const DATA: Token = Token(2);
const CONTINUE: Token = Token(3);

const FRAMES_STORAGE: usize = 32;

type ThreadHandle = JoinHandle<Result<(), Error>>;

#[derive(Clone, Debug)]
pub(crate) struct IoLoopHandle {
    handle: Arc<Mutex<Option<ThreadHandle>>>,
}

impl Default for IoLoopHandle {
    fn default() -> Self {
        Self {
            handle: Arc::new(Mutex::new(None)),
        }
    }
}

impl IoLoopHandle {
    pub(crate) fn register(&self, handle: JoinHandle<Result<(), Error>>) {
        *self.handle.lock() = Some(handle);
    }

    pub(crate) fn wait(&self) -> Result<(), Error> {
        if let Some(handle) = self.handle.lock().take() {
            handle.join().expect("io loop")?
        }
        Ok(())
    }
}

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

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    registration: Registration,
    set_readiness: SetReadiness,
    hb_handle: Option<JoinHandle<()>>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    has_data: bool,
    send_heartbeat: Arc<AtomicBool>,
    poll_timeout: Option<Duration>,
}

impl<T: Evented + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self, Error> {
        let (poll, registered) = poll.map(|t| Ok((t.0, true))).unwrap_or_else(|| {
            Poll::new()
                .map(|poll| (poll, false))
                .map_err(Error::IOError)
        })?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let (registration, set_readiness) = Registration::new2();
        let inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            registration,
            set_readiness,
            hb_handle: None,
            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,
            has_data: false,
            send_heartbeat: Arc::new(AtomicBool::new(false)),
            poll_timeout: None,
        };
        if registered {
            inner
                .poll
                .reregister(&inner.socket, SOCKET, Ready::all(), PollOpt::edge())
                .map_err(Error::IOError)?;
        } else {
            inner
                .poll
                .register(&inner.socket, SOCKET, Ready::all(), PollOpt::edge())
                .map_err(Error::IOError)?;
        }
        inner
            .poll
            .register(&inner.connection, DATA, Ready::readable(), PollOpt::edge())
            .map_err(Error::IOError)?;
        inner
            .poll
            .register(
                &inner.registration,
                CONTINUE,
                Ready::readable(),
                PollOpt::edge(),
            )
            .map_err(Error::IOError)?;
        Ok(inner)
    }

    fn start_heartbeat(&mut self, interval: Duration) -> Result<(), Error> {
        let connection = self.connection.clone();
        let send_hartbeat = self.send_heartbeat.clone();
        let hb_handle = ThreadBuilder::new()
            .name("heartbeat".to_owned())
            .spawn(move || {
                while connection.status().connected() {
                    let start = Instant::now();
                    let mut remaining = interval;

                    loop {
                        thread::park_timeout(remaining);
                        let elapsed = start.elapsed();
                        if elapsed >= remaining {
                            break;
                        }
                        remaining -= interval - elapsed;
                    }

                    send_hartbeat.store(true, Ordering::Relaxed);
                }
            })
            .map_err(Error::IOError)?;
        self.hb_handle = Some(hb_handle);
        Ok(())
    }

    fn heartbeat(&mut self) -> Result<(), Error> {
        trace!("send heartbeat");
        self.connection.send_heartbeat()?;
        self.send_heartbeat.store(false, Ordering::Relaxed);
        Ok(())
    }

    fn ensure_setup(&mut self) -> Result<(), Error> {
        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 {
                trace!("io_loop: start heartbeat");
                let heartbeat = Duration::from_secs(u64::from(heartbeat));
                self.start_heartbeat(heartbeat)?;
                self.poll_timeout = Some(heartbeat);
                trace!("io_loop: heartbeat started");
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    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 run(mut self) -> Result<(), Error> {
        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.do_run(&mut events)?;
                    }
                    if let Some(hb_handle) = self.hb_handle.take() {
                        hb_handle.thread().unpark();
                        hb_handle.join().expect("heartbeat loop failed");
                    }
                    Ok(())
                })
                .map_err(Error::IOError)?,
        );
        Ok(())
    }

    fn do_run(&mut self, events: &mut Events) -> Result<(), Error> {
        // First, update our internal state
        trace!("io_loop run");
        self.ensure_setup()?;
        trace!("io_loop poll");
        self.poll
            .poll(events, self.poll_timeout)
            .map_err(Error::IOError)?;
        trace!("io_loop poll done");
        for event in events.iter() {
            match event.token() {
                SOCKET => {
                    if event.readiness().is_readable() {
                        self.can_read = true;
                    }
                    if event.readiness().is_writable() {
                        self.can_write = true;
                    }
                }
                DATA => self.has_data = true,
                _ => {}
            }
        }

        // Then, actually run an iteration of the event loop
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data
        );
        loop {
            if self.send_heartbeat.load(Ordering::Relaxed) {
                self.heartbeat()?;
            }
            if self.can_write() {
                if let Err(e) = self.write_to_stream() {
                    if e.wouldblock() {
                        self.can_write = false
                    } else {
                        error!("error writing: {:?}", e);
                        if let ConnectionState::SentProtocolHeader(wait_handle, ..) =
                            self.connection.status().state()
                        {
                            wait_handle.error(Error::ConnectionRefused);
                            self.status = Status::Stop;
                        }
                        self.connection.set_error()?;
                        return Err(e);
                    }
                }
                self.send_buffer.shift_unless_available(self.frame_size);
            }
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() && self.can_read() {
                if let Err(e) = self.read_from_stream() {
                    if e.wouldblock() {
                        self.can_read = false
                    } else {
                        error!("error reading: {:?}", e);
                        self.connection.set_error()?;
                        return Err(e);
                    }
                }
                self.receive_buffer.shift_unless_available(self.frame_size);
            }
            if self.can_parse() {
                self.parse()?;
            }
            if !self.can_read()
                || !self.can_write()
                || self.status == Status::Stop
                || self.connection.status().errored()
            {
                if self.status != Status::Stop
                    && (self.can_read() || self.can_parse() || self.can_write())
                {
                    trace!(
                        "io_loop send continue; can_read={}, can_write={}, has_data={}",
                        self.can_read,
                        self.can_write,
                        self.has_data
                    );
                    self.send_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 send_continue(&mut self) -> Result<(), Error> {
        self.set_readiness
            .set_readiness(Ready::readable())
            .map_err(Error::IOError)
    }

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

        self.socket
            .write(&self.send_buffer.data())
            .map(|sz| {
                trace!("wrote {} bytes", sz);
                self.send_buffer.consume(sz);
            })
            .map_err(Error::IOError)
    }

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

    fn serialize(&mut self) -> Result<(), Error> {
        if let Some((send_id, next_msg)) = 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(_) => {
                    self.connection.mark_sent(send_id);
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame(send_id, next_msg)?;
                            self.send_buffer.shift();
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.connection.set_error()?;
                            Err(Error::SerialisationError(e))
                        }
                    }
                }
            }
        } else {
            self.has_data = false;
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<(), Error> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.data().offset(i);
                self.receive_buffer.consume(consumed);

                if let Err(e) = self.connection.handle_frame(f) {
                    self.connection.set_error()?;
                    Err(e)
                } else {
                    Ok(())
                }
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(())
                } else {
                    error!("parse error: {:?}", e);
                    self.connection.set_error()?;
                    Err(Error::ParsingError(format!("{:?}", e)))
                }
            }
        }
    }
}
<commit_msg>io_loop: split parse and handle_frame<commit_after>use crate::{buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error};
use amq_protocol::frame::{gen_frame, parse_frame, GenError, Offset, AMQPFrame};
use log::{error, trace};
use mio::{Evented, Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use parking_lot::Mutex;
use std::{
    io::{Read, Write},
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    thread::{self, Builder as ThreadBuilder, JoinHandle},
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const DATA: Token = Token(2);
const CONTINUE: Token = Token(3);

const FRAMES_STORAGE: usize = 32;

type ThreadHandle = JoinHandle<Result<(), Error>>;

#[derive(Clone, Debug)]
pub(crate) struct IoLoopHandle {
    handle: Arc<Mutex<Option<ThreadHandle>>>,
}

impl Default for IoLoopHandle {
    fn default() -> Self {
        Self {
            handle: Arc::new(Mutex::new(None)),
        }
    }
}

impl IoLoopHandle {
    pub(crate) fn register(&self, handle: JoinHandle<Result<(), Error>>) {
        *self.handle.lock() = Some(handle);
    }

    pub(crate) fn wait(&self) -> Result<(), Error> {
        if let Some(handle) = self.handle.lock().take() {
            handle.join().expect("io loop")?
        }
        Ok(())
    }
}

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

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    registration: Registration,
    set_readiness: SetReadiness,
    hb_handle: Option<JoinHandle<()>>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    has_data: bool,
    send_heartbeat: Arc<AtomicBool>,
    poll_timeout: Option<Duration>,
}

impl<T: Evented + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self, Error> {
        let (poll, registered) = poll.map(|t| Ok((t.0, true))).unwrap_or_else(|| {
            Poll::new()
                .map(|poll| (poll, false))
                .map_err(Error::IOError)
        })?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let (registration, set_readiness) = Registration::new2();
        let inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            registration,
            set_readiness,
            hb_handle: None,
            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,
            has_data: false,
            send_heartbeat: Arc::new(AtomicBool::new(false)),
            poll_timeout: None,
        };
        if registered {
            inner
                .poll
                .reregister(&inner.socket, SOCKET, Ready::all(), PollOpt::edge())
                .map_err(Error::IOError)?;
        } else {
            inner
                .poll
                .register(&inner.socket, SOCKET, Ready::all(), PollOpt::edge())
                .map_err(Error::IOError)?;
        }
        inner
            .poll
            .register(&inner.connection, DATA, Ready::readable(), PollOpt::edge())
            .map_err(Error::IOError)?;
        inner
            .poll
            .register(
                &inner.registration,
                CONTINUE,
                Ready::readable(),
                PollOpt::edge(),
            )
            .map_err(Error::IOError)?;
        Ok(inner)
    }

    fn start_heartbeat(&mut self, interval: Duration) -> Result<(), Error> {
        let connection = self.connection.clone();
        let send_hartbeat = self.send_heartbeat.clone();
        let hb_handle = ThreadBuilder::new()
            .name("heartbeat".to_owned())
            .spawn(move || {
                while connection.status().connected() {
                    let start = Instant::now();
                    let mut remaining = interval;

                    loop {
                        thread::park_timeout(remaining);
                        let elapsed = start.elapsed();
                        if elapsed >= remaining {
                            break;
                        }
                        remaining -= interval - elapsed;
                    }

                    send_hartbeat.store(true, Ordering::Relaxed);
                }
            })
            .map_err(Error::IOError)?;
        self.hb_handle = Some(hb_handle);
        Ok(())
    }

    fn heartbeat(&mut self) -> Result<(), Error> {
        trace!("send heartbeat");
        self.connection.send_heartbeat()?;
        self.send_heartbeat.store(false, Ordering::Relaxed);
        Ok(())
    }

    fn ensure_setup(&mut self) -> Result<(), Error> {
        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 {
                trace!("io_loop: start heartbeat");
                let heartbeat = Duration::from_secs(u64::from(heartbeat));
                self.start_heartbeat(heartbeat)?;
                self.poll_timeout = Some(heartbeat);
                trace!("io_loop: heartbeat started");
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    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 run(mut self) -> Result<(), Error> {
        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.do_run(&mut events)?;
                    }
                    if let Some(hb_handle) = self.hb_handle.take() {
                        hb_handle.thread().unpark();
                        hb_handle.join().expect("heartbeat loop failed");
                    }
                    Ok(())
                })
                .map_err(Error::IOError)?,
        );
        Ok(())
    }

    fn do_run(&mut self, events: &mut Events) -> Result<(), Error> {
        // First, update our internal state
        trace!("io_loop run");
        self.ensure_setup()?;
        trace!("io_loop poll");
        self.poll
            .poll(events, self.poll_timeout)
            .map_err(Error::IOError)?;
        trace!("io_loop poll done");
        for event in events.iter() {
            match event.token() {
                SOCKET => {
                    if event.readiness().is_readable() {
                        self.can_read = true;
                    }
                    if event.readiness().is_writable() {
                        self.can_write = true;
                    }
                }
                DATA => self.has_data = true,
                _ => {}
            }
        }

        // Then, actually run an iteration of the event loop
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data
        );
        loop {
            if self.send_heartbeat.load(Ordering::Relaxed) {
                self.heartbeat()?;
            }
            if self.can_write() {
                if let Err(e) = self.write_to_stream() {
                    if e.wouldblock() {
                        self.can_write = false
                    } else {
                        error!("error writing: {:?}", e);
                        if let ConnectionState::SentProtocolHeader(wait_handle, ..) =
                            self.connection.status().state()
                        {
                            wait_handle.error(Error::ConnectionRefused);
                            self.status = Status::Stop;
                        }
                        self.connection.set_error()?;
                        return Err(e);
                    }
                }
                self.send_buffer.shift_unless_available(self.frame_size);
            }
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() && self.can_read() {
                if let Err(e) = self.read_from_stream() {
                    if e.wouldblock() {
                        self.can_read = false
                    } else {
                        error!("error reading: {:?}", e);
                        self.connection.set_error()?;
                        return Err(e);
                    }
                }
                self.receive_buffer.shift_unless_available(self.frame_size);
            }
            if self.can_parse() {
                if let Some(frame) = self.parse()? {
                    self.handle_frame(frame)?;
                }
            }
            if !self.can_read()
                || !self.can_write()
                || self.status == Status::Stop
                || self.connection.status().errored()
            {
                if self.status != Status::Stop
                    && (self.can_read() || self.can_parse() || self.can_write())
                {
                    trace!(
                        "io_loop send continue; can_read={}, can_write={}, has_data={}",
                        self.can_read,
                        self.can_write,
                        self.has_data
                    );
                    self.send_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 send_continue(&mut self) -> Result<(), Error> {
        self.set_readiness
            .set_readiness(Ready::readable())
            .map_err(Error::IOError)
    }

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

        self.socket
            .write(&self.send_buffer.data())
            .map(|sz| {
                trace!("wrote {} bytes", sz);
                self.send_buffer.consume(sz);
            })
            .map_err(Error::IOError)
    }

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

    fn serialize(&mut self) -> Result<(), Error> {
        if let Some((send_id, next_msg)) = 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(_) => {
                    self.connection.mark_sent(send_id);
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame(send_id, next_msg)?;
                            self.send_buffer.shift();
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.connection.set_error()?;
                            Err(Error::SerialisationError(e))
                        }
                    }
                }
            }
        } else {
            self.has_data = false;
            Ok(())
        }
    }

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

    fn handle_frame(&mut self, frame: AMQPFrame) -> Result<(), Error> {
        if let Err(e) = self.connection.handle_frame(frame) {
            self.connection.set_error()?;
            Err(e)
        } else {
            Ok(())
        }
    }
}
<|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::{self, 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() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&mut self, error: Error) -> Result<()> {
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    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 {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                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.socket.flush()?;
        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()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            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));
                        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.send_continue()?;
            } else {
                self.socket.flush()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        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() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            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>Revert "flush the socket more often"<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::{self, 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() {
                        if let Err(err) = self.run(&mut events) {
                            self.cancel_serilized_frames(err)?;
                        }
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn cancel_serilized_frames(&mut self, error: Error) -> Result<()> {
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        Err(error)
    }

    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 {
                trace!("Got mio event for socket: {:?}", event);
                if event.is_read_closed() || event.is_write_closed() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
                }
                if event.is_error() {
                    self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
                }
                // Due to a bug in epoll/mio, it doesn't seem like we can trust this, it's sometimes missing when it should be there
                /*
                if event.is_readable() {
                    self.can_read = true;
                }
                */
                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()?;
                // Update last_write so that if we cannot write yet to the socket, we don't enqueue countless heartbeats
                self.last_write = Instant::now();
            }
            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));
                        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.send_continue()?;
            } else {
                self.socket.flush()?;
            }
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.can_write = false;
        }
        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() {
            // FIXME: having the body here is too verbose, impl Display on frame instead: trace!("will write to buffer: {:?}", next_msg);
            trace!("Will write to buffer");
            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|>


<commit_before>extern crate hyper;
extern crate serde_json;

use request;

pub use types::{Account, CreateAccountResponse, EditAccountResponse};

impl Account {
    pub fn new(short_name: &str, author_name: &str, author_url: &str) -> Account {
        Account {
            short_name: short_name.to_string(),
            author_name: author_name.to_string(),
            author_url: author_url.to_string(),
            access_token: "".to_string(),
            auth_url: "".to_string(),
            page_count: 0,
        }
    }

    pub fn set_short_name(&mut self, short_name: &str) -> &mut Account {
        self.short_name = short_name.to_string();
        self
    }

    pub fn set_author_name(&mut self, author_name: &str) -> &mut Account {
        self.author_name = author_name.to_string();
        self
    }

    pub fn set_author_url(&mut self, author_url: &str) -> &mut Account {
        self.author_url = author_url.to_string();
        self
    }

    pub fn set_auth_url(&mut self, auth_url: &str) -> &mut Account {
        self.auth_url = auth_url.to_string();
        self
    }

    pub fn set_page_count(&mut self, page_count: &i32) -> &mut Account {
        self.page_count = page_count.clone();
        self
    }

    pub fn set_access_token(&mut self, access_token: &str) -> &mut Account {
        self.access_token = access_token.to_string();
        self
    }

    fn update(&mut self, new_acc: &Account) -> &mut Account {
        self.set_short_name(&new_acc.short_name);
        self.set_author_name(&new_acc.author_name);
        self.set_author_url(&new_acc.author_url);

        self
    }

    pub fn create_account(&mut self) -> Result<CreateAccountResponse, &str> {
        let url = format!(
            "https://api.telegra.ph/createAccount?short_name={}&author_name={}",
            self.short_name,
            self.author_name
        );
        let res_json = request::get(&url);
        let decoded: CreateAccountResponse = serde_json::from_str(&res_json).unwrap();

        self.set_access_token(&decoded.result.access_token)
            .set_auth_url(&decoded.result.auth_url);

        Ok(decoded)
    }

    pub fn edit_accout_info(&mut self, acc: &Account) -> Result<EditAccountResponse, &str> {
        let url = String::from("https://api.telegra.ph/editAccountInfo?");
        let mut params: Vec<String> = vec![];

        let access_token = "access_token=".to_string() + &if acc.access_token.len() == 0 {
                self.access_token.clone()
            } else {
                acc.access_token.clone()
            };
        params.push(access_token);

        if acc.short_name.len() > 0 {
            let short_name = "short_name=".to_string() + &acc.short_name;
            params.push(short_name);
        }
        if acc.author_name.len() > 0 {
            let author_name = "author_name=".to_string() + &acc.author_name;
            params.push(author_name);
        }
        if acc.author_url.len() > 0 {
            let author_url = "author_url=".to_string() + &acc.author_url;
            params.push(author_url);
        }

        let url = url + &params.join("&");
        let res_json = request::get(&url);
        let decoded: EditAccountResponse = serde_json::from_str(&res_json).unwrap();

        self.update(&decoded.result);

        Ok(decoded)
    }

    pub fn get_account_info<'a>(access_token: &String, fields_list: &Vec<&str>) -> Result<EditAccountResponse, &'a str> {
        let fields = fields_list.iter()
            .map(|&field| "\"".to_string() + field + "\"")
            .collect::<Vec<String>>()
            .join(",");
        let fields = "[".to_string() + &fields + "]";
        let url = format!(
            "https://api.telegra.ph/getAccountInfo?access_token={}&fields={}",
            access_token,
            fields
        );
        let res_json = request::get(&url);
        let decoded: EditAccountResponse = serde_json::from_str(&res_json).unwrap();

        Ok(decoded)
    }

    pub fn revoke_access_token(&mut self) -> Result<EditAccountResponse, &str> {
        let url = format!(
            "https://api.telegra.ph/revokeAccessToken?access_token={}",
            self.access_token
        );
        let res_json = request::get(&url);
        let decoded: EditAccountResponse = serde_json::from_str(&res_json).unwrap();

        self.set_access_token(&decoded.result.access_token);

        Ok(decoded)
    }

}
<commit_msg>refactor get_account_info method<commit_after>extern crate hyper;
extern crate serde_json;

use request;

pub use types::{Account, CreateAccountResponse, EditAccountResponse};

impl Account {
    pub fn new(short_name: &str, author_name: &str, author_url: &str) -> Account {
        Account {
            short_name: short_name.to_string(),
            author_name: author_name.to_string(),
            author_url: author_url.to_string(),
            access_token: "".to_string(),
            auth_url: "".to_string(),
            page_count: 0,
        }
    }

    pub fn set_short_name(&mut self, short_name: &str) -> &mut Account {
        self.short_name = short_name.to_string();
        self
    }

    pub fn set_author_name(&mut self, author_name: &str) -> &mut Account {
        self.author_name = author_name.to_string();
        self
    }

    pub fn set_author_url(&mut self, author_url: &str) -> &mut Account {
        self.author_url = author_url.to_string();
        self
    }

    pub fn set_auth_url(&mut self, auth_url: &str) -> &mut Account {
        self.auth_url = auth_url.to_string();
        self
    }

    pub fn set_page_count(&mut self, page_count: &i32) -> &mut Account {
        self.page_count = page_count.clone();
        self
    }

    pub fn set_access_token(&mut self, access_token: &str) -> &mut Account {
        self.access_token = access_token.to_string();
        self
    }

    fn update(&mut self, new_acc: &Account) -> &mut Account {
        self.set_short_name(&new_acc.short_name);
        self.set_author_name(&new_acc.author_name);
        self.set_author_url(&new_acc.author_url);

        self
    }

    pub fn create_account(&mut self) -> Result<CreateAccountResponse, &str> {
        let url = format!(
            "https://api.telegra.ph/createAccount?short_name={}&author_name={}",
            self.short_name,
            self.author_name
        );
        let res_json = request::get(&url);
        let decoded: CreateAccountResponse = serde_json::from_str(&res_json).unwrap();

        self.set_access_token(&decoded.result.access_token)
            .set_auth_url(&decoded.result.auth_url);

        Ok(decoded)
    }

    pub fn edit_accout_info(&mut self, acc: &Account) -> Result<EditAccountResponse, &str> {
        let url = String::from("https://api.telegra.ph/editAccountInfo?");
        let mut params: Vec<String> = vec![];

        let access_token = "access_token=".to_string() + &if acc.access_token.len() == 0 {
                self.access_token.clone()
            } else {
                acc.access_token.clone()
            };
        params.push(access_token);

        if acc.short_name.len() > 0 {
            let short_name = "short_name=".to_string() + &acc.short_name;
            params.push(short_name);
        }
        if acc.author_name.len() > 0 {
            let author_name = "author_name=".to_string() + &acc.author_name;
            params.push(author_name);
        }
        if acc.author_url.len() > 0 {
            let author_url = "author_url=".to_string() + &acc.author_url;
            params.push(author_url);
        }

        let url = url + &params.join("&");
        let res_json = request::get(&url);
        let decoded: EditAccountResponse = serde_json::from_str(&res_json).unwrap();

        self.update(&decoded.result);

        Ok(decoded)
    }

    pub fn get_account_info<'a>(&self, fields_list: &Vec<&str>) -> Result<EditAccountResponse, &'a str> {
        let fields = fields_list.iter()
            .map(|&field| "\"".to_string() + field + "\"")
            .collect::<Vec<String>>()
            .join(",");
        let fields = "[".to_string() + &fields + "]";
        let url = format!(
            "https://api.telegra.ph/getAccountInfo?access_token={}&fields={}",
            self.access_token,
            fields
        );
        let res_json = request::get(&url);
        let decoded: EditAccountResponse = serde_json::from_str(&res_json).unwrap();

        Ok(decoded)
    }

    pub fn revoke_access_token(&mut self) -> Result<EditAccountResponse, &str> {
        let url = format!(
            "https://api.telegra.ph/revokeAccessToken?access_token={}",
            self.access_token
        );
        let res_json = request::get(&url);
        let decoded: EditAccountResponse = serde_json::from_str(&res_json).unwrap();

        self.set_access_token(&decoded.result.access_token);

        Ok(decoded)
    }

}
<|endoftext|>