<commit_before>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError, Offset};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
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 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>,
    hb_handle: Option<JoinHandle<()>>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    send_heartbeat: Arc<AtomicBool>,
    poll_timeout: Option<Duration>,
}

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,
            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,
            send_heartbeat: Arc::new(AtomicBool::new(false)),
            poll_timeout: None,
        };
        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 start_heartbeat(&mut self, interval: Duration) -> Result<()> {
        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);
                }
            })?;
        self.hb_handle = Some(hb_handle);
        Ok(())
    }

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

    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 {
                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 has_data(&self) -> bool {
        self.connection.has_pending_frames() || self.send_buffer.available_data() > 0
    }

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

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

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

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

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    if let Some(hb_handle) = self.hb_handle.take() {
                        hb_handle.thread().unpark();
                        hb_handle.join().expect("heartbeat loop failed");
                    }
                    Ok(())
                })?,
            waker,
        )
    }

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

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

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.heartbeat()?;
            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.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 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)?;
                }
            }
            self.send_buffer.shift_unless_available(self.frame_size);
        }
        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.socket.write(&self.send_buffer.data())?;

        trace!("wrote {} bytes", sz);
        self.send_buffer.consume(sz);
        if sz > 0 && self.send_buffer.available_data() > 0 {
            // We didn't write all the data yet
            self.send_continue()?;
        }
        Ok(())
    }

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

    fn serialize(&mut self) -> Result<()> {
        if let Some((send_id, next_msg, resolver)) = self.connection.next_frame() {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(_) => {
                    if let Some(resolver) = resolver {
                        resolver.swear(Ok(())); // FIXME: do that only once written all
                    }
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection
                                .requeue_frame((send_id, next_msg, resolver))?;
                            self.send_buffer.shift();
                            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.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);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>IoLoop: wait for frame to be 100% sent before notifying as such<commit_after>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error,
    PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError, Offset};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    collections::VecDeque,
    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 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>,
    hb_handle: Option<JoinHandle<()>>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    send_heartbeat: Arc<AtomicBool>,
    poll_timeout: Option<Duration>,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

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,
            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,
            send_heartbeat: Arc::new(AtomicBool::new(false)),
            poll_timeout: None,
            serialized_frames: VecDeque::default(),
        };
        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 start_heartbeat(&mut self, interval: Duration) -> Result<()> {
        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);
                }
            })?;
        self.hb_handle = Some(hb_handle);
        Ok(())
    }

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

    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 {
                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 has_data(&self) -> bool {
        self.connection.has_pending_frames() || self.send_buffer.available_data() > 0
    }

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

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

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

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

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    if let Some(hb_handle) = self.hb_handle.take() {
                        hb_handle.thread().unpark();
                        hb_handle.join().expect("heartbeat loop failed");
                    }
                    Ok(())
                })?,
            waker,
        )
    }

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

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

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.heartbeat()?;
            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.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 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)?;
                }
            }
            self.send_buffer.shift_unless_available(self.frame_size);
        }
        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.socket.write(&self.send_buffer.data())?;

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

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

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

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

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


<commit_before>//! # cli
//!
//! The cargo-make cli
//!

#[cfg(test)]
#[path = "cli_test.rs"]
mod cli_test;

use crate::cli_commands;
use crate::config;
use crate::descriptor;
use crate::environment;
use crate::logger;
use crate::logger::LoggerOptions;
use crate::profile;
use crate::recursion_level;
use crate::runner;
use crate::types::{CliArgs, GlobalConfig};
use crate::version;
use clap::{App, Arg, ArgMatches, SubCommand};

static VERSION: &str = env!("CARGO_PKG_VERSION");
static AUTHOR: &str = env!("CARGO_PKG_AUTHORS");
static DESCRIPTION: &str = env!("CARGO_PKG_DESCRIPTION");
static DEFAULT_TOML: &str = "Makefile.toml";
static DEFAULT_LOG_LEVEL: &str = "info";
static DEFAULT_TASK_NAME: &str = "default";
static DEFAULT_OUTPUT_FORMAT: &str = "default";

fn run(cli_args: CliArgs, global_config: &GlobalConfig) {
    recursion_level::increment();

    logger::init(&LoggerOptions {
        level: cli_args.log_level.clone(),
        color: !cli_args.disable_color,
    });

    if recursion_level::is_top() {
        info!("{} {}", &cli_args.command, &VERSION);
        debug!("Written By {}", &AUTHOR);
    }

    debug!("Cli Args {:#?}", &cli_args);
    debug!("Global Configuration {:#?}", &global_config);

    // only run check for updates if we are not in a CI env and user didn't ask to skip the check
    if !cli_args.disable_check_for_updates
        && !ci_info::is_ci()
        && version::should_check(&global_config)
    {
        version::check();
    }

    let cwd_string_option = match cli_args.cwd.clone() {
        Some(value) => Some(value),
        None => match global_config.search_project_root {
            Some(search) => {
                if search {
                    match environment::get_project_root() {
                        Some(value) => Some(value.clone()),
                        None => None,
                    }
                } else {
                    None
                }
            }
            None => None,
        },
    };
    let cwd = match cwd_string_option {
        Some(ref value) => Some(value.as_ref()),
        None => None,
    };
    let home = environment::setup_cwd(cwd);

    let force_makefile = cli_args.build_file.is_some();
    let build_file = &cli_args
        .build_file
        .clone()
        .unwrap_or(DEFAULT_TOML.to_string());
    let task = &cli_args.task;
    let profile_name = &cli_args
        .profile
        .clone()
        .unwrap_or(profile::DEFAULT_PROFILE.to_string());
    let normalized_profile_name = profile::set(&profile_name);

    environment::load_env_file(cli_args.env_file.clone());

    let env = cli_args.env.clone();

    let experimental = cli_args.experimental;
    let descriptor_load_result = descriptor::load(&build_file, force_makefile, env, experimental);

    let config = match descriptor_load_result {
        Ok(config) => config,
        Err(ref min_version) => {
            error!(
                "{} version: {} does not meet minimum required version: {}",
                &cli_args.command, &VERSION, &min_version
            );
            panic!(
                "{} version: {} does not meet minimum required version: {}",
                &cli_args.command, &VERSION, &min_version
            );
        }
    };

    match config.config.additional_profiles {
        Some(ref profiles) => profile::set_additional(profiles),
        None => profile::set_additional(&vec![]),
    };

    let env_info = environment::setup_env(&cli_args, &config, &task, home);

    let crate_name = envmnt::get_or("CARGO_MAKE_CRATE_NAME", "");
    if crate_name.len() > 0 {
        info!("Project: {}", &crate_name);
    }
    info!("Build File: {}", &build_file);
    info!("Task: {}", &task);
    info!("Profile: {}", &normalized_profile_name);

    // ensure profile env was not overridden
    profile::set(&normalized_profile_name);

    if cli_args.list_all_steps {
        cli_commands::list_steps::run(&config, &cli_args.output_format, &cli_args.output_file);
    } else if cli_args.diff_execution_plan {
        let default_config = descriptor::load_internal_descriptors(true, experimental, None);
        cli_commands::diff_steps::run(&default_config, &config, &task, &cli_args);
    } else if cli_args.print_only {
        cli_commands::print_steps::print(
            &config,
            &task,
            &cli_args.output_format,
            cli_args.disable_workspace,
        );
    } else {
        runner::run(config, &task, env_info, &cli_args);
    }
}

