<commit_before>use super::super::{Device, RawDevice, RawSurface, Surface};
use super::Error;

use drm::control::{connector, crtc, framebuffer, Device as ControlDevice, Mode};
use failure::ResultExt;
use gbm::{self, BufferObject, BufferObjectFlags, Format as GbmFormat, SurfaceBufferHandle};
use image::{ImageBuffer, Rgba};

use std::cell::{Cell, RefCell};
use std::rc::Rc;

use crate::backend::graphics::CursorBackend;

pub(super) struct GbmSurfaceInternal<D: RawDevice + 'static> {
    pub(super) dev: Rc<RefCell<gbm::Device<D>>>,
    pub(super) surface: RefCell<gbm::Surface<framebuffer::Handle>>,
    pub(super) crtc: <D as Device>::Surface,
    pub(super) cursor: Cell<(BufferObject<()>, (u32, u32))>,
    pub(super) current_frame_buffer: Cell<Option<framebuffer::Handle>>,
    pub(super) front_buffer: Cell<Option<SurfaceBufferHandle<framebuffer::Handle>>>,
    pub(super) next_buffer: Cell<Option<SurfaceBufferHandle<framebuffer::Handle>>>,
    pub(super) recreated: Cell<bool>,
    pub(super) logger: ::slog::Logger,
}

impl<D: RawDevice + 'static> GbmSurfaceInternal<D> {
    pub(super) fn unlock_buffer(&self) {
        // after the page swap is finished we need to release the rendered buffer.
        // this is called from the PageFlipHandler
        trace!(self.logger, "Releasing old front buffer");
        self.front_buffer.set(self.next_buffer.replace(None));
        // drop and release the old buffer
    }

    pub unsafe fn page_flip(&self) -> Result<(), Error<<<D as Device>::Surface as Surface>::Error>> {
        let res = {
            let nb = self.next_buffer.take();
            let res = nb.is_some();
            self.next_buffer.set(nb);
            res
        };
        if res {
            // We cannot call lock_front_buffer anymore without releasing the previous buffer, which will happen when the page flip is done
            warn!(self.logger, "Tried to swap with an already queued flip");
            return Err(Error::FrontBuffersExhausted);
        }

        // supporting only one buffer would cause a lot of inconvinience and
        // would most likely result in a lot of flickering.
        // neither weston, wlc or wlroots bother with that as well.
        // so we just assume we got at least two buffers to do flipping.
        let mut next_bo = self
            .surface
            .borrow()
            .lock_front_buffer()
            .map_err(|_| Error::FrontBufferLockFailed)?;

        // create a framebuffer if the front buffer does not have one already
        // (they are reused by gbm)
        let maybe_fb = next_bo
            .userdata()
            .map_err(|_| Error::InvalidInternalState)?
            .cloned();
        let fb = if let Some(info) = maybe_fb {
            info
        } else {
            let fb = self
                .crtc
                .add_planar_framebuffer(&*next_bo, &[0; 4], 0)
                .compat()
                .map_err(Error::FramebufferCreationFailed)?;
            next_bo.set_userdata(fb).unwrap();
            fb
        };
        self.next_buffer.set(Some(next_bo));

        if self.recreated.get() {
            debug!(self.logger, "Commiting new state");
            self.crtc.commit(fb).map_err(Error::Underlying)?;
            self.recreated.set(false);
        } else {
            trace!(self.logger, "Queueing Page flip");
            RawSurface::page_flip(&self.crtc, fb).map_err(Error::Underlying)?;
        }

        self.current_frame_buffer.set(Some(fb));

        Ok(())
    }

    pub fn recreate(&self) -> Result<(), Error<<<D as Device>::Surface as Surface>::Error>> {
        let (w, h) = self.pending_mode().size();

        // Recreate the surface and the related resources to match the new
        // resolution.
        debug!(self.logger, "(Re-)Initializing surface (with mode: {}:{})", w, h);
        let surface = self
            .dev
            .borrow_mut()
            .create_surface(
                w as u32,
                h as u32,
                GbmFormat::XRGB8888,
                BufferObjectFlags::SCANOUT | BufferObjectFlags::RENDERING,
            )
            .map_err(Error::SurfaceCreationFailed)?;

        // Clean up buffers
        if let Some(Ok(Some(fb))) = self.next_buffer.take().map(|mut bo| bo.take_userdata()) {
            if let Err(err) = self.crtc.destroy_framebuffer(fb) {
                warn!(
                    self.logger,
                    "Error releasing old back_buffer framebuffer: {:?}", err
                );
            }
        }

        if let Some(Ok(Some(fb))) = self.front_buffer.take().map(|mut bo| bo.take_userdata()) {
            if let Err(err) = self.crtc.destroy_framebuffer(fb) {
                warn!(
                    self.logger,
                    "Error releasing old front_buffer framebuffer: {:?}", err
                );
            }
        }

        // Drop the old surface after cleanup
        *self.surface.borrow_mut() = surface;

        self.recreated.set(true);

        Ok(())
    }
}

impl<D: RawDevice + 'static> Surface for GbmSurfaceInternal<D> {
    type Connectors = <<D as Device>::Surface as Surface>::Connectors;
    type Error = Error<<<D as Device>::Surface as Surface>::Error>;

    fn crtc(&self) -> crtc::Handle {
        self.crtc.crtc()
    }

    fn current_connectors(&self) -> Self::Connectors {
        self.crtc.current_connectors()
    }

    fn pending_connectors(&self) -> Self::Connectors {
        self.crtc.pending_connectors()
    }

    fn add_connector(&self, connector: connector::Handle) -> Result<(), Self::Error> {
        self.crtc.add_connector(connector).map_err(Error::Underlying)
    }

    fn remove_connector(&self, connector: connector::Handle) -> Result<(), Self::Error> {
        self.crtc.remove_connector(connector).map_err(Error::Underlying)
    }

    fn set_connectors(&self, connectors: &[connector::Handle]) -> Result<(), Self::Error> {
        self.crtc.set_connectors(connectors).map_err(Error::Underlying)
    }

    fn current_mode(&self) -> Mode {
        self.crtc.current_mode()
    }

    fn pending_mode(&self) -> Mode {
        self.crtc.pending_mode()
    }

    fn use_mode(&self, mode: Mode) -> Result<(), Self::Error> {
        self.crtc.use_mode(mode).map_err(Error::Underlying)
    }
}

#[cfg(feature = "backend_drm")]
impl<D: RawDevice + 'static> CursorBackend for GbmSurfaceInternal<D>
where
    <D as RawDevice>::Surface: CursorBackend<CursorFormat = dyn drm::buffer::Buffer>,
    <<D as RawDevice>::Surface as CursorBackend>::Error: ::std::error::Error + Send,
{
    type CursorFormat = ImageBuffer<Rgba<u8>, Vec<u8>>;
    type Error = Error<<<D as Device>::Surface as CursorBackend>::Error>;

    fn set_cursor_position(&self, x: u32, y: u32) -> Result<(), Self::Error> {
        self.crtc.set_cursor_position(x, y).map_err(Error::Underlying)
    }

    fn set_cursor_representation(
        &self,
        buffer: &ImageBuffer<Rgba<u8>, Vec<u8>>,
        hotspot: (u32, u32),
    ) -> Result<(), Self::Error> {
        let (w, h) = buffer.dimensions();
        debug!(self.logger, "Importing cursor");

        // import the cursor into a buffer we can render
        let mut cursor = self
            .dev
            .borrow_mut()
            .create_buffer_object(
                w,
                h,
                GbmFormat::ARGB8888,
                BufferObjectFlags::CURSOR | BufferObjectFlags::WRITE,
            )
            .map_err(Error::BufferCreationFailed)?;

        cursor
            .write(&**buffer)
            .map_err(|_| Error::DeviceDestroyed)?
            .map_err(Error::BufferWriteFailed)?;

        trace!(self.logger, "Setting the new imported cursor");

        self.crtc
            .set_cursor_representation(&cursor, hotspot)
            .map_err(Error::Underlying)?;

        // and store it
        self.cursor.set((cursor, hotspot));
        Ok(())
    }
}

impl<D: RawDevice + 'static> Drop for GbmSurfaceInternal<D> {
    fn drop(&mut self) {
        // Drop framebuffers attached to the userdata of the gbm surface buffers.
        // (They don't implement drop, as they need the device)
        if let Ok(Some(fb)) = {
            if let Some(mut next) = self.next_buffer.take() {
                next.take_userdata()
            } else {
                Ok(None)
            }
        } {
            // ignore failure at this point
            let _ = self.crtc.destroy_framebuffer(fb);
        }

        if let Ok(Some(fb)) = {
            if let Some(mut next) = self.front_buffer.take() {
                next.take_userdata()
            } else {
                Ok(None)
            }
        } {
            // ignore failure at this point
            let _ = self.crtc.destroy_framebuffer(fb);
        }
    }
}

/// Gbm surface for rendering
pub struct GbmSurface<D: RawDevice + 'static>(pub(super) Rc<GbmSurfaceInternal<D>>);

impl<D: RawDevice + 'static> GbmSurface<D> {
    /// Flips the underlying buffers.
    ///
    /// The surface will report being already flipped until the matching event
    /// was processed either by calling [`Device::process_events`] manually after the flip
    /// (bad idea performance-wise) or by binding the device to an event-loop by using
    /// [`device_bind`](::backend::drm::device_bind).
    ///
    /// *Note*: This might trigger a full modeset on the underlying device,
    /// potentially causing some flickering. In that case this operation is
    /// blocking until the crtc is in the desired state.
    ///
    /// # Safety
    ///
    /// When used in conjunction with an EGL context, this must be called exactly once
    /// after page-flipping the associated context.
    pub unsafe fn page_flip(&self) -> ::std::result::Result<(), <Self as Surface>::Error> {
        self.0.page_flip()
    }

    /// Recreate underlying gbm resources.
    ///
    /// This recreates the gbm surfaces resources, which might be needed after e.g.
    /// calling [`Surface::use_mode`](Surface::use_mode).
    /// You may check if your [`GbmSurface`] needs recreation through
    /// [`needs_recreation`](GbmSurface::needs_recreation).
    pub fn recreate(&self) -> Result<(), <Self as Surface>::Error> {
        self.0.recreate()
    }

    /// Check if underlying gbm resources need to be recreated.
    pub fn needs_recreation(&self) -> bool {
        self.0.crtc.commit_pending()
    }
}