/// Handles the command line arguments and executes the runner.
fn run_for_args(
    matches: ArgMatches,
    global_config: &GlobalConfig,
    command_name: &String,
    sub_command: bool,
) {
    let cmd_matches = if sub_command {
        match matches.subcommand_matches(command_name) {
            Some(value) => value,
            None => panic!("cargo-{} not invoked via cargo command.", &command_name),
        }
    } else {
        &matches
    };

    let mut cli_args = CliArgs::new();

    cli_args.command = if sub_command {
        let mut binary = "cargo ".to_string();
        binary.push_str(&command_name);
        binary
    } else {
        command_name.clone()
    };

    cli_args.env = cmd_matches.values_of_lossy("env");

    cli_args.build_file = if cmd_matches.occurrences_of("makefile") == 0 {
        None
    } else {
        let makefile = cmd_matches
            .value_of("makefile")
            .unwrap_or(&DEFAULT_TOML)
            .to_string();
        Some(makefile)
    };

    cli_args.cwd = match cmd_matches.value_of("cwd") {
        Some(value) => Some(value.to_string()),
        None => None,
    };

    let default_log_level = match global_config.log_level {
        Some(ref value) => value.as_str(),
        None => &DEFAULT_LOG_LEVEL,
    };
    cli_args.log_level = if cmd_matches.is_present("v") {
        "verbose".to_string()
    } else {
        cmd_matches
            .value_of("loglevel")
            .unwrap_or(default_log_level)
            .to_string()
    };

    let default_disable_color = match global_config.disable_color {
        Some(value) => value,
        None => false,
    };
    cli_args.disable_color = cmd_matches.is_present("no-color")
        || envmnt::is("CARGO_MAKE_DISABLE_COLOR")
        || default_disable_color;

    cli_args.print_time_summary =
        cmd_matches.is_present("time-summary") || envmnt::is("CARGO_MAKE_PRINT_TIME_SUMMARY");

    cli_args.env_file = match cmd_matches.value_of("envfile") {
        Some(value) => Some(value.to_string()),
        None => None,
    };

    cli_args.output_format = cmd_matches
        .value_of("output-format")
        .unwrap_or(DEFAULT_OUTPUT_FORMAT)
        .to_string();

    cli_args.output_file = match cmd_matches.value_of("output_file") {
        Some(value) => Some(value.to_string()),
        None => None,
    };

    let profile_name = cmd_matches
        .value_of("profile".to_string())
        .unwrap_or(profile::DEFAULT_PROFILE);
    cli_args.profile = Some(profile_name.to_string());

    cli_args.disable_check_for_updates = cmd_matches.is_present("disable-check-for-updates");
    cli_args.experimental = cmd_matches.is_present("experimental");
    cli_args.print_only = cmd_matches.is_present("print-steps");
    cli_args.disable_workspace = cmd_matches.is_present("no-workspace");
    cli_args.disable_on_error = cmd_matches.is_present("no-on-error");
    cli_args.allow_private = cmd_matches.is_present("allow-private");
    cli_args.skip_init_end_tasks = cmd_matches.is_present("skip-init-end-tasks");
    cli_args.list_all_steps = cmd_matches.is_present("list-steps");
    cli_args.diff_execution_plan = cmd_matches.is_present("diff-steps");

    let default_task_name = match global_config.default_task_name {
        Some(ref value) => value.as_str(),
        None => &DEFAULT_TASK_NAME,
    };
    let task = cmd_matches.value_of("task").unwrap_or(default_task_name);
    cli_args.task = cmd_matches.value_of("TASK").unwrap_or(task).to_string();

    cli_args.arguments = match cmd_matches.values_of("TASK_ARGS") {
        Some(values) => {
            let args_str: Vec<&str> = values.collect();
            let args_strings = args_str.iter().map(|item| item.to_string()).collect();
            Some(args_strings)
        }
        None => None,
    };

    run(cli_args, global_config);
}

fn create_cli<'a, 'b>(
    global_config: &'a GlobalConfig,
    command_name: &String,
    sub_command: bool,
) -> App<'a, 'b> {
    let default_task_name = match global_config.default_task_name {
        Some(ref value) => value.as_str(),
        None => &DEFAULT_TASK_NAME,
    };
    let default_log_level = match global_config.log_level {
        Some(ref value) => value.as_str(),
        None => &DEFAULT_LOG_LEVEL,
    };

    let mut cli_app = if sub_command {
        SubCommand::with_name(&command_name)
    } else {
        let name = command_name.as_str();
        App::new(name).bin_name(name)
    };

    cli_app = cli_app
        .version(VERSION)
        .author(AUTHOR)
        .about(DESCRIPTION)
        .arg(
            Arg::with_name("makefile")
                .long("--makefile")
                .value_name("FILE")
                .help("The optional toml file containing the tasks definitions")
                .default_value(&DEFAULT_TOML),
        )
        .arg(
            Arg::with_name("task")
                .short("-t")
                .long("--task")
                .value_name("TASK")
                .help(
                    "The task name to execute \
                     (can omit the flag if the task name is the last argument)",
                )
                .default_value(default_task_name),
        )
        .arg(
            Arg::with_name("profile")
                .short("-p")
                .long("--profile")
                .value_name("PROFILE")
                .help(
                    "The profile name (will be converted to lower case)",
                )
                .default_value(&profile::DEFAULT_PROFILE),
        )
        .arg(
            Arg::with_name("cwd")
                .long("--cwd")
                .value_name("DIRECTORY")
                .help(
                    "Will set the current working directory. \
                     The search for the makefile will be from this directory if defined.",
                ),
        )
        .arg(Arg::with_name("no-workspace").long("--no-workspace").help(
            "Disable workspace support (tasks are triggered on workspace and not on members)",
        ))
        .arg(
            Arg::with_name("no-on-error")
                .long("--no-on-error")
                .help("Disable on error flow even if defined in config sections"),
        )
        .arg(
            Arg::with_name("allow-private")
                .long("--allow-private")
                .help("Allow invocation of private tasks"),
        )
        .arg(
            Arg::with_name("skip-init-end-tasks")
                .long("--skip-init-end-tasks")
                .help("If set, init and end tasks are skipped"),
        )
        .arg(
            Arg::with_name("envfile")
                .long("--env-file")
                .value_name("FILE")
                .help("Set environment variables from provided file"),
        )
        .arg(
            Arg::with_name("env")
                .long("--env")
                .short("-e")
                .value_name("ENV")
                .multiple(true)
                .takes_value(true)
                .number_of_values(1)
                .help("Set environment variables"),
        )
        .arg(
            Arg::from_usage("-l, --loglevel=[LOG LEVEL] 'The log level'")
                .possible_values(&["verbose", "info", "error"])
                .default_value(default_log_level),
        )
        .arg(
            Arg::with_name("v")
                .short("-v")
                .long("--verbose")
                .help("Sets the log level to verbose (shorthand for --loglevel verbose)"),
        )
        .arg(
            Arg::with_name("no-color")
                .long("--no-color")
                .help("Disables colorful output"),
        )
        .arg(
            Arg::with_name("time-summary")
                .long("--time-summary")
                .help("Print task level time summary at end of flow"),
        )
        .arg(
            Arg::with_name("experimental")
                .long("--experimental")
                .help("Allows access unsupported experimental predefined tasks."),
        )
        .arg(
            Arg::with_name("disable-check-for-updates")
                .long("--disable-check-for-updates")
                .help("Disables the update check during startup"),
        )
        .arg(
            Arg::from_usage("--output-format=[OUTPUT FORMAT] 'The print/list steps format (some operations do not support all formats)'")
                .possible_values(&["default", "short-description", "markdown", "markdown-single-page", "markdown-sub-section"])
                .default_value(DEFAULT_OUTPUT_FORMAT),
        )
        .arg(
            Arg::with_name("output_file")
                .long("--output-file")
                .value_name("OUTPUT_FILE")
                .help("The list steps output file name"),
        )
        .arg(Arg::with_name("print-steps").long("--print-steps").help(
            "Only prints the steps of the build in the order they will \
             be invoked but without invoking them",
        ))
        .arg(
            Arg::with_name("list-steps")
                .long("--list-all-steps")
                .help("Lists all known steps"),
        )
        .arg(
            Arg::with_name("diff-steps")
                .long("--diff-steps")
                .help("Runs diff between custom flow and prebuilt flow (requires git)"),
        )
        .arg(Arg::with_name("TASK").help("The task name to execute"))
        .arg(
            Arg::with_name("TASK_ARGS")
                .multiple(true)
                .help("Task arguments which can be accessed in the task itself."),
        );

    if sub_command {
        App::new("cargo").bin_name("cargo").subcommand(cli_app)
    } else {
        cli_app
    }
}

/// Handles the command line arguments and executes the runner.
pub(crate) fn run_cli(command_name: String, sub_command: bool) {
    let global_config = config::load();

    let app = create_cli(&global_config, &command_name, sub_command);

    let matches = app.get_matches();

    run_for_args(matches, &global_config, &command_name, sub_command);
}
<commit_msg>feat: Allow leading hyphen for arguments<commit_after>//! # cli
//!
//! The cargo-make cli
//!

#[cfg(test)]
#[path = "cli_test.rs"]
mod cli_test;

use crate::cli_commands;
use crate::config;
use crate::descriptor;
use crate::environment;
use crate::logger;
use crate::logger::LoggerOptions;
use crate::profile;
use crate::recursion_level;
use crate::runner;
use crate::types::{CliArgs, GlobalConfig};
use crate::version;
use clap::{App, AppSettings, Arg, ArgMatches, SubCommand};

static VERSION: &str = env!("CARGO_PKG_VERSION");
static AUTHOR: &str = env!("CARGO_PKG_AUTHORS");
static DESCRIPTION: &str = env!("CARGO_PKG_DESCRIPTION");
static DEFAULT_TOML: &str = "Makefile.toml";
static DEFAULT_LOG_LEVEL: &str = "info";
static DEFAULT_TASK_NAME: &str = "default";
static DEFAULT_OUTPUT_FORMAT: &str = "default";