impl<D: RawDevice + 'static> Surface for GbmSurface<D> {
    type Connectors = <<D as Device>::Surface as Surface>::Connectors;
    type Error = Error<<<D as Device>::Surface as Surface>::Error>;

    fn crtc(&self) -> crtc::Handle {
        self.0.crtc()
    }

    fn current_connectors(&self) -> Self::Connectors {
        self.0.current_connectors()
    }

    fn pending_connectors(&self) -> Self::Connectors {
        self.0.pending_connectors()
    }

    fn add_connector(&self, connector: connector::Handle) -> Result<(), Self::Error> {
        self.0.add_connector(connector)
    }

    fn remove_connector(&self, connector: connector::Handle) -> Result<(), Self::Error> {
        self.0.remove_connector(connector)
    }

    fn set_connectors(&self, connectors: &[connector::Handle]) -> Result<(), Self::Error> {
        self.0.set_connectors(connectors)
    }

    fn current_mode(&self) -> Mode {
        self.0.current_mode()
    }

    fn pending_mode(&self) -> Mode {
        self.0.pending_mode()
    }

    fn use_mode(&self, mode: Mode) -> Result<(), Self::Error> {
        self.0.use_mode(mode)
    }
}

#[cfg(feature = "backend_drm")]
impl<D: RawDevice + 'static> CursorBackend for GbmSurface<D>
where
    <D as RawDevice>::Surface: CursorBackend<CursorFormat = dyn drm::buffer::Buffer>,
    <<D as RawDevice>::Surface as CursorBackend>::Error: ::std::error::Error + Send,
{
    type CursorFormat = ImageBuffer<Rgba<u8>, Vec<u8>>;
    type Error = Error<<<D as Device>::Surface as CursorBackend>::Error>;

    fn set_cursor_position(&self, x: u32, y: u32) -> Result<(), Self::Error> {
        self.0.set_cursor_position(x, y)
    }

    fn set_cursor_representation(
        &self,
        buffer: &ImageBuffer<Rgba<u8>, Vec<u8>>,
        hotspot: (u32, u32),
    ) -> Result<(), Self::Error> {
        self.0.set_cursor_representation(buffer, hotspot)
    }
}
<commit_msg>gbm: split out clear_framebuffers<commit_after>use super::super::{Device, RawDevice, RawSurface, Surface};
use super::Error;

use drm::control::{connector, crtc, framebuffer, Device as ControlDevice, Mode};
use failure::ResultExt;
use gbm::{self, BufferObject, BufferObjectFlags, Format as GbmFormat, SurfaceBufferHandle};
use image::{ImageBuffer, Rgba};

use std::cell::{Cell, RefCell};
use std::rc::Rc;

use crate::backend::graphics::CursorBackend;

pub(super) struct GbmSurfaceInternal<D: RawDevice + 'static> {
    pub(super) dev: Rc<RefCell<gbm::Device<D>>>,
    pub(super) surface: RefCell<gbm::Surface<framebuffer::Handle>>,
    pub(super) crtc: <D as Device>::Surface,
    pub(super) cursor: Cell<(BufferObject<()>, (u32, u32))>,
    pub(super) current_frame_buffer: Cell<Option<framebuffer::Handle>>,
    pub(super) front_buffer: Cell<Option<SurfaceBufferHandle<framebuffer::Handle>>>,
    pub(super) next_buffer: Cell<Option<SurfaceBufferHandle<framebuffer::Handle>>>,
    pub(super) recreated: Cell<bool>,
    pub(super) logger: ::slog::Logger,
}

impl<D: RawDevice + 'static> GbmSurfaceInternal<D> {
    pub(super) fn unlock_buffer(&self) {
        // after the page swap is finished we need to release the rendered buffer.
        // this is called from the PageFlipHandler
        trace!(self.logger, "Releasing old front buffer");
        self.front_buffer.set(self.next_buffer.replace(None));
        // drop and release the old buffer
    }

    pub unsafe fn page_flip(&self) -> Result<(), Error<<<D as Device>::Surface as Surface>::Error>> {
        let res = {
            let nb = self.next_buffer.take();
            let res = nb.is_some();
            self.next_buffer.set(nb);
            res
        };
        if res {
            // We cannot call lock_front_buffer anymore without releasing the previous buffer, which will happen when the page flip is done
            warn!(self.logger, "Tried to swap with an already queued flip");
            return Err(Error::FrontBuffersExhausted);
        }

        // supporting only one buffer would cause a lot of inconvinience and
        // would most likely result in a lot of flickering.
        // neither weston, wlc or wlroots bother with that as well.
        // so we just assume we got at least two buffers to do flipping.
        let mut next_bo = self
            .surface
            .borrow()
            .lock_front_buffer()
            .map_err(|_| Error::FrontBufferLockFailed)?;

        // create a framebuffer if the front buffer does not have one already
        // (they are reused by gbm)
        let maybe_fb = next_bo
            .userdata()
            .map_err(|_| Error::InvalidInternalState)?
            .cloned();
        let fb = if let Some(info) = maybe_fb {
            info
        } else {
            let fb = self
                .crtc
                .add_planar_framebuffer(&*next_bo, &[0; 4], 0)
                .compat()
                .map_err(Error::FramebufferCreationFailed)?;
            next_bo.set_userdata(fb).unwrap();
            fb
        };
        self.next_buffer.set(Some(next_bo));

        if self.recreated.get() {
            debug!(self.logger, "Commiting new state");
            self.crtc.commit(fb).map_err(Error::Underlying)?;
            self.recreated.set(false);
        } else {
            trace!(self.logger, "Queueing Page flip");
            RawSurface::page_flip(&self.crtc, fb).map_err(Error::Underlying)?;
        }

        self.current_frame_buffer.set(Some(fb));

        Ok(())
    }

    pub fn recreate(&self) -> Result<(), Error<<<D as Device>::Surface as Surface>::Error>> {
        let (w, h) = self.pending_mode().size();

        // Recreate the surface and the related resources to match the new
        // resolution.
        debug!(self.logger, "(Re-)Initializing surface (with mode: {}:{})", w, h);
        let surface = self
            .dev
            .borrow_mut()
            .create_surface(
                w as u32,
                h as u32,
                GbmFormat::XRGB8888,
                BufferObjectFlags::SCANOUT | BufferObjectFlags::RENDERING,
            )
            .map_err(Error::SurfaceCreationFailed)?;

        // Clean up buffers
        self.clear_framebuffers();

        // Drop the old surface after cleanup
        *self.surface.borrow_mut() = surface;

        self.recreated.set(true);

        Ok(())
    }

    pub fn clear_framebuffers(&self) {
        if let Some(Ok(Some(fb))) = self.next_buffer.take().map(|mut bo| bo.take_userdata()) {
            if let Err(err) = self.crtc.destroy_framebuffer(fb) {
                warn!(
                    self.logger,
                    "Error releasing old back_buffer framebuffer: {:?}", err
                );
            }
        }

        if let Some(Ok(Some(fb))) = self.front_buffer.take().map(|mut bo| bo.take_userdata()) {
            if let Err(err) = self.crtc.destroy_framebuffer(fb) {
                warn!(
                    self.logger,
                    "Error releasing old front_buffer framebuffer: {:?}", err
                );
            }
        }
    }
}

impl<D: RawDevice + 'static> Surface for GbmSurfaceInternal<D> {
    type Connectors = <<D as Device>::Surface as Surface>::Connectors;
    type Error = Error<<<D as Device>::Surface as Surface>::Error>;

    fn crtc(&self) -> crtc::Handle {
        self.crtc.crtc()
    }

    fn current_connectors(&self) -> Self::Connectors {
        self.crtc.current_connectors()
    }

    fn pending_connectors(&self) -> Self::Connectors {
        self.crtc.pending_connectors()
    }

    fn add_connector(&self, connector: connector::Handle) -> Result<(), Self::Error> {
        self.crtc.add_connector(connector).map_err(Error::Underlying)
    }

    fn remove_connector(&self, connector: connector::Handle) -> Result<(), Self::Error> {
        self.crtc.remove_connector(connector).map_err(Error::Underlying)
    }

    fn set_connectors(&self, connectors: &[connector::Handle]) -> Result<(), Self::Error> {
        self.crtc.set_connectors(connectors).map_err(Error::Underlying)
    }

    fn current_mode(&self) -> Mode {
        self.crtc.current_mode()
    }

    fn pending_mode(&self) -> Mode {
        self.crtc.pending_mode()
    }

    fn use_mode(&self, mode: Mode) -> Result<(), Self::Error> {
        self.crtc.use_mode(mode).map_err(Error::Underlying)
    }
}