fn run(cli_args: CliArgs, global_config: &GlobalConfig) {
    recursion_level::increment();

    logger::init(&LoggerOptions {
        level: cli_args.log_level.clone(),
        color: !cli_args.disable_color,
    });

    if recursion_level::is_top() {
        info!("{} {}", &cli_args.command, &VERSION);
        debug!("Written By {}", &AUTHOR);
    }

    debug!("Cli Args {:#?}", &cli_args);
    debug!("Global Configuration {:#?}", &global_config);

    // only run check for updates if we are not in a CI env and user didn't ask to skip the check
    if !cli_args.disable_check_for_updates
        && !ci_info::is_ci()
        && version::should_check(&global_config)
    {
        version::check();
    }

    let cwd_string_option = match cli_args.cwd.clone() {
        Some(value) => Some(value),
        None => match global_config.search_project_root {
            Some(search) => {
                if search {
                    match environment::get_project_root() {
                        Some(value) => Some(value.clone()),
                        None => None,
                    }
                } else {
                    None
                }
            }
            None => None,
        },
    };
    let cwd = match cwd_string_option {
        Some(ref value) => Some(value.as_ref()),
        None => None,
    };
    let home = environment::setup_cwd(cwd);

    let force_makefile = cli_args.build_file.is_some();
    let build_file = &cli_args
        .build_file
        .clone()
        .unwrap_or(DEFAULT_TOML.to_string());
    let task = &cli_args.task;
    let profile_name = &cli_args
        .profile
        .clone()
        .unwrap_or(profile::DEFAULT_PROFILE.to_string());
    let normalized_profile_name = profile::set(&profile_name);

    environment::load_env_file(cli_args.env_file.clone());

    let env = cli_args.env.clone();

    let experimental = cli_args.experimental;
    let descriptor_load_result = descriptor::load(&build_file, force_makefile, env, experimental);

    let config = match descriptor_load_result {
        Ok(config) => config,
        Err(ref min_version) => {
            error!(
                "{} version: {} does not meet minimum required version: {}",
                &cli_args.command, &VERSION, &min_version
            );
            panic!(
                "{} version: {} does not meet minimum required version: {}",
                &cli_args.command, &VERSION, &min_version
            );
        }
    };

    match config.config.additional_profiles {
        Some(ref profiles) => profile::set_additional(profiles),
        None => profile::set_additional(&vec![]),
    };

    let env_info = environment::setup_env(&cli_args, &config, &task, home);

    let crate_name = envmnt::get_or("CARGO_MAKE_CRATE_NAME", "");
    if crate_name.len() > 0 {
        info!("Project: {}", &crate_name);
    }
    info!("Build File: {}", &build_file);
    info!("Task: {}", &task);
    info!("Profile: {}", &normalized_profile_name);

    // ensure profile env was not overridden
    profile::set(&normalized_profile_name);

    if cli_args.list_all_steps {
        cli_commands::list_steps::run(&config, &cli_args.output_format, &cli_args.output_file);
    } else if cli_args.diff_execution_plan {
        let default_config = descriptor::load_internal_descriptors(true, experimental, None);
        cli_commands::diff_steps::run(&default_config, &config, &task, &cli_args);
    } else if cli_args.print_only {
        cli_commands::print_steps::print(
            &config,
            &task,
            &cli_args.output_format,
            cli_args.disable_workspace,
        );
    } else {
        runner::run(config, &task, env_info, &cli_args);
    }
}

/// Handles the command line arguments and executes the runner.
fn run_for_args(
    matches: ArgMatches,
    global_config: &GlobalConfig,
    command_name: &String,
    sub_command: bool,
) {
    let cmd_matches = if sub_command {
        match matches.subcommand_matches(command_name) {
            Some(value) => value,
            None => panic!("cargo-{} not invoked via cargo command.", &command_name),
        }
    } else {
        &matches
    };

    let mut cli_args = CliArgs::new();

    cli_args.command = if sub_command {
        let mut binary = "cargo ".to_string();
        binary.push_str(&command_name);
        binary
    } else {
        command_name.clone()
    };

    cli_args.env = cmd_matches.values_of_lossy("env");

    cli_args.build_file = if cmd_matches.occurrences_of("makefile") == 0 {
        None
    } else {
        let makefile = cmd_matches
            .value_of("makefile")
            .unwrap_or(&DEFAULT_TOML)
            .to_string();
        Some(makefile)
    };

    cli_args.cwd = match cmd_matches.value_of("cwd") {
        Some(value) => Some(value.to_string()),
        None => None,
    };

    let default_log_level = match global_config.log_level {
        Some(ref value) => value.as_str(),
        None => &DEFAULT_LOG_LEVEL,
    };
    cli_args.log_level = if cmd_matches.is_present("v") {
        "verbose".to_string()
    } else {
        cmd_matches
            .value_of("loglevel")
            .unwrap_or(default_log_level)
            .to_string()
    };

    let default_disable_color = match global_config.disable_color {
        Some(value) => value,
        None => false,
    };
    cli_args.disable_color = cmd_matches.is_present("no-color")
        || envmnt::is("CARGO_MAKE_DISABLE_COLOR")
        || default_disable_color;

    cli_args.print_time_summary =
        cmd_matches.is_present("time-summary") || envmnt::is("CARGO_MAKE_PRINT_TIME_SUMMARY");

    cli_args.env_file = match cmd_matches.value_of("envfile") {
        Some(value) => Some(value.to_string()),
        None => None,
    };

    cli_args.output_format = cmd_matches
        .value_of("output-format")
        .unwrap_or(DEFAULT_OUTPUT_FORMAT)
        .to_string();

    cli_args.output_file = match cmd_matches.value_of("output_file") {
        Some(value) => Some(value.to_string()),
        None => None,
    };

    let profile_name = cmd_matches
        .value_of("profile".to_string())
        .unwrap_or(profile::DEFAULT_PROFILE);
    cli_args.profile = Some(profile_name.to_string());

    cli_args.disable_check_for_updates = cmd_matches.is_present("disable-check-for-updates");
    cli_args.experimental = cmd_matches.is_present("experimental");
    cli_args.print_only = cmd_matches.is_present("print-steps");
    cli_args.disable_workspace = cmd_matches.is_present("no-workspace");
    cli_args.disable_on_error = cmd_matches.is_present("no-on-error");
    cli_args.allow_private = cmd_matches.is_present("allow-private");
    cli_args.skip_init_end_tasks = cmd_matches.is_present("skip-init-end-tasks");
    cli_args.list_all_steps = cmd_matches.is_present("list-steps");
    cli_args.diff_execution_plan = cmd_matches.is_present("diff-steps");

    let default_task_name = match global_config.default_task_name {
        Some(ref value) => value.as_str(),
        None => &DEFAULT_TASK_NAME,
    };
    let task = cmd_matches.value_of("task").unwrap_or(default_task_name);
    cli_args.task = cmd_matches.value_of("TASK").unwrap_or(task).to_string();

    cli_args.arguments = match cmd_matches.values_of("TASK_ARGS") {
        Some(values) => {
            let args_str: Vec<&str> = values.collect();
            let args_strings = args_str.iter().map(|item| item.to_string()).collect();
            Some(args_strings)
        }
        None => None,
    };

    run(cli_args, global_config);
}