#[cfg(feature = "backend_drm")]
impl<D: RawDevice + 'static> CursorBackend for GbmSurfaceInternal<D>
where
    <D as RawDevice>::Surface: CursorBackend<CursorFormat = dyn drm::buffer::Buffer>,
    <<D as RawDevice>::Surface as CursorBackend>::Error: ::std::error::Error + Send,
{
    type CursorFormat = ImageBuffer<Rgba<u8>, Vec<u8>>;
    type Error = Error<<<D as Device>::Surface as CursorBackend>::Error>;

    fn set_cursor_position(&self, x: u32, y: u32) -> Result<(), Self::Error> {
        self.crtc.set_cursor_position(x, y).map_err(Error::Underlying)
    }

    fn set_cursor_representation(
        &self,
        buffer: &ImageBuffer<Rgba<u8>, Vec<u8>>,
        hotspot: (u32, u32),
    ) -> Result<(), Self::Error> {
        let (w, h) = buffer.dimensions();
        debug!(self.logger, "Importing cursor");

        // import the cursor into a buffer we can render
        let mut cursor = self
            .dev
            .borrow_mut()
            .create_buffer_object(
                w,
                h,
                GbmFormat::ARGB8888,
                BufferObjectFlags::CURSOR | BufferObjectFlags::WRITE,
            )
            .map_err(Error::BufferCreationFailed)?;

        cursor
            .write(&**buffer)
            .map_err(|_| Error::DeviceDestroyed)?
            .map_err(Error::BufferWriteFailed)?;

        trace!(self.logger, "Setting the new imported cursor");

        self.crtc
            .set_cursor_representation(&cursor, hotspot)
            .map_err(Error::Underlying)?;

        // and store it
        self.cursor.set((cursor, hotspot));
        Ok(())
    }
}

impl<D: RawDevice + 'static> Drop for GbmSurfaceInternal<D> {
    fn drop(&mut self) {
        // Drop framebuffers attached to the userdata of the gbm surface buffers.
        // (They don't implement drop, as they need the device)
        self.clear_framebuffers();
    }
}

/// Gbm surface for rendering
pub struct GbmSurface<D: RawDevice + 'static>(pub(super) Rc<GbmSurfaceInternal<D>>);

impl<D: RawDevice + 'static> GbmSurface<D> {
    /// Flips the underlying buffers.
    ///
    /// The surface will report being already flipped until the matching event
    /// was processed either by calling [`Device::process_events`] manually after the flip
    /// (bad idea performance-wise) or by binding the device to an event-loop by using
    /// [`device_bind`](::backend::drm::device_bind).
    ///
    /// *Note*: This might trigger a full modeset on the underlying device,
    /// potentially causing some flickering. In that case this operation is
    /// blocking until the crtc is in the desired state.
    ///
    /// # Safety
    ///
    /// When used in conjunction with an EGL context, this must be called exactly once
    /// after page-flipping the associated context.
    pub unsafe fn page_flip(&self) -> ::std::result::Result<(), <Self as Surface>::Error> {
        self.0.page_flip()
    }

    /// Recreate underlying gbm resources.
    ///
    /// This recreates the gbm surfaces resources, which might be needed after e.g.
    /// calling [`Surface::use_mode`](Surface::use_mode).
    /// You may check if your [`GbmSurface`] needs recreation through
    /// [`needs_recreation`](GbmSurface::needs_recreation).
    pub fn recreate(&self) -> Result<(), <Self as Surface>::Error> {
        self.0.recreate()
    }

    /// Check if underlying gbm resources need to be recreated.
    pub fn needs_recreation(&self) -> bool {
        self.0.crtc.commit_pending()
    }
}

impl<D: RawDevice + 'static> Surface for GbmSurface<D> {
    type Connectors = <<D as Device>::Surface as Surface>::Connectors;
    type Error = Error<<<D as Device>::Surface as Surface>::Error>;

    fn crtc(&self) -> crtc::Handle {
        self.0.crtc()
    }

    fn current_connectors(&self) -> Self::Connectors {
        self.0.current_connectors()
    }

    fn pending_connectors(&self) -> Self::Connectors {
        self.0.pending_connectors()
    }

    fn add_connector(&self, connector: connector::Handle) -> Result<(), Self::Error> {
        self.0.add_connector(connector)
    }

    fn remove_connector(&self, connector: connector::Handle) -> Result<(), Self::Error> {
        self.0.remove_connector(connector)
    }

    fn set_connectors(&self, connectors: &[connector::Handle]) -> Result<(), Self::Error> {
        self.0.set_connectors(connectors)
    }

    fn current_mode(&self) -> Mode {
        self.0.current_mode()
    }

    fn pending_mode(&self) -> Mode {
        self.0.pending_mode()
    }

    fn use_mode(&self, mode: Mode) -> Result<(), Self::Error> {
        self.0.use_mode(mode)
    }
}

#[cfg(feature = "backend_drm")]
impl<D: RawDevice + 'static> CursorBackend for GbmSurface<D>
where
    <D as RawDevice>::Surface: CursorBackend<CursorFormat = dyn drm::buffer::Buffer>,
    <<D as RawDevice>::Surface as CursorBackend>::Error: ::std::error::Error + Send,
{
    type CursorFormat = ImageBuffer<Rgba<u8>, Vec<u8>>;
    type Error = Error<<<D as Device>::Surface as CursorBackend>::Error>;

    fn set_cursor_position(&self, x: u32, y: u32) -> Result<(), Self::Error> {
        self.0.set_cursor_position(x, y)
    }

    fn set_cursor_representation(
        &self,
        buffer: &ImageBuffer<Rgba<u8>, Vec<u8>>,
        hotspot: (u32, u32),
    ) -> Result<(), Self::Error> {
        self.0.set_cursor_representation(buffer, hotspot)
    }
}
<|endoftext|>


<commit_before>use super::file_watcher::{create_file_watcher,
                          Callbacks};
use crate::manager::service::Service;
use habitat_common::{liveliness_checker,
                     outputln,
                     templating::config::UserConfigPath};
use habitat_core::{fs::USER_CONFIG_FILE,
                   service::ServiceGroup};
use std::{collections::HashMap,
          io,
          path::{Path,
                 PathBuf},
          sync::{mpsc::{channel,
                        sync_channel,
                        Receiver,
                        Sender,
                        SyncSender,
                        TryRecvError,
                        TrySendError},
                 Arc,
                 Mutex},
          thread::{self,
                   Builder as ThreadBuilder},
          time::Duration};

static LOGKEY: &str = "UCW";

// This trait exists to ease the testing of functions that receive a Service. Creating Services
// requires a lot of ceremony, so we work around this with this trait.
pub trait Serviceable {
    fn name(&self) -> &str;
    fn user_config_path(&self) -> &UserConfigPath;
    fn service_group(&self) -> &ServiceGroup;
}

impl Serviceable for Service {
    fn name(&self) -> &str { &self.pkg.name }

    fn user_config_path(&self) -> &UserConfigPath { &self.cfg.user_config_path }

    fn service_group(&self) -> &ServiceGroup { &self.service_group }
}

// WorkerState contains the channels the worker uses to communicate
// with the Watcher.
struct WorkerState {
    // This receiver is used by the watcher to be notified when a
    // worker has events.  The channel is a SyncChannel with buffer
    // size 1, as we are only interested in the fact that there were
    // events, not how many there were.
    have_events:      Receiver<()>,
    // This sender is used by the watcher to notify a worker to stop
    // running.  It is an async channel because we never want the
    // UserConfigWatcher to block, even if the receiver end of the
    // channel somehow dies and/or fails to consume the message.
    stop_running:     Sender<()>,
    // This receiver is used by the watcher tests to be notified when
    // a worker finished setting up the watcher and is about to
    // starting looping it.
    //
    // Silence the dead code warnings from rustc, because it is only
    // used in tests for synchronization purposes.
    #[allow(dead_code)]
    started_watching: Receiver<()>,
}

type ServiceName = String;
#[derive(Clone)]
pub struct UserConfigWatcher {
    // We use Arc/Mutex here, because this needs to be shareable
    // across threads so we can remove watchers from futures; the
    // Sender and Receiver members of WorkerState aren't themselves
    // shareable.
    states: Arc<Mutex<HashMap<ServiceName, WorkerState>>>,
}

impl UserConfigWatcher {
    pub fn new() -> Self { Self { states: Arc::new(Mutex::new(HashMap::new())), } }

    /// Adds a service to the User Config Watcher, thereby starting a watcher thread.
    pub fn add<T: Serviceable>(&mut self, service: &T) -> io::Result<()> {
        // It isn't possible to use the `or_insert_with` function here because it can't have a
        // return value, which we need to return the error from `Worker::run`.
        let mut states = self.states.lock().expect("states lock was poisoned");
        if states.get(service.name()).is_none() {
            let user_toml_path = match service.user_config_path() {
                UserConfigPath::Recommended(ref p) => p.join(USER_CONFIG_FILE),
                UserConfigPath::Deprecated(ref p) => {
                    outputln!(
                        preamble service.service_group(),
                        "Not watching {}, because it is located in deprecated path ({}).",
                        USER_CONFIG_FILE,
                        p.display(),
                    );
                    return Ok(());
                }
            };
            // Establish bi-directional communication with the worker by creating two channels.
            // The sync_channel's buffer size is 1 because we want to use it as a boolean, i.e. we
            // are not interested in the events themselves, but only whether at least one has
            // happened.
            let (events_tx, events_rx) = sync_channel(1);
            let (running_tx, running_rx) = channel();
            let (watching_tx, watching_rx) = sync_channel(1);

            Worker::run(user_toml_path, events_tx, running_rx, watching_tx)?;

            outputln!(preamble service.service_group(), "Watching {}", USER_CONFIG_FILE);

            let state = WorkerState { have_events:      events_rx,
                                      stop_running:     running_tx,
                                      started_watching: watching_rx, };

            states.insert(service.name().to_owned(), state);
        }

        Ok(())
    }

    /// Removes a service from the User Config Watcher, and sends a message to the watcher thread
    /// to stop running.
    pub fn remove<T: Serviceable>(&mut self, service: &T) {
        if let Some(state) = self.states
                                 .lock()
                                 .expect("states lock was poisoned")
                                 .remove(service.name())
        {
            if let Err(e) = state.stop_running.send(()) {
                debug!("Error stopping user-config watcher thread for service {}: {:?}",
                       service.name(),
                       e);
            }
        }
    }

    /// Checks whether the watcher for the specified service has observed any events.
    ///
    /// This also consumes the events.
    pub fn have_events_for<T: Serviceable>(&self, service: &T) -> bool {
        if let Some(state) = self.states
                                 .lock()
                                 .expect("states lock was poisoned")
                                 .get(service.name())
        {
            let rx = &state.have_events;

            match rx.try_recv() {
                Ok(_) => {
                    return true;
                }
                Err(TryRecvError::Empty) => return false,
                Err(TryRecvError::Disconnected) => {
                    debug!("UserConfigWatcher worker has died!");
                    return false;
                }
            }
        }

        false
    }
}

struct UserConfigCallbacks {
    have_events: SyncSender<()>,
}

impl Callbacks for UserConfigCallbacks {
    fn file_appeared(&mut self, _: &Path) { self.perform(); }

    fn file_modified(&mut self, _: &Path) { self.perform(); }

    fn file_disappeared(&mut self, _: &Path) { self.perform(); }
}

impl UserConfigCallbacks {
    fn perform(&self) {
        if let Err(TrySendError::Disconnected(_)) = self.have_events.try_send(()) {
            debug!("Worker could not notify Manager of event");
        }
    }
}

struct Worker;

impl Worker {
    // starts a new thread with the file watcher tracking the service's user-config file
    pub fn run(path: PathBuf,
               have_events: SyncSender<()>,
               stop_running: Receiver<()>,
               started_watching: SyncSender<()>)
               -> io::Result<()> {
        ThreadBuilder::new().name(format!("user-config-watcher-{}", path.display()))
                            .spawn(move || -> liveliness_checker::ThreadUnregistered<(), String> {
                                let checked_thread = liveliness_checker::mark_thread_alive();

                                debug!("UserConfigWatcher({}) worker thread starting",
                                       path.display(),);
                                let callbacks = UserConfigCallbacks { have_events };
                                let mut file_watcher =
                                    match create_file_watcher(&path,
                                                               callbacks,
                                                               false)
                                    {
                                        Ok(w) => w,
                                        Err(e) => {
                                            let msg = format!("UserConfigWatcher({}) could not \
                                                               start notifier, ending thread ({})",
                                                              path.display(),
                                                              e,);
                                            outputln!("{}", msg);
                                            return checked_thread.unregister(Err(msg));
                                        }
                                    };

                                let _ = started_watching.try_send(());

                                loop {
                                    let checked_thread = liveliness_checker::mark_thread_alive();

                                    match stop_running.try_recv() {
                                        // As long as the `stop_running` channel is
                                        // empty, this branch will execute on every
                                        // iteration.
                                        Err(TryRecvError::Empty) => {
                                            if let Err(e) = file_watcher.single_iteration() {
                                                let msg = format!("UserConfigWatcher({}) could \
                                                                   not run notifier, ending \
                                                                   thread ({})",
                                                                  path.display(),
                                                                  e,);
                                                outputln!("{}", msg);
                                                break checked_thread.unregister(Err(msg));
                                            };
                                        }

                                        // If we receive a message on the channel, we stop.
                                        Ok(_) => break checked_thread.unregister(Ok(())),

                                        // If the channel is disconnected, we stop as well.
                                        Err(TryRecvError::Disconnected) => {
                                            let msg = format!("UserConfigWatcher({}) worker \
                                                               thread failed to receive on \
                                                               channel",
                                                              path.display(),);
                                            debug!("{}", msg);
                                            break checked_thread.unregister(Err(msg));
                                        }
                                    }

                                    thread::sleep(Duration::from_secs(1));
                                }
                            })?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::{env,
              fs::{remove_file,
                   File},
              io::Write,
              str::FromStr,
              thread,
              time::{Duration,
                     Instant}};

    use crate::manager::file_watcher::WATCHER_DELAY_MS;

    use tempfile::TempDir;

    #[test]
    fn test_events_with_poll_watcher() {
        env::set_var("HAB_STUDIO_HOST_ARCH", "aarch64-darwin");
        no_events_at_first();
        events_present_after_adding_config();
        events_present_after_changing_config();
        events_present_after_removing_config();
        env::remove_var("HAB_STUDIO_HOST_ARCH");
    }

    #[test]
    fn no_events_at_first() {
        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");

        assert!(!ucm.have_events_for(&service));
    }

    #[test]
    fn events_present_after_adding_config() {
        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));

        File::create(service.user_config_path().get_path().join(USER_CONFIG_FILE))
            .expect("creating file");

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_changing_config() {
        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        let mut file = File::create(&file_path).expect("creating file");

        file.write_all(b"42").expect(USER_CONFIG_FILE);

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_removing_config() {
        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        File::create(&file_path).expect("creating file");

        // Allow the watcher to notice that a file was created.
        wait_for_events(&ucm, &service);

        remove_file(&file_path).expect("removing file");

        assert!(wait_for_events(&ucm, &service));
    }

    fn wait_for_watcher<T: Serviceable>(ucm: &UserConfigWatcher, service: &T) -> bool {
        let start = Instant::now();
        let timeout = Duration::from_secs(10);

        while start.elapsed() < timeout {
            let states = ucm.states.lock().expect("states lock was poisoned");
            let state = states.get(service.name()).expect("service added");
            match state.started_watching.try_recv() {
                Ok(_) => {
                    println!("Received data on the start_watching channel. Returning true.");
                    return true;
                }
                Err(TryRecvError::Empty) => {
                    println!("Received nothing on the start_watching channel. Returning ().");
                }
                Err(TryRecvError::Disconnected) => {
                    println!("The start_watching channel was disconnected. Returning false.");
                    return false;
                }
            }

            thread::sleep(Duration::from_millis(100));
        }

        false
    }

    fn wait_for_events<T: Serviceable>(ucm: &UserConfigWatcher, service: &T) -> bool {
        let start = Instant::now();
        let timeout = Duration::from_millis(WATCHER_DELAY_MS * 5);

        while start.elapsed() < timeout {
            if ucm.have_events_for(service) {
                return true;
            }

            thread::sleep(Duration::from_millis(100));
        }

        false
    }

    struct TestService {
        // It is only used as a way to remove the temporary directory
        // when the test finishes. Not used otherwise.
        #[allow(dead_code)]
        tmp:              TempDir,
        name:             String,
        user_config_path: UserConfigPath,
        service_group:    ServiceGroup,
    }

    impl Serviceable for TestService {
        fn name(&self) -> &str { &self.name }

        fn user_config_path(&self) -> &UserConfigPath { &self.user_config_path }

        fn service_group(&self) -> &ServiceGroup { &self.service_group }
    }

    impl Default for TestService {
        fn default() -> Self {
            let tmp = TempDir::new().expect("creating temp dir");
            let path = UserConfigPath::Recommended(tmp.path().to_path_buf());
            Self { tmp,
                   name: String::from("foo"),
                   user_config_path: path,
                   service_group: ServiceGroup::from_str("foo.bar@yoyodine").unwrap() }
        }
    }
}
<commit_msg>update user_config_watcher<commit_after>use super::file_watcher::{create_file_watcher,
                          Callbacks};
use crate::manager::service::Service;
use habitat_common::{liveliness_checker,
                     outputln,
                     templating::config::UserConfigPath};
use habitat_core::{fs::USER_CONFIG_FILE,
                   service::ServiceGroup};
use std::{collections::HashMap,
          io,
          path::{Path,
                 PathBuf},
          sync::{mpsc::{channel,
                        sync_channel,
                        Receiver,
                        Sender,
                        SyncSender,
                        TryRecvError,
                        TrySendError},
                 Arc,
                 Mutex},
          thread::{self,
                   Builder as ThreadBuilder},
          time::Duration};

static LOGKEY: &str = "UCW";

// This trait exists to ease the testing of functions that receive a Service. Creating Services
// requires a lot of ceremony, so we work around this with this trait.
pub trait Serviceable {
    fn name(&self) -> &str;
    fn user_config_path(&self) -> &UserConfigPath;
    fn service_group(&self) -> &ServiceGroup;
}

impl Serviceable for Service {
    fn name(&self) -> &str { &self.pkg.name }

    fn user_config_path(&self) -> &UserConfigPath { &self.cfg.user_config_path }

    fn service_group(&self) -> &ServiceGroup { &self.service_group }
}

// WorkerState contains the channels the worker uses to communicate
// with the Watcher.
struct WorkerState {
    // This receiver is used by the watcher to be notified when a
    // worker has events.  The channel is a SyncChannel with buffer
    // size 1, as we are only interested in the fact that there were
    // events, not how many there were.
    have_events:      Receiver<()>,
    // This sender is used by the watcher to notify a worker to stop
    // running.  It is an async channel because we never want the
    // UserConfigWatcher to block, even if the receiver end of the
    // channel somehow dies and/or fails to consume the message.
    stop_running:     Sender<()>,
    // This receiver is used by the watcher tests to be notified when
    // a worker finished setting up the watcher and is about to
    // starting looping it.
    //
    // Silence the dead code warnings from rustc, because it is only
    // used in tests for synchronization purposes.
    #[allow(dead_code)]
    started_watching: Receiver<()>,
}

type ServiceName = String;
#[derive(Clone)]
pub struct UserConfigWatcher {
    // We use Arc/Mutex here, because this needs to be shareable
    // across threads so we can remove watchers from futures; the
    // Sender and Receiver members of WorkerState aren't themselves
    // shareable.
    states: Arc<Mutex<HashMap<ServiceName, WorkerState>>>,
}

impl UserConfigWatcher {
    pub fn new() -> Self { Self { states: Arc::new(Mutex::new(HashMap::new())), } }