fn create_cli<'a, 'b>(
    global_config: &'a GlobalConfig,
    command_name: &String,
    sub_command: bool,
) -> App<'a, 'b> {
    let default_task_name = match global_config.default_task_name {
        Some(ref value) => value.as_str(),
        None => &DEFAULT_TASK_NAME,
    };
    let default_log_level = match global_config.log_level {
        Some(ref value) => value.as_str(),
        None => &DEFAULT_LOG_LEVEL,
    };

    let mut cli_app = if sub_command {
        SubCommand::with_name(&command_name)
    } else {
        let name = command_name.as_str();
        App::new(name).bin_name(name)
    };

    cli_app = cli_app
        .version(VERSION)
        .author(AUTHOR)
        .about(DESCRIPTION)
        .setting(AppSettings::AllowLeadingHyphen)
        .arg(
            Arg::with_name("makefile")
                .long("--makefile")
                .value_name("FILE")
                .help("The optional toml file containing the tasks definitions")
                .default_value(&DEFAULT_TOML),
        )
        .arg(
            Arg::with_name("task")
                .short("-t")
                .long("--task")
                .value_name("TASK")
                .help(
                    "The task name to execute \
                     (can omit the flag if the task name is the last argument)",
                )
                .default_value(default_task_name),
        )
        .arg(
            Arg::with_name("profile")
                .short("-p")
                .long("--profile")
                .value_name("PROFILE")
                .help(
                    "The profile name (will be converted to lower case)",
                )
                .default_value(&profile::DEFAULT_PROFILE),
        )
        .arg(
            Arg::with_name("cwd")
                .long("--cwd")
                .value_name("DIRECTORY")
                .help(
                    "Will set the current working directory. \
                     The search for the makefile will be from this directory if defined.",
                ),
        )
        .arg(Arg::with_name("no-workspace").long("--no-workspace").help(
            "Disable workspace support (tasks are triggered on workspace and not on members)",
        ))
        .arg(
            Arg::with_name("no-on-error")
                .long("--no-on-error")
                .help("Disable on error flow even if defined in config sections"),
        )
        .arg(
            Arg::with_name("allow-private")
                .long("--allow-private")
                .help("Allow invocation of private tasks"),
        )
        .arg(
            Arg::with_name("skip-init-end-tasks")
                .long("--skip-init-end-tasks")
                .help("If set, init and end tasks are skipped"),
        )
        .arg(
            Arg::with_name("envfile")
                .long("--env-file")
                .value_name("FILE")
                .help("Set environment variables from provided file"),
        )
        .arg(
            Arg::with_name("env")
                .long("--env")
                .short("-e")
                .value_name("ENV")
                .multiple(true)
                .takes_value(true)
                .number_of_values(1)
                .help("Set environment variables"),
        )
        .arg(
            Arg::from_usage("-l, --loglevel=[LOG LEVEL] 'The log level'")
                .possible_values(&["verbose", "info", "error"])
                .default_value(default_log_level),
        )
        .arg(
            Arg::with_name("v")
                .short("-v")
                .long("--verbose")
                .help("Sets the log level to verbose (shorthand for --loglevel verbose)"),
        )
        .arg(
            Arg::with_name("no-color")
                .long("--no-color")
                .help("Disables colorful output"),
        )
        .arg(
            Arg::with_name("time-summary")
                .long("--time-summary")
                .help("Print task level time summary at end of flow"),
        )
        .arg(
            Arg::with_name("experimental")
                .long("--experimental")
                .help("Allows access unsupported experimental predefined tasks."),
        )
        .arg(
            Arg::with_name("disable-check-for-updates")
                .long("--disable-check-for-updates")
                .help("Disables the update check during startup"),
        )
        .arg(
            Arg::from_usage("--output-format=[OUTPUT FORMAT] 'The print/list steps format (some operations do not support all formats)'")
                .possible_values(&["default", "short-description", "markdown", "markdown-single-page", "markdown-sub-section"])
                .default_value(DEFAULT_OUTPUT_FORMAT),
        )
        .arg(
            Arg::with_name("output_file")
                .long("--output-file")
                .value_name("OUTPUT_FILE")
                .help("The list steps output file name"),
        )
        .arg(Arg::with_name("print-steps").long("--print-steps").help(
            "Only prints the steps of the build in the order they will \
             be invoked but without invoking them",
        ))
        .arg(
            Arg::with_name("list-steps")
                .long("--list-all-steps")
                .help("Lists all known steps"),
        )
        .arg(
            Arg::with_name("diff-steps")
                .long("--diff-steps")
                .help("Runs diff between custom flow and prebuilt flow (requires git)"),
        )
        .arg(Arg::with_name("TASK").help("The task name to execute"))
        .arg(
            Arg::with_name("TASK_ARGS")
                .multiple(true)
                .help("Task arguments which can be accessed in the task itself."),
        );

    if sub_command {
        App::new("cargo").bin_name("cargo").subcommand(cli_app)
    } else {
        cli_app
    }
}