    /// Adds a service to the User Config Watcher, thereby starting a watcher thread.
    pub fn add<T: Serviceable>(&mut self, service: &T) -> io::Result<()> {
        // It isn't possible to use the `or_insert_with` function here because it can't have a
        // return value, which we need to return the error from `Worker::run`.
        let mut states = self.states.lock().expect("states lock was poisoned");
        if states.get(service.name()).is_none() {
            let user_toml_path = match service.user_config_path() {
                UserConfigPath::Recommended(ref p) => p.join(USER_CONFIG_FILE),
                UserConfigPath::Deprecated(ref p) => {
                    outputln!(
                        preamble service.service_group(),
                        "Not watching {}, because it is located in deprecated path ({}).",
                        USER_CONFIG_FILE,
                        p.display(),
                    );
                    return Ok(());
                }
            };
            // Establish bi-directional communication with the worker by creating two channels.
            // The sync_channel's buffer size is 1 because we want to use it as a boolean, i.e. we
            // are not interested in the events themselves, but only whether at least one has
            // happened.
            let (events_tx, events_rx) = sync_channel(1);
            let (running_tx, running_rx) = channel();
            let (watching_tx, watching_rx) = sync_channel(1);

            Worker::run(user_toml_path, events_tx, running_rx, watching_tx)?;

            outputln!(preamble service.service_group(), "Watching {}", USER_CONFIG_FILE);

            let state = WorkerState { have_events:      events_rx,
                                      stop_running:     running_tx,
                                      started_watching: watching_rx, };

            states.insert(service.name().to_owned(), state);
        }

        Ok(())
    }

    /// Removes a service from the User Config Watcher, and sends a message to the watcher thread
    /// to stop running.
    pub fn remove<T: Serviceable>(&mut self, service: &T) {
        if let Some(state) = self.states
                                 .lock()
                                 .expect("states lock was poisoned")
                                 .remove(service.name())
        {
            if let Err(e) = state.stop_running.send(()) {
                debug!("Error stopping user-config watcher thread for service {}: {:?}",
                       service.name(),
                       e);
            }
        }
    }

    /// Checks whether the watcher for the specified service has observed any events.
    ///
    /// This also consumes the events.
    pub fn have_events_for<T: Serviceable>(&self, service: &T) -> bool {
        if let Some(state) = self.states
                                 .lock()
                                 .expect("states lock was poisoned")
                                 .get(service.name())
        {
            let rx = &state.have_events;

            match rx.try_recv() {
                Ok(_) => {
                    return true;
                }
                Err(TryRecvError::Empty) => return false,
                Err(TryRecvError::Disconnected) => {
                    debug!("UserConfigWatcher worker has died!");
                    return false;
                }
            }
        }

        false
    }
}

struct UserConfigCallbacks {
    have_events: SyncSender<()>,
}

impl Callbacks for UserConfigCallbacks {
    fn file_appeared(&mut self, _: &Path) { self.perform(); }

    fn file_modified(&mut self, _: &Path) { self.perform(); }

    fn file_disappeared(&mut self, _: &Path) { self.perform(); }
}

impl UserConfigCallbacks {
    fn perform(&self) {
        if let Err(TrySendError::Disconnected(_)) = self.have_events.try_send(()) {
            debug!("Worker could not notify Manager of event");
        }
    }
}

struct Worker;

impl Worker {
    // starts a new thread with the file watcher tracking the service's user-config file
    pub fn run(path: PathBuf,
               have_events: SyncSender<()>,
               stop_running: Receiver<()>,
               started_watching: SyncSender<()>)
               -> io::Result<()> {
        ThreadBuilder::new().name(format!("user-config-watcher-{}", path.display()))
                            .spawn(move || -> liveliness_checker::ThreadUnregistered<(), String> {
                                let checked_thread = liveliness_checker::mark_thread_alive();

                                debug!("UserConfigWatcher({}) worker thread starting",
                                       path.display(),);
                                let callbacks = UserConfigCallbacks { have_events };
                                let mut file_watcher =
                                    match create_file_watcher(&path,
                                                               callbacks,
                                                               false)
                                    {
                                        Ok(w) => w,
                                        Err(e) => {
                                            let msg = format!("UserConfigWatcher({}) could not \
                                                               start notifier, ending thread ({})",
                                                              path.display(),
                                                              e,);
                                            outputln!("{}", msg);
                                            return checked_thread.unregister(Err(msg));
                                        }
                                    };

                                let _ = started_watching.try_send(());

                                loop {
                                    let checked_thread = liveliness_checker::mark_thread_alive();

                                    match stop_running.try_recv() {
                                        // As long as the `stop_running` channel is
                                        // empty, this branch will execute on every
                                        // iteration.
                                        Err(TryRecvError::Empty) => {
                                            if let Err(e) = file_watcher.single_iteration() {
                                                let msg = format!("UserConfigWatcher({}) could \
                                                                   not run notifier, ending \
                                                                   thread ({})",
                                                                  path.display(),
                                                                  e,);
                                                outputln!("{}", msg);
                                                break checked_thread.unregister(Err(msg));
                                            };
                                        }

                                        // If we receive a message on the channel, we stop.
                                        Ok(_) => break checked_thread.unregister(Ok(())),

                                        // If the channel is disconnected, we stop as well.
                                        Err(TryRecvError::Disconnected) => {
                                            let msg = format!("UserConfigWatcher({}) worker \
                                                               thread failed to receive on \
                                                               channel",
                                                              path.display(),);
                                            debug!("{}", msg);
                                            break checked_thread.unregister(Err(msg));
                                        }
                                    }

                                    thread::sleep(Duration::from_secs(1));
                                }
                            })?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use habitat_core::locked_env_var;
    use std::{fs::{remove_file,
                   File},
              io::Write,
              str::FromStr,
              thread,
              time::{Duration,
                     Instant}};

    use crate::manager::file_watcher::WATCHER_DELAY_MS;

    use tempfile::TempDir;

    locked_env_var!(HAB_STUDIO_HOST_ARCH, lock_env_var);

    #[test]
    fn no_events_at_first() {
        let lock = lock_env_var();
        lock.set("");

        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");

        lock.unset();

        assert!(!ucm.have_events_for(&service));
    }

    #[test]
    fn no_events_at_first_poll_watcher() {
        let lock = lock_env_var();
        lock.set("aarch64-darwin");

        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");

        lock.unset();

        assert!(!ucm.have_events_for(&service));
    }

    #[test]
    fn events_present_after_adding_config() {
        let lock = lock_env_var();
        lock.set("");

        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));

        File::create(service.user_config_path().get_path().join(USER_CONFIG_FILE))
            .expect("creating file");

        lock.unset();

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_adding_config_poll_watcher() {
        let lock = lock_env_var();
        lock.set("aarch64-darwin");

        let service = TestService::default();
        let mut ucm = UserConfigWatcher::new();
        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));

        File::create(service.user_config_path().get_path().join(USER_CONFIG_FILE))
            .expect("creating file");

        lock.unset();

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_changing_config() {
        let lock = lock_env_var();
        lock.set("");

        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        let mut file = File::create(&file_path).expect("creating file");

        file.write_all(b"42").expect(USER_CONFIG_FILE);

        lock.unset();

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_changing_config_poll_watcher() {
        let lock = lock_env_var();
        lock.set("aarch64-darwin");

        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        let mut file = File::create(&file_path).expect("creating file");

        file.write_all(b"42").expect(USER_CONFIG_FILE);

        lock.unset();

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_removing_config() {
        let lock = lock_env_var();
        lock.set("");

        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        File::create(&file_path).expect("creating file");

        // Allow the watcher to notice that a file was created.
        wait_for_events(&ucm, &service);

        remove_file(&file_path).expect("removing file");

        lock.unset();

        assert!(wait_for_events(&ucm, &service));
    }

    #[test]
    fn events_present_after_removing_config_poll_watcher() {
        let lock = lock_env_var();
        lock.set("aarch64-darwin");

        let service = TestService::default();
        let file_path = service.user_config_path().get_path().join(USER_CONFIG_FILE);
        let mut ucm = UserConfigWatcher::new();

        ucm.add(&service).expect("adding service");
        assert!(wait_for_watcher(&ucm, &service));
        File::create(&file_path).expect("creating file");

        // Allow the watcher to notice that a file was created.
        wait_for_events(&ucm, &service);

        remove_file(&file_path).expect("removing file");

        lock.unset();

        assert!(wait_for_events(&ucm, &service));
    }

    fn wait_for_watcher<T: Serviceable>(ucm: &UserConfigWatcher, service: &T) -> bool {
        let start = Instant::now();
        let timeout = Duration::from_secs(10);

        while start.elapsed() < timeout {
            let states = ucm.states.lock().expect("states lock was poisoned");
            let state = states.get(service.name()).expect("service added");
            match state.started_watching.try_recv() {
                Ok(_) => {
                    println!("Received data on the start_watching channel. Returning true.");
                    return true;
                }
                Err(TryRecvError::Empty) => {
                    println!("Received nothing on the start_watching channel. Returning ().");
                }
                Err(TryRecvError::Disconnected) => {
                    println!("The start_watching channel was disconnected. Returning false.");
                    return false;
                }
            }

            thread::sleep(Duration::from_millis(100));
        }

        false
    }

    fn wait_for_events<T: Serviceable>(ucm: &UserConfigWatcher, service: &T) -> bool {
        let start = Instant::now();
        let timeout = Duration::from_millis(WATCHER_DELAY_MS * 5);

        while start.elapsed() < timeout {
            if ucm.have_events_for(service) {
                return true;
            }

            thread::sleep(Duration::from_millis(100));
        }

        false
    }

    struct TestService {
        // It is only used as a way to remove the temporary directory
        // when the test finishes. Not used otherwise.
        #[allow(dead_code)]
        tmp:              TempDir,
        name:             String,
        user_config_path: UserConfigPath,
        service_group:    ServiceGroup,
    }

    impl Serviceable for TestService {
        fn name(&self) -> &str { &self.name }

        fn user_config_path(&self) -> &UserConfigPath { &self.user_config_path }

        fn service_group(&self) -> &ServiceGroup { &self.service_group }
    }

    impl Default for TestService {
        fn default() -> Self {
            let tmp = TempDir::new().expect("creating temp dir");
            let path = UserConfigPath::Recommended(tmp.path().to_path_buf());
            Self { tmp,
                   name: String::from("foo"),
                   user_config_path: path,
                   service_group: ServiceGroup::from_str("foo.bar@yoyodine").unwrap() }
        }
    }
}
<|endoftext|>


<commit_before>use serde_json::{self, Value};

/// User ID
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum UserId {
	/// The user is identified by ID
	Id(String),
	/// The user is identified by e-mail address
	Email(String),
}

impl UserId {
	/// Create user ID from `id` given
	pub fn from_id(id: String) -> UserId {
		UserId::Id(id)
	}

	/// Create user ID from `email` given
	pub fn from_email(email: String) -> UserId {
		UserId::Email(email)
	}

	/// Get ID if the user is identified by ID or empty string otherwise.
	pub fn id(&self) -> String {
		match self {
			&UserId::Id(ref id) => id.clone(),
			_ => String::new(),
		}
	}

	/// Get e-mail address if the user is identified by e-mail address or empty string otherwise.
	pub fn email(&self) -> String {
		match self {
			&UserId::Email(ref email) => email.clone(),
			_ => String::new(),
		}
	}

	/// Test if the user identified by ID
	pub fn is_id(&self) -> bool {
		match self {
			&UserId::Id(_) => true,
			_ => false,
		}
	}

	/// Test if the user identified by e-mail address
	pub fn is_email(&self) -> bool {
		match self {
			&UserId::Email(_) => true,
			_ => false,
		}
	}
}

// Convert User ID into JSON Value
impl serde_json::value::ToJson for UserId {
	fn to_json(&self) -> Value {
		use serde_json::value::Map;
		let mut object: Map<String, Value> = Map::new();

		match self {
			&UserId::Id(ref id) => {
				object.insert("Id".to_owned(), Value::String(id.clone()));
			},
			&UserId::Email(ref email) => {
				object.insert("Email".to_owned(), Value::String(email.clone()));
			}
		};

		Value::Object(object)
	}
}

#[cfg(test)]
mod tests {
	use super::UserId;
	use serde_json::value::ToJson;

	#[test]
	fn user_id_from_id() {
		let user = UserId::from_id(String::from("ID!"));
		assert_eq!(user.is_id(), true);
		assert_eq!(user.id(), String::from("ID!"));
		assert_eq!(user.is_email(), false);
		assert_eq!(user.email(), String::new());
	}

	#[test]
	fn user_by_id_to_json() {
		let user = UserId::from_id(String::from("ID!"));
		let value = user.to_json();
		assert!(value.is_object());

		if let Some(object) = value.as_object() {
			match object.get("Id") {
				Some(id) => {
					assert!(id.is_string());
					if let Some(s) = id.as_str() {
						assert_eq!(s, "ID!");
					}
					else {
						panic!("Cannot get contained value");
					}
				},
				None => panic!("Must contain Id property")
			}
		}
		else {
			panic!("Cannot get contained object");
		}
	}

	#[test]
	fn user_id_from_email() {
		let user = UserId::from_email(String::from("email@address.com"));
		assert_eq!(user.is_id(), false);
		assert_eq!(user.id(), String::new());
		assert_eq!(user.is_email(), true);
		assert_eq!(user.email(), String::from("email@address.com"));
	}

	#[test]
	fn user_by_email_to_json() {
		let user = UserId::from_email(String::from("email@address.com"));
		let value = user.to_json();
		assert!(value.is_object());

		if let Some(object) = value.as_object() {
			match object.get("Email") {
				Some(email) => {
					assert!(email.is_string());
					if let Some(s) = email.as_str() {
						assert_eq!(s, "email@address.com");
					}
					else {
						panic!("Cannot get contained value");
					}
				},
				None => panic!("Must contain Id property")
			}
		}
		else {
			panic!("Cannot get contained object");
		}
	}
}
<commit_msg>Change tabs to spaces<commit_after>use serde_json::{self, Value};

/// User ID
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum UserId {
    /// The user is identified by ID
    Id(String),
    /// The user is identified by e-mail address
    Email(String),
}

impl UserId {
    /// Create user ID from `id` given
    pub fn from_id(id: String) -> UserId {
        UserId::Id(id)
    }

    /// Create user ID from `email` given
    pub fn from_email(email: String) -> UserId {
        UserId::Email(email)
    }

    /// Get ID if the user is identified by ID or empty string otherwise.
    pub fn id(&self) -> String {
        match self {
            &UserId::Id(ref id) => id.clone(),
            _ => String::new(),
        }
    }

    /// Get e-mail address if the user is identified by e-mail address or empty string otherwise.
    pub fn email(&self) -> String {
        match self {
            &UserId::Email(ref email) => email.clone(),
            _ => String::new(),
        }
    }

    /// Test if the user identified by ID
    pub fn is_id(&self) -> bool {
        match self {
            &UserId::Id(_) => true,
            _ => false,
        }
    }

    /// Test if the user identified by e-mail address
    pub fn is_email(&self) -> bool {
        match self {
            &UserId::Email(_) => true,
            _ => false,
        }
    }
}

// Convert User ID into JSON Value
impl serde_json::value::ToJson for UserId {
    fn to_json(&self) -> Value {
        use serde_json::value::Map;
        let mut object: Map<String, Value> = Map::new();

        match self {
            &UserId::Id(ref id) => {
                object.insert("Id".to_owned(), Value::String(id.clone()));
            },
            &UserId::Email(ref email) => {
                object.insert("Email".to_owned(), Value::String(email.clone()));
            }
        };

        Value::Object(object)
    }
}

#[cfg(test)]
mod tests {
    use super::UserId;
    use serde_json::value::ToJson;

    #[test]
    fn user_id_from_id() {
        let user = UserId::from_id(String::from("ID!"));
        assert_eq!(user.is_id(), true);
        assert_eq!(user.id(), String::from("ID!"));
        assert_eq!(user.is_email(), false);
        assert_eq!(user.email(), String::new());
    }

    #[test]
    fn user_by_id_to_json() {
        let user = UserId::from_id(String::from("ID!"));
        let value = user.to_json();
        assert!(value.is_object());

        if let Some(object) = value.as_object() {
            match object.get("Id") {
                Some(id) => {
                    assert!(id.is_string());
                    if let Some(s) = id.as_str() {
                        assert_eq!(s, "ID!");
                    }
                    else {
                        panic!("Cannot get contained value");
                    }
                },
                None => panic!("Must contain Id property")
            }
        }
        else {
            panic!("Cannot get contained object");
        }
    }

    #[test]
    fn user_id_from_email() {
        let user = UserId::from_email(String::from("email@address.com"));
        assert_eq!(user.is_id(), false);
        assert_eq!(user.id(), String::new());
        assert_eq!(user.is_email(), true);
        assert_eq!(user.email(), String::from("email@address.com"));
    }

    #[test]
    fn user_by_email_to_json() {
        let user = UserId::from_email(String::from("email@address.com"));
        let value = user.to_json();
        assert!(value.is_object());

        if let Some(object) = value.as_object() {
            match object.get("Email") {
                Some(email) => {
                    assert!(email.is_string());
                    if let Some(s) = email.as_str() {
                        assert_eq!(s, "email@address.com");
                    }
                    else {
                        panic!("Cannot get contained value");
                    }
                },
                None => panic!("Must contain Id property")
            }
        }
        else {
            panic!("Cannot get contained object");
        }
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(any(debug, test), deny(warnings))]

extern crate chrono;
extern crate clap;
extern crate diesel;
extern crate dotenv;
extern crate diesel_infer_schema;

mod database_error;
#[macro_use]
mod database;
mod cli;
mod pretty_printing;

use chrono::*;
use clap::{ArgMatches,Shell};
use diesel::migrations::schema::*;
use diesel::types::{FromSql, VarChar};
use diesel::{migrations, Connection, Insertable};
use std::any::Any;
use std::error::Error;
use std::io::stdout;
use std::path::{PathBuf, Path};
use std::{env, fs};

use self::database_error::{DatabaseError, DatabaseResult};

fn main() {
    use self::dotenv::dotenv;
    dotenv().ok();

    let matches = cli::build_cli().get_matches();

    match matches.subcommand() {
        ("migration", Some(matches)) => run_migration_command(matches),
        ("setup", Some(matches)) => run_setup_command(matches),
        ("database", Some(matches)) => run_database_command(matches),
        ("bash-completion", Some(matches)) => generate_bash_completion_command(matches),
        ("print-schema", Some(matches)) => run_infer_schema(matches),
        _ => unreachable!("The cli parser should prevent reaching here"),
    }
}

fn run_migration_command(matches: &ArgMatches) {
    match matches.subcommand() {
        ("run", Some(args)) => {
            let database_url = database::database_url(matches);
            let dir = migrations_dir(args);
            call_with_conn!(database_url, migrations::run_pending_migrations_in_directory(&dir, &mut stdout()))
                .unwrap_or_else(handle_error);
        }
        ("revert", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, migrations::revert_latest_migration)
                .unwrap_or_else(handle_error);
        }
        ("redo", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, redo_latest_migration);
        }
        ("generate", Some(args)) => {
            let migration_name = args.value_of("MIGRATION_NAME").unwrap();
            let version = migration_version(args);
            let versioned_name = format!("{}_{}", version, migration_name);
            let migration_dir = migrations_dir(args).join(versioned_name);
            fs::create_dir(&migration_dir).unwrap();

            let migration_dir_relative = convert_absolute_path_to_relative(
                &migration_dir,
                &env::current_dir().unwrap()
            );

            let up_path = migration_dir.join("up.sql");
            println!("Creating {}", migration_dir_relative.join("up.sql").display());
            fs::File::create(up_path).unwrap();
            let down_path = migration_dir.join("down.sql");
            println!("Creating {}", migration_dir_relative.join("down.sql").display());
            fs::File::create(down_path).unwrap();
        }
        _ => unreachable!("The cli parser should prevent reaching here"),
    }
}

use std::fmt::Display;
fn migration_version<'a>(matches: &'a ArgMatches) -> Box<Display + 'a> {
    matches.value_of("MIGRATION_VERSION").map(|s| Box::new(s) as Box<Display>)
        .unwrap_or_else(|| Box::new(UTC::now().format("%Y%m%d%H%M%S")))
}