/// Handles the command line arguments and executes the runner.
pub(crate) fn run_cli(command_name: String, sub_command: bool) {
    let global_config = config::load();

    let app = create_cli(&global_config, &command_name, sub_command);

    let matches = app.get_matches();

    run_for_args(matches, &global_config, &command_name, sub_command);
}
<|endoftext|>


<commit_before>//! Lua functionality

use hlua;
use hlua::{Lua, LuaError};
use hlua::any::AnyLuaValue;

use std::thread;

use std::fs::{File};
use std::path::Path;
use std::io::Write;

use std::sync::{Mutex, RwLock};
use std::sync::mpsc::{channel, Sender, Receiver};

#[macro_use]
mod funcs;
#[cfg(test)]
mod tests;

lazy_static! {
    /// Sends requests to the lua thread
    static ref SENDER: Mutex<Option<Sender<LuaQuery>>> = Mutex::new(None);

    /// Receives data back from the lua thread
    /// This should only be accessed by the lua thread itself.
    static ref RECEIVER: Mutex<Option<Receiver<LuaResponse>>> = Mutex::new(None);

    /// Whether the lua thread is currently running
    pub static ref RUNNING: RwLock<bool> = RwLock::new(false);
}

/// Messages sent to the lua thread
pub enum LuaQuery {
    /// Halt the lua thread
    Terminate,
    // Restart the lua thread
    Restart,
    /// Execute a string
    Execute(String),
    /// Execute a file
    ExecuteFile(String),
    /// Get a variable
    GetVariable(String),
    /// Set a value
    SetValue {
        name: Box<::std::borrow::Borrow<str> + Sized>,
        val: Box<hlua::Push<&'static mut Lua<'static>> + Sized>
    },
    /// Create a new array
    EmptyArray(String),
    /// Message to ping the lua thread
    Ping,
    /// Unused send type
    Unused,
}

/// Messages received from lua thread
pub enum LuaResponse {
    /// Lua variable obtained
    Variable(Option<AnyLuaValue>),
    /// Lua error
    Error(hlua::LuaError),
    /// A function is returned
    Function(hlua::functions_read::LuaFunction<String>),
    /// Pong response from lua ping
    Pong,
    /// Unused response type
    Unused,
}

unsafe impl Send for LuaQuery { }
unsafe impl Send for LuaResponse { }
unsafe impl Sync for LuaQuery { }
unsafe impl Sync for LuaResponse { }

/// Whether the lua thread is currently available
pub fn thread_running() -> bool {
    *RUNNING.read().unwrap()
}

/// Errors which may arise from attempting
/// to sending a message to the lua thread.
#[derive(Debug)]
pub enum LuaSendError {
    /// The thread crashed, was shut down, or rebooted.
    ThreadClosed,
    /// The thread has not been initialized yet (maybe not used)
    ThreadUninitialized,
    /// The sender had an issue, most likey because the thread panicked
    Sender
}

/// Attemps to send a LuaQuery to the lua thread.
pub fn try_send(query: LuaQuery) -> Result<(), LuaSendError> {
    if !thread_running() {
        Err(LuaSendError::ThreadClosed)
    }
    else if let Some(ref sender) = *SENDER.lock().unwrap() {
        match sender.send(query) {
            Ok(_) => Ok(()),
            Err(_) => Err(LuaSendError::Sender)
        }
    }
    else {
        Err(LuaSendError::ThreadUninitialized)
    }
}

/// Initialize the lua thread
pub fn init() {
    trace!("Initializing...");
    let (query_tx, query_rx) = channel::<LuaQuery>();
    let (answer_tx, answer_rx) = channel::<LuaResponse>();
    {
        let mut sender = SENDER.lock().unwrap();
        let mut receiver = RECEIVER.lock().unwrap();

        *sender = Some(query_tx);
        *receiver = Some(answer_rx);
    }

    thread::spawn(move || {
        thread_init(answer_tx, query_rx);
    });
    trace!("Created thread. Init finished.");
}

fn thread_init(sender: Sender<LuaResponse>, receiver: Receiver<LuaQuery>) {
    trace!("thread: initializing.");
    let mut lua = Lua::new();
    //unsafe {
    //    hlua_ffi::lua_atpanic(&mut lua.as_mut_lua().0, thread_on_panic);
    //}
    debug!("thread: Loading Lua libraries...");
    lua.openlibs();
    trace!("thread: Loading way-cooler lua extensions...");
    // We should have some good file handling, read files from /usr by default,
    // but for now we're reading directly from the source.
    lua.execute_from_reader::<(), File>(
        File::open("lib/lua/init.lua").unwrap()
    ).unwrap();
    trace!("thread: loading way-cooler libraries...");
    funcs::register_libraries(&mut lua);
    // Only ready after loading libs
    *RUNNING.write().unwrap() = true;
    debug!("thread: entering main loop...");
    thread_main_loop(sender, receiver, &mut lua);
}

fn thread_main_loop(sender: Sender<LuaResponse>, receiver: Receiver<LuaQuery>,
                    lua: &mut Lua) {
    loop {
        let request = receiver.recv();
        match request {
            Err(e) => {
                error!("Lua thread: unable to receive message: {}", e);
                error!("Lua thread: now panicking!");
                *RUNNING.write().unwrap() = false;

                panic!("Lua thread: lost contact with host, exiting!");
            }
            Ok(message) => {
                trace!("Handling a request");
                thread_handle_message(&sender, message, lua);
            }
        }
    }
}