fn migrations_dir(matches: &ArgMatches) -> PathBuf {
    matches.value_of("MIGRATION_DIRECTORY")
        .map(PathBuf::from)
        .or_else(|| {
            env::var("MIGRATION_DIRECTORY").map(PathBuf::from).ok()
        }).unwrap_or_else(|| {
            migrations::find_migrations_directory()
                .unwrap_or_else(handle_error)
        })
}

fn run_setup_command(matches: &ArgMatches) {
    migrations::find_migrations_directory()
        .unwrap_or_else(|_| {
            create_migrations_directory()
                .unwrap_or_else(handle_error)
        });

    database::setup_database(matches).unwrap_or_else(handle_error);
}

fn run_database_command(matches: &ArgMatches) {
    match matches.subcommand() {
        ("setup", Some(args)) => database::setup_database(args).unwrap_or_else(handle_error),
        ("reset", Some(args)) => database::reset_database(args).unwrap_or_else(handle_error),
        ("drop", Some(args)) => database::drop_database_command(args).unwrap_or_else(handle_error),
        _ => unreachable!("The cli parser should prevent reaching here"),
    };
}

fn generate_bash_completion_command(_: &ArgMatches) {
    cli::build_cli().gen_completions_to("diesel", Shell::Bash, &mut stdout());
}

/// Looks for a migrations directory in the current path and all parent paths,
/// and creates one in the same directory as the Cargo.toml if it can't find
/// one. It also sticks a .gitkeep in the directory so git will pick it up.
/// Returns a `DatabaseError::CargoTomlNotFound` if no Cargo.toml is found.
fn create_migrations_directory() -> DatabaseResult<PathBuf> {
    let project_root = try!(find_project_root());
    println!("Creating migrations/ directory at: {}", project_root
                                                        .join("migrations")
                                                        .display());
    try!(fs::create_dir(project_root.join("migrations")));
    try!(fs::File::create(project_root.join("migrations/.gitkeep")));
    Ok(project_root)
}

fn find_project_root() -> DatabaseResult<PathBuf> {
    search_for_cargo_toml_directory(&try!(env::current_dir()))
}

/// Searches for the directory that holds the project's Cargo.toml, and returns
/// the path if it found it, or returns a `DatabaseError::CargoTomlNotFound`.
fn search_for_cargo_toml_directory(path: &Path) -> DatabaseResult<PathBuf> {
    let toml_path = path.join("Cargo.toml");
    if toml_path.is_file() {
        Ok(path.to_owned())
    } else {
        path.parent().map(search_for_cargo_toml_directory)
            .unwrap_or(Err(DatabaseError::CargoTomlNotFound))
    }
}

/// Reverts the most recent migration, and then runs it again, all in a
/// transaction. If either part fails, the transaction is not committed.
fn redo_latest_migration<Conn>(conn: &Conn) where
        Conn: Connection + Any,
        String: FromSql<VarChar, Conn::Backend>,
        for<'a> &'a NewMigration<'a>:
            Insertable<__diesel_schema_migrations::table, Conn::Backend>,
{
    let migration_inner = || {
        let reverted_version = try!(migrations::revert_latest_migration(conn));
        migrations::run_migration_with_version(conn, &reverted_version, &mut stdout())
    };
    if should_redo_migration_in_transaction(conn) {
        conn.transaction(migration_inner).unwrap_or_else(handle_error);
    } else {
        migration_inner().unwrap_or_else(handle_error);
    }
}

#[cfg(feature="mysql")]
fn should_redo_migration_in_transaction(t: &Any) -> bool {
    !t.is::<::diesel::mysql::MysqlConnection>()
}

#[cfg(not(feature="mysql"))]
fn should_redo_migration_in_transaction(t: &Any) -> bool {
    true
}

fn handle_error<E: Error, T>(error: E) -> T {
    println!("{}", error);
    ::std::process::exit(1);
}

// Converts an absolute path to a relative path, with the restriction that the
// target path must be in the same directory or above the current path.
fn convert_absolute_path_to_relative(target_path: &Path, mut current_path: &Path)
    -> PathBuf
{
    let mut result = PathBuf::new();

    while !target_path.starts_with(current_path) {
        result.push("..");
        match current_path.parent() {
            Some(parent) => current_path = parent,
            None => return target_path.into(),
        }
    }

    result.join(target_path.strip_prefix(current_path).unwrap())
}

fn run_infer_schema(matches: &ArgMatches) {
    let database_url = database::database_url(matches);
    let schema_name = matches.value_of("schema");

    let filtering_tables = matches.values_of("table-name").map(|values| {
        values.map(|table_name| {
            if cfg!(not(feature = "postgres")) || table_name.contains('.') {
                String::from(table_name)
            } else {
                format!("{}.{}", schema_name.unwrap_or("public"), table_name)
            }
        }).collect()
    }).unwrap_or_else(::std::collections::HashSet::new);

    let is_whitelist = matches.is_present("whitelist");
    let is_blacklist = matches.is_present("blacklist");

    let table_names = diesel_infer_schema::load_table_names(&database_url, schema_name)
        .expect(&format!("Could not load table names from database `{}`{}",
            database_url,
            if let Some(name) = schema_name {
                format!(" with schema `{}`", name)
            } else {
                "".into()
            }
        ));

    let tables = table_names.iter()
        .filter_map(|table| {
            let table_name = table.to_string();
            if is_whitelist && !filtering_tables.contains(&table_name[..]) {
                return None;
            }
            if is_blacklist && filtering_tables.contains(&table_name[..]) {
                return None;
            }
            Some(diesel_infer_schema::expand_infer_table_from_schema(&database_url, &table)
                .expect(&format!("Could not load table `{}`", table.to_string())))
        });

    let schema = diesel_infer_schema::handle_schema(tables, schema_name);

    let pretty = pretty_printing::format_schema(schema.as_str())
        .expect("Could not write to stdout");

    println!("{}", pretty);
}

#[cfg(test)]
mod tests {
    extern crate tempdir;

    use database_error::DatabaseError;

    use self::tempdir::TempDir;

    use std::fs;
    use std::path::PathBuf;

    use super::convert_absolute_path_to_relative;
    use super::search_for_cargo_toml_directory;

    #[test]
    fn toml_directory_find_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();
        let toml_path = temp_path.join("Cargo.toml");

        fs::File::create(&toml_path).unwrap();

        assert_eq!(Ok(temp_path.clone()), search_for_cargo_toml_directory(&temp_path));
    }

    #[test]
    fn cargo_toml_not_found_if_no_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();

        assert_eq!(Err(DatabaseError::CargoTomlNotFound),
            search_for_cargo_toml_directory(&temp_path));
    }

    #[test]
    fn convert_absolute_path_to_relative_works() {
        assert_eq!(PathBuf::from("migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo")));
        assert_eq!(PathBuf::from("../migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/src")));
        assert_eq!(PathBuf::from("../../../migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/src/controllers/errors")));
        assert_eq!(PathBuf::from("12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/migrations")));
        assert_eq!(PathBuf::from("../12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/migrations/67890_create_post")));
    }
}
<commit_msg>Fix CI warning<commit_after>#![cfg_attr(any(debug, test), deny(warnings))]

extern crate chrono;
extern crate clap;
extern crate diesel;
extern crate dotenv;
extern crate diesel_infer_schema;

mod database_error;
#[macro_use]
mod database;
mod cli;
mod pretty_printing;

use chrono::*;
use clap::{ArgMatches,Shell};
use diesel::migrations::schema::*;
use diesel::types::{FromSql, VarChar};
use diesel::{migrations, Connection, Insertable};
use std::any::Any;
use std::error::Error;
use std::io::stdout;
use std::path::{PathBuf, Path};
use std::{env, fs};

use self::database_error::{DatabaseError, DatabaseResult};

fn main() {
    use self::dotenv::dotenv;
    dotenv().ok();

    let matches = cli::build_cli().get_matches();

    match matches.subcommand() {
        ("migration", Some(matches)) => run_migration_command(matches),
        ("setup", Some(matches)) => run_setup_command(matches),
        ("database", Some(matches)) => run_database_command(matches),
        ("bash-completion", Some(matches)) => generate_bash_completion_command(matches),
        ("print-schema", Some(matches)) => run_infer_schema(matches),
        _ => unreachable!("The cli parser should prevent reaching here"),
    }
}

fn run_migration_command(matches: &ArgMatches) {
    match matches.subcommand() {
        ("run", Some(args)) => {
            let database_url = database::database_url(matches);
            let dir = migrations_dir(args);
            call_with_conn!(database_url, migrations::run_pending_migrations_in_directory(&dir, &mut stdout()))
                .unwrap_or_else(handle_error);
        }
        ("revert", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, migrations::revert_latest_migration)
                .unwrap_or_else(handle_error);
        }
        ("redo", Some(_)) => {
            let database_url = database::database_url(matches);
            call_with_conn!(database_url, redo_latest_migration);
        }
        ("generate", Some(args)) => {
            let migration_name = args.value_of("MIGRATION_NAME").unwrap();
            let version = migration_version(args);
            let versioned_name = format!("{}_{}", version, migration_name);
            let migration_dir = migrations_dir(args).join(versioned_name);
            fs::create_dir(&migration_dir).unwrap();

            let migration_dir_relative = convert_absolute_path_to_relative(
                &migration_dir,
                &env::current_dir().unwrap()
            );

            let up_path = migration_dir.join("up.sql");
            println!("Creating {}", migration_dir_relative.join("up.sql").display());
            fs::File::create(up_path).unwrap();
            let down_path = migration_dir.join("down.sql");
            println!("Creating {}", migration_dir_relative.join("down.sql").display());
            fs::File::create(down_path).unwrap();
        }
        _ => unreachable!("The cli parser should prevent reaching here"),
    }
}

use std::fmt::Display;
fn migration_version<'a>(matches: &'a ArgMatches) -> Box<Display + 'a> {
    matches.value_of("MIGRATION_VERSION").map(|s| Box::new(s) as Box<Display>)
        .unwrap_or_else(|| Box::new(UTC::now().format("%Y%m%d%H%M%S")))
}

fn migrations_dir(matches: &ArgMatches) -> PathBuf {
    matches.value_of("MIGRATION_DIRECTORY")
        .map(PathBuf::from)
        .or_else(|| {
            env::var("MIGRATION_DIRECTORY").map(PathBuf::from).ok()
        }).unwrap_or_else(|| {
            migrations::find_migrations_directory()
                .unwrap_or_else(handle_error)
        })
}

fn run_setup_command(matches: &ArgMatches) {
    migrations::find_migrations_directory()
        .unwrap_or_else(|_| {
            create_migrations_directory()
                .unwrap_or_else(handle_error)
        });

    database::setup_database(matches).unwrap_or_else(handle_error);
}

fn run_database_command(matches: &ArgMatches) {
    match matches.subcommand() {
        ("setup", Some(args)) => database::setup_database(args).unwrap_or_else(handle_error),
        ("reset", Some(args)) => database::reset_database(args).unwrap_or_else(handle_error),
        ("drop", Some(args)) => database::drop_database_command(args).unwrap_or_else(handle_error),
        _ => unreachable!("The cli parser should prevent reaching here"),
    };
}

fn generate_bash_completion_command(_: &ArgMatches) {
    cli::build_cli().gen_completions_to("diesel", Shell::Bash, &mut stdout());
}

/// Looks for a migrations directory in the current path and all parent paths,
/// and creates one in the same directory as the Cargo.toml if it can't find
/// one. It also sticks a .gitkeep in the directory so git will pick it up.
/// Returns a `DatabaseError::CargoTomlNotFound` if no Cargo.toml is found.
fn create_migrations_directory() -> DatabaseResult<PathBuf> {
    let project_root = try!(find_project_root());
    println!("Creating migrations/ directory at: {}", project_root
                                                        .join("migrations")
                                                        .display());
    try!(fs::create_dir(project_root.join("migrations")));
    try!(fs::File::create(project_root.join("migrations/.gitkeep")));
    Ok(project_root)
}

fn find_project_root() -> DatabaseResult<PathBuf> {
    search_for_cargo_toml_directory(&try!(env::current_dir()))
}

/// Searches for the directory that holds the project's Cargo.toml, and returns
/// the path if it found it, or returns a `DatabaseError::CargoTomlNotFound`.
fn search_for_cargo_toml_directory(path: &Path) -> DatabaseResult<PathBuf> {
    let toml_path = path.join("Cargo.toml");
    if toml_path.is_file() {
        Ok(path.to_owned())
    } else {
        path.parent().map(search_for_cargo_toml_directory)
            .unwrap_or(Err(DatabaseError::CargoTomlNotFound))
    }
}

/// Reverts the most recent migration, and then runs it again, all in a
/// transaction. If either part fails, the transaction is not committed.
fn redo_latest_migration<Conn>(conn: &Conn) where
        Conn: Connection + Any,
        String: FromSql<VarChar, Conn::Backend>,
        for<'a> &'a NewMigration<'a>:
            Insertable<__diesel_schema_migrations::table, Conn::Backend>,
{
    let migration_inner = || {
        let reverted_version = try!(migrations::revert_latest_migration(conn));
        migrations::run_migration_with_version(conn, &reverted_version, &mut stdout())
    };
    if should_redo_migration_in_transaction(conn) {
        conn.transaction(migration_inner).unwrap_or_else(handle_error);
    } else {
        migration_inner().unwrap_or_else(handle_error);
    }
}

#[cfg(feature="mysql")]
fn should_redo_migration_in_transaction(t: &Any) -> bool {
    !t.is::<::diesel::mysql::MysqlConnection>()
}

#[cfg(not(feature="mysql"))]
fn should_redo_migration_in_transaction(_t: &Any) -> bool {
    true
}

fn handle_error<E: Error, T>(error: E) -> T {
    println!("{}", error);
    ::std::process::exit(1);
}

// Converts an absolute path to a relative path, with the restriction that the
// target path must be in the same directory or above the current path.
fn convert_absolute_path_to_relative(target_path: &Path, mut current_path: &Path)
    -> PathBuf
{
    let mut result = PathBuf::new();

    while !target_path.starts_with(current_path) {
        result.push("..");
        match current_path.parent() {
            Some(parent) => current_path = parent,
            None => return target_path.into(),
        }
    }

    result.join(target_path.strip_prefix(current_path).unwrap())
}

fn run_infer_schema(matches: &ArgMatches) {
    let database_url = database::database_url(matches);
    let schema_name = matches.value_of("schema");

    let filtering_tables = matches.values_of("table-name").map(|values| {
        values.map(|table_name| {
            if cfg!(not(feature = "postgres")) || table_name.contains('.') {
                String::from(table_name)
            } else {
                format!("{}.{}", schema_name.unwrap_or("public"), table_name)
            }
        }).collect()
    }).unwrap_or_else(::std::collections::HashSet::new);

    let is_whitelist = matches.is_present("whitelist");
    let is_blacklist = matches.is_present("blacklist");

    let table_names = diesel_infer_schema::load_table_names(&database_url, schema_name)
        .expect(&format!("Could not load table names from database `{}`{}",
            database_url,
            if let Some(name) = schema_name {
                format!(" with schema `{}`", name)
            } else {
                "".into()
            }
        ));

    let tables = table_names.iter()
        .filter_map(|table| {
            let table_name = table.to_string();
            if is_whitelist && !filtering_tables.contains(&table_name[..]) {
                return None;
            }
            if is_blacklist && filtering_tables.contains(&table_name[..]) {
                return None;
            }
            Some(diesel_infer_schema::expand_infer_table_from_schema(&database_url, &table)
                .expect(&format!("Could not load table `{}`", table.to_string())))
        });

    let schema = diesel_infer_schema::handle_schema(tables, schema_name);

    let pretty = pretty_printing::format_schema(schema.as_str())
        .expect("Could not write to stdout");

    println!("{}", pretty);
}

#[cfg(test)]
mod tests {
    extern crate tempdir;

    use database_error::DatabaseError;

    use self::tempdir::TempDir;

    use std::fs;
    use std::path::PathBuf;

    use super::convert_absolute_path_to_relative;
    use super::search_for_cargo_toml_directory;

    #[test]
    fn toml_directory_find_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();
        let toml_path = temp_path.join("Cargo.toml");

        fs::File::create(&toml_path).unwrap();

        assert_eq!(Ok(temp_path.clone()), search_for_cargo_toml_directory(&temp_path));
    }

    #[test]
    fn cargo_toml_not_found_if_no_cargo_toml() {
        let dir = TempDir::new("diesel").unwrap();
        let temp_path = dir.path().canonicalize().unwrap();

        assert_eq!(Err(DatabaseError::CargoTomlNotFound),
            search_for_cargo_toml_directory(&temp_path));
    }

    #[test]
    fn convert_absolute_path_to_relative_works() {
        assert_eq!(PathBuf::from("migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo")));
        assert_eq!(PathBuf::from("../migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/src")));
        assert_eq!(PathBuf::from("../../../migrations/12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/src/controllers/errors")));
        assert_eq!(PathBuf::from("12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/migrations")));
        assert_eq!(PathBuf::from("../12345_create_user"),
                   convert_absolute_path_to_relative(&PathBuf::from("projects/foo/migrations/12345_create_user"),
                                                     &PathBuf::from("projects/foo/migrations/67890_create_post")));
    }
}
<|endoftext|>


<commit_before>#![no_std]

use libtock_low_level_debug::{AlertCode, LowLevelDebug};
use libtock_platform::Syscalls;
use libtock_runtime::TockSyscalls;

#[panic_handler]
fn panic_handler(_info: &core::panic::PanicInfo) -> ! {
    // Signal a panic using the LowLevelDebug capsule (if available).
    LowLevelDebug::<TockSyscalls>::print_alert_code(AlertCode::Panic);

    // Exit with a non-zero exit code to indicate failure.
    TockSyscalls::exit_terminate(1);
}
<commit_msg>Use the ErrorCode enum directly<commit_after>#![no_std]

use libtock_low_level_debug::{AlertCode, LowLevelDebug};
use libtock_platform::{ErrorCode, Syscalls};
use libtock_runtime::TockSyscalls;

#[panic_handler]
fn panic_handler(_info: &core::panic::PanicInfo) -> ! {
    // Signal a panic using the LowLevelDebug capsule (if available).
    LowLevelDebug::<TockSyscalls>::print_alert_code(AlertCode::Panic);

    // Exit with a non-zero exit code to indicate failure.
    TockSyscalls::exit_terminate(ErrorCode::Fail as u32);
}
<|endoftext|>