fn thread_handle_message(sender: &Sender<LuaResponse>,
                         request: LuaQuery, lua: &mut Lua) {
    match request {
        LuaQuery::Terminate => {
            trace!("thread: Received terminate signal");
            *RUNNING.write().unwrap() = false;

            info!("thread: Lua thread terminating!");
            return;
        },

        LuaQuery::Restart => {
            trace!("thread: Received restart signal!");
            error!("thread: Lua thread restart not supported!");

            *RUNNING.write().unwrap() = false;

            panic!("Lua thread: Restart not supported!");
        },

        LuaQuery::Execute(code) => {
            trace!("thread: Received request to execute code");
            trace!("thread: Executing {:?}", code);

            match lua.execute::<()>(&code) {
                Err(error) => {
                    warn!("thread: Error executing code: {:?}", error);
                    let response = LuaResponse::Error(error);

                    thread_send(&sender, response);
                }
                Ok(_) => {
                    // This is gonna be really spammy one day
                    trace!("thread: Code executed okay.");
                }
            }
        },

        LuaQuery::ExecuteFile(name) => {
            trace!("thread: Received request to execute file {}", name);
            info!("thread: Executing {}", name);

            let path = Path::new(&name);
            let try_file = File::open(path);

            if let Ok(file) = try_file {
                let result = lua.execute_from_reader::<(), File>(file);
                if let Err(err) = result {
                    warn!("thread: Error executing {}!", name);

                    thread_send(&sender, LuaResponse::Error(err));
                }
                else {
                    trace!("thread: Execution of {} successful.", name);
                }
            }
            else { // Could not open file
                // Unwrap_err is used because we're in the else of let
                let read_error =
                    LuaError::ReadError(try_file.unwrap_err());

                thread_send(&sender, LuaResponse::Error(read_error));
            }
        },

        LuaQuery::GetVariable(varname) => {
            trace!("thread: Received request to get variable {}", varname);
            let var_result = lua.get(varname.as_str());

            match var_result {
                Some(var) => {
                    thread_send(&sender, LuaResponse::Variable(Some(var)));
                }
                None => {
                    warn!("thread: Unable to get variable {}", varname);

                    thread_send(&sender, LuaResponse::Variable(None));
                }
            }
        },

        LuaQuery::SetValue { name: _name, val: _val } => {
            panic!("thread: unimplemented LuaQuery::SetValue!");
        },

        LuaQuery::EmptyArray(_name) => {
            panic!("thread: unimplemented LuaQuery::EmptyArray!");
        },

        _ => {
            panic!("Unimplemented send type for lua thread!");
        }
    }
}

fn thread_send(sender: &Sender<LuaResponse>, response: LuaResponse) {
    trace!("Called thread_send");
    match sender.send(response) {
        Err(_) => {
            error!("thread: Unable to broadcast response!");
            error!("thread: Shutting down in response to inability \
                    to continue!");
            panic!("Lua thread unable to communicate with main thread, \
                    shutting down!");
        }
        Ok(_) => {}
    }
}
<commit_msg>Improve LuaQuery enum<commit_after>//! Lua functionality

use hlua;
use hlua::{Lua, LuaError};
use hlua::any::AnyLuaValue;

use rustc_serialize::json::Json;

use std::thread;

use std::fs::{File};
use std::path::Path;
use std::io::Write;

use std::sync::{Mutex, RwLock};
use std::sync::mpsc::{channel, Sender, Receiver};

#[macro_use]
mod funcs;
#[cfg(test)]
mod tests;

lazy_static! {
    /// Sends requests to the lua thread
    static ref SENDER: Mutex<Option<Sender<LuaQuery>>> = Mutex::new(None);

    /// Receives data back from the lua thread
    /// This should only be accessed by the lua thread itself.
    static ref RECEIVER: Mutex<Option<Receiver<LuaResponse>>> = Mutex::new(None);

    /// Whether the lua thread is currently running
    pub static ref RUNNING: RwLock<bool> = RwLock::new(false);
}

/// Represents an identifier for dealing with nested tables.
///
/// To access foo.bar.baz, use vec!["foo", "bar", "baz"].
///
/// To access foo[2], use vec!["foo", 2].
pub type LuaIdentifier = Vec<AnyLuaValue>;

/// Messages sent to the lua thread
#[derive(Debug)]
pub enum LuaQuery {
    /// Pings the lua thread
    Ping,
    /// Halt the lua thread
    Terminate,
    // Restart the lua thread
    Restart,

    /// Execute a string
    Execute(String),
    /// Execute a file
    ExecFile(String),

    /// Get a variable, expecting an AnyLuaValue
    GetValue(LuaIdentifier),
    /// Invoke a function found at the position,
    /// with the specified arguments.
    Invoke(LuaIdentifier, Vec<AnyLuaValue>),
    /// Set a value
    SetValue {
        /// The name of the thing to stuff
        name: LuaIdentifier,
        /// The value to store.
        val: Json
    },
    /// Create a new table
    NewTable(LuaIdentifier),
}

/// Messages received from lua thread
pub enum LuaResponse {
    /// Lua variable obtained
    Variable(Option<AnyLuaValue>),
    /// Lua error
    Error(hlua::LuaError),
    /// A function is returned
    Function(hlua::functions_read::LuaFunction<String>),
    /// Pong response from lua ping
    Pong,
}

unsafe impl Send for LuaQuery { }
unsafe impl Send for LuaResponse { }
unsafe impl Sync for LuaQuery { }
unsafe impl Sync for LuaResponse { }

/// Whether the lua thread is currently available
pub fn thread_running() -> bool {
    *RUNNING.read().unwrap()
}

/// Errors which may arise from attempting
/// to sending a message to the lua thread.
#[derive(Debug)]
pub enum LuaSendError {
    /// The thread crashed, was shut down, or rebooted.
    ThreadClosed,
    /// The thread has not been initialized yet (maybe not used)
    ThreadUninitialized,
    /// The sender had an issue, most likey because the thread panicked
    Sender
}

/// Attemps to send a LuaQuery to the lua thread.
pub fn try_send(query: LuaQuery) -> Result<(), LuaSendError> {
    if !thread_running() {
        Err(LuaSendError::ThreadClosed)
    }
    else if let Some(ref sender) = *SENDER.lock().unwrap() {
        match sender.send(query) {
            Ok(_) => Ok(()),
            Err(_) => Err(LuaSendError::Sender)
        }
    }
    else {
        Err(LuaSendError::ThreadUninitialized)
    }
}

/// Initialize the lua thread
pub fn init() {
    trace!("Initializing...");
    let (query_tx, query_rx) = channel::<LuaQuery>();
    let (answer_tx, answer_rx) = channel::<LuaResponse>();
    {
        let mut sender = SENDER.lock().unwrap();
        let mut receiver = RECEIVER.lock().unwrap();

        *sender = Some(query_tx);
        *receiver = Some(answer_rx);
    }

    thread::spawn(move || {
        thread_init(answer_tx, query_rx);
    });
    trace!("Created thread. Init finished.");
}

fn thread_init(sender: Sender<LuaResponse>, receiver: Receiver<LuaQuery>) {
    trace!("thread: initializing.");
    let mut lua = Lua::new();
    //unsafe {
    //    hlua_ffi::lua_atpanic(&mut lua.as_mut_lua().0, thread_on_panic);
    //}
    debug!("thread: Loading Lua libraries...");
    lua.openlibs();
    trace!("thread: Loading way-cooler lua extensions...");
    // We should have some good file handling, read files from /usr by default,
    // but for now we're reading directly from the source.
    lua.execute_from_reader::<(), File>(
        File::open("lib/lua/init.lua").unwrap()
    ).unwrap();
    trace!("thread: loading way-cooler libraries...");
    funcs::register_libraries(&mut lua);
    // Only ready after loading libs
    *RUNNING.write().unwrap() = true;
    debug!("thread: entering main loop...");
    thread_main_loop(sender, receiver, &mut lua);
}

fn thread_main_loop(sender: Sender<LuaResponse>, receiver: Receiver<LuaQuery>,
                    lua: &mut Lua) {
    loop {
        let request = receiver.recv();
        match request {
            Err(e) => {
                error!("Lua thread: unable to receive message: {}", e);
                error!("Lua thread: now panicking!");
                *RUNNING.write().unwrap() = false;

                panic!("Lua thread: lost contact with host, exiting!");
            }
            Ok(message) => {
                trace!("Handling a request");
                thread_handle_message(&sender, message, lua);
            }
        }
    }
}

fn thread_handle_message(sender: &Sender<LuaResponse>,
                         request: LuaQuery, lua: &mut Lua) {
    match request {
        LuaQuery::Terminate => {
            trace!("thread: Received terminate signal");
            *RUNNING.write().unwrap() = false;

            info!("thread: Lua thread terminating!");
            return;
        },

        LuaQuery::Restart => {
            trace!("thread: Received restart signal!");
            error!("thread: Lua thread restart not supported!");

            *RUNNING.write().unwrap() = false;

            panic!("Lua thread: Restart not supported!");
        },

        LuaQuery::Execute(code) => {
            trace!("thread: Received request to execute code");
            trace!("thread: Executing {:?}", code);

            match lua.execute::<()>(&code) {
                Err(error) => {
                    warn!("thread: Error executing code: {:?}", error);
                    let response = LuaResponse::Error(error);

                    thread_send(&sender, response);
                }
                Ok(_) => {
                    // This is gonna be really spammy one day
                    trace!("thread: Code executed okay.");
                }
            }
        },

        LuaQuery::ExecFile(name) => {
            trace!("thread: Received request to execute file {}", name);
            info!("thread: Executing {}", name);

            let path = Path::new(&name);
            let try_file = File::open(path);

            if let Ok(file) = try_file {
                let result = lua.execute_from_reader::<(), File>(file);
                if let Err(err) = result {
                    warn!("thread: Error executing {}!", name);

                    thread_send(&sender, LuaResponse::Error(err));
                }
                else {
                    trace!("thread: Execution of {} successful.", name);
                }
            }
            else { // Could not open file
                // Unwrap_err is used because we're in the else of let
                let read_error =
                    LuaError::ReadError(try_file.unwrap_err());

                thread_send(&sender, LuaResponse::Error(read_error));
            }
        },

        LuaQuery::GetValue(varname) => {
            trace!("thread: Received request to get variable {:?}", varname);
            let var_result = lua.get(format!("{:?}", varname));

            match var_result {
                Some(var) => {
                    thread_send(&sender, LuaResponse::Variable(Some(var)));
                }
                None => {
                    warn!("thread: Unable to get variable {:?}", varname);

                    thread_send(&sender, LuaResponse::Variable(None));
                }
            }
        },

        LuaQuery::SetValue { name: _name, val: _val } => {
            panic!("thread: unimplemented LuaQuery::SetValue!");
        },

        LuaQuery::NewTable(_name) => {
            panic!("thread: unimplemented LuaQuery::NewTable!");
        },

        LuaQuery::Ping => {
            panic!("thread: unimplemented LuaQuery::Ping!");
        },
        _ => {
            panic!("Unimplemented send type for lua thread!");
        }
    }
}

fn thread_send(sender: &Sender<LuaResponse>, response: LuaResponse) {
    trace!("Called thread_send");
    match sender.send(response) {
        Err(_) => {
            error!("thread: Unable to broadcast response!");
            error!("thread: Shutting down in response to inability \
                    to continue!");
            panic!("Lua thread unable to communicate with main thread, \
                    shutting down!");
        }
        Ok(_) => {}
    }
}
<|endoftext|>