<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use cell::Cell;
use error::{Error};
use fmt;
use marker::Reflect;
use thread;

pub struct Flag { failed: Cell<bool> }

// This flag is only ever accessed with a lock previously held. Note that this
// a totally private structure.
unsafe impl Send for Flag {}
unsafe impl Sync for Flag {}

impl Flag {
    pub const fn new() -> Flag {
        Flag { failed: Cell::new(false) }
    }

    #[inline]
    pub fn borrow(&self) -> LockResult<Guard> {
        let ret = Guard { panicking: thread::panicking() };
        if self.get() {
            Err(PoisonError::new(ret))
        } else {
            Ok(ret)
        }
    }

    #[inline]
    pub fn done(&self, guard: &Guard) {
        if !guard.panicking && thread::panicking() {
            self.failed.set(true);
        }
    }

    #[inline]
    pub fn get(&self) -> bool {
        self.failed.get()
    }
}

pub struct Guard {
    panicking: bool,
}

/// A type of error which can be returned whenever a lock is acquired.
///
/// Both Mutexes and RwLocks are poisoned whenever a thread fails while the lock
/// is held. The precise semantics for when a lock is poisoned is documented on
/// each lock, but once a lock is poisoned then all future acquisitions will
/// return this error.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct PoisonError<T> {
    guard: T,
}

/// An enumeration of possible errors which can occur while calling the
/// `try_lock` method.
#[stable(feature = "rust1", since = "1.0.0")]
pub enum TryLockError<T> {
    /// The lock could not be acquired because another thread failed while holding
    /// the lock.
    #[stable(feature = "rust1", since = "1.0.0")]
    Poisoned(#[cfg_attr(not(stage0), stable(feature = "rust1", since = "1.0.0"))] PoisonError<T>),
    /// The lock could not be acquired at this time because the operation would
    /// otherwise block.
    #[stable(feature = "rust1", since = "1.0.0")]
    WouldBlock,
}

/// A type alias for the result of a lock method which can be poisoned.
///
/// The `Ok` variant of this result indicates that the primitive was not
/// poisoned, and the `Guard` is contained within. The `Err` variant indicates
/// that the primitive was poisoned. Note that the `Err` variant *also* carries
/// the associated guard, and it can be acquired through the `into_inner`
/// method.
#[stable(feature = "rust1", since = "1.0.0")]
pub type LockResult<Guard> = Result<Guard, PoisonError<Guard>>;

/// A type alias for the result of a nonblocking locking method.
///
/// For more information, see `LockResult`. A `TryLockResult` doesn't
/// necessarily hold the associated guard in the `Err` type as the lock may not
/// have been acquired for other reasons.
#[stable(feature = "rust1", since = "1.0.0")]
pub type TryLockResult<Guard> = Result<Guard, TryLockError<Guard>>;

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> fmt::Debug for PoisonError<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        "PoisonError { inner: .. }".fmt(f)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> fmt::Display for PoisonError<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        "poisoned lock: another task failed inside".fmt(f)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T: Send + Reflect> Error for PoisonError<T> {
    fn description(&self) -> &str {
        "poisoned lock: another task failed inside"
    }
}

impl<T> PoisonError<T> {
    /// Creates a `PoisonError`.
    #[stable(feature = "sync_poison", since = "1.2.0")]
    pub fn new(guard: T) -> PoisonError<T> {
        PoisonError { guard: guard }
    }

    /// Consumes this error indicating that a lock is poisoned, returning the
    /// underlying guard to allow access regardless.
    #[stable(feature = "sync_poison", since = "1.2.0")]
    pub fn into_inner(self) -> T { self.guard }

    /// Reaches into this error indicating that a lock is poisoned, returning a
    /// reference to the underlying guard to allow access regardless.
    #[stable(feature = "sync_poison", since = "1.2.0")]
    pub fn get_ref(&self) -> &T { &self.guard }

    /// Reaches into this error indicating that a lock is poisoned, returning a
    /// mutable reference to the underlying guard to allow access regardless.
    #[stable(feature = "sync_poison", since = "1.2.0")]
    pub fn get_mut(&mut self) -> &mut T { &mut self.guard }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> From<PoisonError<T>> for TryLockError<T> {
    fn from(err: PoisonError<T>) -> TryLockError<T> {
        TryLockError::Poisoned(err)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> fmt::Debug for TryLockError<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            TryLockError::Poisoned(..) => "Poisoned(..)".fmt(f),
            TryLockError::WouldBlock => "WouldBlock".fmt(f)
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T: Send + Reflect> fmt::Display for TryLockError<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.description().fmt(f)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T: Send + Reflect> Error for TryLockError<T> {
    fn description(&self) -> &str {
        match *self {
            TryLockError::Poisoned(ref p) => p.description(),
            TryLockError::WouldBlock => "try_lock failed because the operation would block"
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            TryLockError::Poisoned(ref p) => Some(p),
            _ => None
        }
    }
}

pub fn map_result<T, U, F>(result: LockResult<T>, f: F)
                           -> LockResult<U>
                           where F: FnOnce(T) -> U {
    match result {
        Ok(t) => Ok(f(t)),
        Err(PoisonError { guard }) => Err(PoisonError::new(f(guard)))
    }
}
<commit_msg>Remove unnecessary bounds on Error and Display implementations for TryLockError and PoisonError.<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use cell::Cell;
use error::{Error};
use fmt;
use marker::Reflect;
use thread;

pub struct Flag { failed: Cell<bool> }

// This flag is only ever accessed with a lock previously held. Note that this
// a totally private structure.
unsafe impl Send for Flag {}
unsafe impl Sync for Flag {}

impl Flag {
    pub const fn new() -> Flag {
        Flag { failed: Cell::new(false) }
    }

    #[inline]
    pub fn borrow(&self) -> LockResult<Guard> {
        let ret = Guard { panicking: thread::panicking() };
        if self.get() {
            Err(PoisonError::new(ret))
        } else {
            Ok(ret)
        }
    }

    #[inline]
    pub fn done(&self, guard: &Guard) {
        if !guard.panicking && thread::panicking() {
            self.failed.set(true);
        }
    }

    #[inline]
    pub fn get(&self) -> bool {
        self.failed.get()
    }
}

pub struct Guard {
    panicking: bool,
}

/// A type of error which can be returned whenever a lock is acquired.
///
/// Both Mutexes and RwLocks are poisoned whenever a thread fails while the lock
/// is held. The precise semantics for when a lock is poisoned is documented on
/// each lock, but once a lock is poisoned then all future acquisitions will
/// return this error.
#[stable(feature = "rust1", since = "1.0.0")]
pub struct PoisonError<T> {
    guard: T,
}

/// An enumeration of possible errors which can occur while calling the
/// `try_lock` method.
#[stable(feature = "rust1", since = "1.0.0")]
pub enum TryLockError<T> {
    /// The lock could not be acquired because another thread failed while holding
    /// the lock.
    #[stable(feature = "rust1", since = "1.0.0")]
    Poisoned(#[cfg_attr(not(stage0), stable(feature = "rust1", since = "1.0.0"))] PoisonError<T>),
    /// The lock could not be acquired at this time because the operation would
    /// otherwise block.
    #[stable(feature = "rust1", since = "1.0.0")]
    WouldBlock,
}

/// A type alias for the result of a lock method which can be poisoned.
///
/// The `Ok` variant of this result indicates that the primitive was not
/// poisoned, and the `Guard` is contained within. The `Err` variant indicates
/// that the primitive was poisoned. Note that the `Err` variant *also* carries
/// the associated guard, and it can be acquired through the `into_inner`
/// method.
#[stable(feature = "rust1", since = "1.0.0")]
pub type LockResult<Guard> = Result<Guard, PoisonError<Guard>>;

/// A type alias for the result of a nonblocking locking method.
///
/// For more information, see `LockResult`. A `TryLockResult` doesn't
/// necessarily hold the associated guard in the `Err` type as the lock may not
/// have been acquired for other reasons.
#[stable(feature = "rust1", since = "1.0.0")]
pub type TryLockResult<Guard> = Result<Guard, TryLockError<Guard>>;

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> fmt::Debug for PoisonError<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        "PoisonError { inner: .. }".fmt(f)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> fmt::Display for PoisonError<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        "poisoned lock: another task failed inside".fmt(f)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T: Reflect> Error for PoisonError<T> {
    fn description(&self) -> &str {
        "poisoned lock: another task failed inside"
    }
}

impl<T> PoisonError<T> {
    /// Creates a `PoisonError`.
    #[stable(feature = "sync_poison", since = "1.2.0")]
    pub fn new(guard: T) -> PoisonError<T> {
        PoisonError { guard: guard }
    }

    /// Consumes this error indicating that a lock is poisoned, returning the
    /// underlying guard to allow access regardless.
    #[stable(feature = "sync_poison", since = "1.2.0")]
    pub fn into_inner(self) -> T { self.guard }

    /// Reaches into this error indicating that a lock is poisoned, returning a
    /// reference to the underlying guard to allow access regardless.
    #[stable(feature = "sync_poison", since = "1.2.0")]
    pub fn get_ref(&self) -> &T { &self.guard }

    /// Reaches into this error indicating that a lock is poisoned, returning a
    /// mutable reference to the underlying guard to allow access regardless.
    #[stable(feature = "sync_poison", since = "1.2.0")]
    pub fn get_mut(&mut self) -> &mut T { &mut self.guard }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> From<PoisonError<T>> for TryLockError<T> {
    fn from(err: PoisonError<T>) -> TryLockError<T> {
        TryLockError::Poisoned(err)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> fmt::Debug for TryLockError<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            TryLockError::Poisoned(..) => "Poisoned(..)".fmt(f),
            TryLockError::WouldBlock => "WouldBlock".fmt(f)
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T> fmt::Display for TryLockError<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            TryLockError::Poisoned(..) => "poisoned lock: another task failed inside",
            TryLockError::WouldBlock => "try_lock failed because the operation would block"
        }.fmt(f)
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T: Reflect> Error for TryLockError<T> {
    fn description(&self) -> &str {
        match *self {
            TryLockError::Poisoned(ref p) => p.description(),
            TryLockError::WouldBlock => "try_lock failed because the operation would block"
        }
    }

    fn cause(&self) -> Option<&Error> {
        match *self {
            TryLockError::Poisoned(ref p) => Some(p),
            _ => None
        }
    }
}

pub fn map_result<T, U, F>(result: LockResult<T>, f: F)
                           -> LockResult<U>
                           where F: FnOnce(T) -> U {
    match result {
        Ok(t) => Ok(f(t)),
        Err(PoisonError { guard }) => Err(PoisonError::new(f(guard)))
    }
}
<|endoftext|>


<commit_before>#![feature(exit_status)]

extern crate zip;

use std::io;
use std::fs;

fn main()
{
    let args: Vec<_> = std::env::args().collect();
    if args.len() < 2 {
        println!("Usage: {} <filename>", args[0]);
        std::env::set_exit_status(1);
        return;
    }
    let fname = std::path::Path::new(&*args[1]);
    let file = fs::File::open(&fname).unwrap();

    let mut archive = zip::ZipArchive::new(file).unwrap();

    for i in 0..archive.len()
    {
        let mut file = archive.by_index(i).unwrap();
        let outpath = sanitize_filename(file.name());
        println!("{}", outpath.display());

        {
            let comment = file.comment();
            if comment.len() > 0 { println!("  File comment: {}", comment); }
        }

        create_directory(outpath.parent().unwrap_or(std::path::Path::new("")));

        if (&*file.name()).ends_with("/") {
            create_directory(&outpath);
        }
        else {
            write_file(&mut file, &outpath);
        }
    }
}

fn write_file(reader: &mut zip::read::ZipFile, outpath: &std::path::Path)
{
    let mut outfile = fs::File::create(&outpath).unwrap();
    io::copy(reader, &mut outfile).unwrap();
}

fn create_directory(outpath: &std::path::Path)
{
    fs::create_dir_all(&outpath).unwrap();
}

fn sanitize_filename(filename: &str) -> std::path::PathBuf
{
    let no_null_filename = match filename.find('\0') {
        Some(index) => &filename[0..index],
        None => filename,
    };

    std::path::Path::new(no_null_filename)
        .components()
        .filter(|component| *component != std::path::Component::ParentDir)
        .fold(std::path::PathBuf::new("."), |mut path, ref cur| {
            path.push(cur.as_os_str());
            path
        })
}
<commit_msg>Fix the extract test<commit_after>#![feature(exit_status)]

extern crate zip;

use std::io;
use std::fs;

fn main()
{
    let args: Vec<_> = std::env::args().collect();
    if args.len() < 2 {
        println!("Usage: {} <filename>", args[0]);
        std::env::set_exit_status(1);
        return;
    }
    let fname = std::path::Path::new(&*args[1]);
    let file = fs::File::open(&fname).unwrap();

    let mut archive = zip::ZipArchive::new(file).unwrap();

    for i in 0..archive.len()
    {
        let mut file = archive.by_index(i).unwrap();
        let outpath = sanitize_filename(file.name());
        println!("{}", outpath.display());

        {
            let comment = file.comment();
            if comment.len() > 0 { println!("  File comment: {}", comment); }
        }

        create_directory(outpath.parent().unwrap_or(std::path::Path::new("")));

        if (&*file.name()).ends_with("/") {
            create_directory(&outpath);
        }
        else {
            write_file(&mut file, &outpath);
        }
    }
}

fn write_file(reader: &mut zip::read::ZipFile, outpath: &std::path::Path)
{
    let mut outfile = fs::File::create(&outpath).unwrap();
    io::copy(reader, &mut outfile).unwrap();
}

fn create_directory(outpath: &std::path::Path)
{
    fs::create_dir_all(&outpath).unwrap();
}

fn sanitize_filename(filename: &str) -> std::path::PathBuf
{
    let no_null_filename = match filename.find('\0') {
        Some(index) => &filename[0..index],
        None => filename,
    };

    std::path::Path::new(no_null_filename)
        .components()
        .filter(|component| *component != std::path::Component::ParentDir)
        .fold(std::path::PathBuf::new(), |mut path, ref cur| {
            path.push(cur.as_os_str());
            path
        })
}
<|endoftext|>


<commit_before>//! 非同期タスクを管理する。

use crate::sysmod::SystemModules;
use anyhow::Result;
use chrono::prelude::*;
use log::{error, info, trace};
use std::future::Future;
use std::sync::Arc;
use tokio::signal::unix::{signal, SignalKind};
use tokio::sync::mpsc;

type ShutdownTx = mpsc::UnboundedSender<()>;
type ShutdownRx = mpsc::UnboundedReceiver<()>;
/// タスク完了通知側 (複数)
type CompleteSyncChild = mpsc::Sender<()>;
/// タスク完了待ち側 (単数)
type CompleteSyncParent = mpsc::Receiver<()>;

/// [TaskServer] と各非同期タスク間で共有されるデータ。
///
/// インスタンスは1つだけ生成され、[Arc] により共有される。
struct InternalControl {
    /// Tokio ランタイム。
    rt: tokio::runtime::Runtime,
    /// 全システムモジュールのリスト。
    sysmods: SystemModules,
    shutdown_tx: ShutdownTx,
    shutdown_rx: ShutdownRx,
}

/// 各非同期タスクに clone され渡されるコントロールハンドル。
///
/// 各フィールドは private で、メソッドで機能を提供する。
/// [Clone::clone] 可能で、各タスクごとに複製したものが渡される。
#[derive(Clone)]
pub struct Control {
    /// private で [InternalControl] への [Arc] を持つ。
    internal: Arc<InternalControl>,
    /// システムシャットダウン時、全タスクの完了待ちのため、
    /// サーバ側は clone されたこれがすべて drop されるまで待機する。
    complete_sync: Option<CompleteSyncChild>,
}

/// [TaskServer::run] の返す実行終了種別。
pub enum RunResult {
    Shutdown,
    Reboot,
}

/// タスクサーバ本体。
pub struct TaskServer {
    /// 各タスクに clone して渡すオリジナルの [Control]。
    ctrl: Control,
    /// [Control::complete_sync] (Sender) の対向 (Receiver)。
    ///
    /// clone されたすべての [Control::complete_sync] が drop するまで待機する。
    /// <https://tokio.rs/tokio/topics/shutdown>
    complete_sync_wait: CompleteSyncParent,
}

impl Control {
    /// 1回限りのタスクを生成して実行開始する。
    ///
    /// F: [Control] を引数に、T を返す関数。
    /// T: Future<Output = anyhow::Result<()> かつスレッド間移動可能。
    ///
    /// つまり、F は [Control] を引数に、anyhow::Result<()> を返す async function。
    pub fn spawn_oneshot_task<F, T>(&self, name: &str, f: F)
    where
        F: FnOnce(Control) -> T,
        T: Future<Output = Result<()>> + Send + 'static,
    {
        // move するデータを準備する
        let name = name.to_string();
        let ctrl = self.clone();
        let future = f(ctrl);

        self.internal.rt.spawn(async move {
            info!("[{}] start (one-shot)", name);

            let result = future.await;

            if let Err(e) = result {
                error!("[{}] finish (error): {:?}", name, e);
            } else {
                info!("[{}] finish (success)", name);
            }
        });
    }

    /// 周期タスクを生成する。
    ///
    /// wakeup_list: 起動時刻。以下を満たさないと panic する。
    /// * second 以下が 0 である。(分単位)
    /// * 昇順ソート済みである。
    ///
    /// F: [Control] を引数に、T を返す関数。
    /// T: Future<Output = anyhow::Result<()> かつスレッド間移動可能。
    ///
    /// つまり、F は [Control] を引数に、anyhow::Result<()> を返す async function。
    pub fn spawn_periodic_task<F, T>(&self, name: &str, wakeup_list: &[NaiveTime], f: F)
    where
        F: Fn(Control) -> T + Send + 'static,
        T: Future<Output = Result<()>> + Send + 'static,
    {
        // move するデータを準備する
        let name = name.to_string();
        let ctrl = self.clone();

        // 空でなくソート済み、秒以下がゼロなのを確認後
        // 今日のその時刻からなる Local DateTime に変換する
        // TODO: is_sorted() がまだ unstable
        assert!(!wakeup_list.is_empty(), "time list is empty");
        let sorted = wakeup_list.windows(2).all(|t| t[0] <= t[1]);
        assert!(sorted, "time list is not sorted");
        let today = Local::today();
        let mut dt_list: Vec<_> = wakeup_list
            .iter()
            .map(|time| {
                assert_eq!(time.second(), 0);
                assert_eq!(time.nanosecond(), 0);
                today.and_time(*time).unwrap()
            })
            .collect();

        // wakeup time list を最初の LOG_LIMIT 個までログに出力する
        const LOG_LIMIT: usize = 5;
        let log_iter = wakeup_list.iter().take(LOG_LIMIT);
        let mut str = log_iter.enumerate().fold(String::new(), |sum, (i, v)| {
            let str = if i == 0 {
                format!("{}", v)
            } else {
                format!(", {}", v)
            };
            sum + &str
        });
        if wakeup_list.len() > LOG_LIMIT {
            str += &format!(", ... ({} items)", wakeup_list.len());
        }
        info!("[{}] registered as a periodic task", name);
        info!("[{}] wakeup time: {}", name, str);

        self.internal.rt.spawn(async move {
            type CDuration = chrono::Duration;
            type TDuration = tokio::time::Duration;

            loop {
                // 現在時刻を取得して分までに切り捨てる
                let now = Local::now();
                let now_hmd = now.date().and_hms(now.hour(), now.minute(), 0);
                let next_min = now_hmd + CDuration::minutes(1);
                trace!("[{}] periodic task check: {}", name, now_hmd);

                // 起動時刻リスト内で二分探索
                match dt_list.binary_search(&now_hmd) {
                    Ok(ind) => {
                        // 一致するものを発見したので続行
                        trace!("[{}] hit in time list: {}", name, now_hmd);
                    }
                    Err(ind) => {
                        // ind = insertion point
                        trace!("[{}] not found in time list: {}", name, now_hmd);
                        // 起きるべき時刻は dt_list[ind]
                        if ind < dt_list.len() {
                            let target_dt = dt_list[ind] + CDuration::seconds(1);
                            let sleep_duration = target_dt - Local::now();
                            let sleep_sec = sleep_duration.num_seconds().clamp(0, i64::MAX) as u64;
                            trace!("[{}] target: {}, sleep_sec: {}", name, target_dt, sleep_sec);
                            tokio::time::sleep(TDuration::from_secs(sleep_sec)).await;
                            trace!("[{}] wake up", name);
                        } else {
                            // 一番後ろよりも現在時刻が後
                            // 起動時刻リストをすべて1日ずつ後ろにずらす
                            for dt in dt_list.iter_mut() {
                                let tomorrow = dt.date() + CDuration::days(1);
                                let time = dt.time();
                                *dt = tomorrow.and_time(time).unwrap();
                                trace!("[{}] advance time list by 1 day", name);
                            }
                        }
                        continue;
                    }
                }

                let future = f(ctrl.clone());
                info!("[{}] start (periodic)", name);
                let result = future.await;
                if let Err(e) = result {
                    error!("[{}] finish (error): {:?}", name, e);
                } else {
                    info!("[{}] finish (success)", name);
                }

                // 次の "分" を狙って sleep する
                // 目標は安全のため hh:mm:05 を狙う
                let target_dt = next_min + CDuration::seconds(5);
                // タスクの実行に1分以上かかると負になるが、
                // chrono::Duration は負数を許している
                // その場合は 0 に補正する
                let sleep_duration = target_dt - Local::now();
                let sleep_sec = sleep_duration.num_seconds().clamp(0, i64::MAX) as u64;
                trace!("[{}] target: {}, sleep_sec: {}", name, target_dt, sleep_sec);
                tokio::time::sleep(tokio::time::Duration::from_secs(sleep_sec)).await;
                trace!("[{}] wake up", name);
            }
        });
    }

    /// [crate::sysmod::SystemModule] リストへの参照を取得する。
    pub fn sysmods(&self) -> &SystemModules {
        &self.internal.sysmods
    }
}

impl TaskServer {
    /// タスクサーバを生成して初期化する。
    pub fn new(sysmods: SystemModules) -> Self {
        let rt = tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()
            .unwrap();

        let (shutdown_tx, shutdown_rx) = mpsc::unbounded_channel();
        let (complete_sync, complete_sync_wait) = mpsc::channel(1);

        let internal = InternalControl {
            rt,
            sysmods,
            shutdown_tx,
            shutdown_rx,
        };
        let ctrl = Control {
            internal: Arc::new(internal),
            complete_sync: Some(complete_sync),
        };

        TaskServer {
            ctrl,
            complete_sync_wait,
        }
    }

    /// [Control::spawn_oneshot_task] と同じ。
    pub fn spawn_oneshot_task<F, T>(&self, name: &str, f: F)
    where
        F: FnOnce(Control) -> T,
        T: Future<Output = Result<()>> + Send + 'static,
    {
        self.ctrl.spawn_oneshot_task(name, f);
    }

    /// 実行を開始し、完了するまでブロックする。
    ///
    /// self の所有権はを consume するため、一度しか実行できない。
    pub fn run(mut self) -> RunResult {
        // async block へ move するためのコピーを作る
        let ctrl = self.ctrl.clone();
        // オリジナルの complete_sync を self から奪って drop しておく
        drop(self.ctrl.complete_sync.take().unwrap());

        self.ctrl.internal.rt.block_on(async move {
            // SystemModule 全体に on_start イベントを配送
            ctrl.internal.sysmods.on_start(&ctrl).await;

            // この async block をシグナル処理に使う
            let mut sigint = signal(SignalKind::interrupt()).unwrap();
            let mut sigterm = signal(SignalKind::terminate()).unwrap();
            let mut sighup = signal(SignalKind::hangup()).unwrap();

            let run_result;
            tokio::select! {
                _ = sigint.recv() => {
                    info!("[signal] SIGINT");
                    run_result = RunResult::Shutdown;
                },
                _ = sigterm.recv() => {
                    info!("[signal] SIGTERM");
                    run_result = RunResult::Shutdown;
                },
                _ = sighup.recv() => {
                    info!("[signal] SIGHUP");
                    run_result = RunResult::Reboot;
                },
            }

            // 全タスク完了待ち
            // この async block で持っている分の complete_sync を
            // ctrl ごと drop する
            // オリジナルの self.ctrl.complete_sync も drop 済みなので
            // 残りは各 async task に clone された Control
            drop(ctrl);
            // 全 complete_sync (Sender) が drop されるまで待つ
            // その時 None を返す
            // それ以外の用法はしないので、データは送信されてこない
            info!("waiting for all tasks to be completed....");
            let sync_result = self.complete_sync_wait.recv().await;
            assert!(sync_result.is_none());
            info!("OK: all tasks are completed");

            run_result
        })
    }
}
<commit_msg>システムシャットダウン開始を各タスクに伝える<commit_after>//! 非同期タスクを管理する。

use crate::sysmod::SystemModules;
use anyhow::Result;
use chrono::prelude::*;
use log::{error, info, trace};
use std::future::Future;
use std::sync::Arc;
use tokio::signal::unix::{signal, SignalKind};
use tokio::sync::watch;

/// システムシャットダウン開始通知送信側 (単数)
type CancelTx = watch::Sender<bool>;
/// システムシャットダウン開始通知受信側 (複数)
type CancelRx = watch::Receiver<bool>;

/// [TaskServer] と各非同期タスク間で共有されるデータ。
///
/// インスタンスは1つだけ生成され、[Arc] により共有される。
struct InternalControl {
    /// Tokio ランタイム。
    rt: tokio::runtime::Runtime,
    /// 全システムモジュールのリスト。
    sysmods: SystemModules,
}

/// 各非同期タスクに clone され渡されるコントロールハンドル。
///
/// 各フィールドは private で、メソッドで機能を提供する。
/// [Clone::clone] 可能で、各タスクごとに複製したものが渡される。
#[derive(Clone)]
pub struct Control {
    /// private で [InternalControl] への [Arc] を持つ。
    internal: Arc<InternalControl>,
    /// システムシャットダウン時、true が設定送信される。
    ///
    /// また、シャットダウンシーケンスにおいて、全タスクの完了待ちのためにも使う。
    /// サーバ側は clone されたこれがすべて drop されるまで待機する。
    cancel_rx: Option<CancelRx>,
}

/// [TaskServer::run] の返す実行終了種別。
pub enum RunResult {
    Shutdown,
    Reboot,
}

/// タスクサーバ本体。
pub struct TaskServer {
    /// 各タスクに clone して渡すオリジナルの [Control]。
    ctrl: Control,
    /// システムシャットダウン時の中断リクエストの送信側。
    /// <https://tokio.rs/tokio/topics/shutdown>
    /// "Telling things to shut down" + "Waiting for things to finish shutting down"
    cancel_tx: CancelTx,
}

impl Control {
    /// 1回限りのタスクを生成して実行開始する。
    ///
    /// F: [Control] を引数に、T を返す関数。
    /// T: Future<Output = anyhow::Result<()> かつスレッド間移動可能。
    ///
    /// つまり、F は [Control] を引数に、anyhow::Result<()> を返す async function。
    pub fn spawn_oneshot_task<F, T>(&self, name: &str, f: F)
    where
        F: FnOnce(Control) -> T,
        T: Future<Output = Result<()>> + Send + 'static,
    {
        // move するデータを準備する
        let name = name.to_string();
        let ctrl = self.clone();
        let future = f(ctrl);

        self.internal.rt.spawn(async move {
            info!("[{}] start (one-shot)", name);

            let result = future.await;

            if let Err(e) = result {
                error!("[{}] finish (error): {:?}", name, e);
            } else {
                info!("[{}] finish (success)", name);
            }
        });
    }

    /// 周期タスクを生成する。
    ///
    /// wakeup_list: 起動時刻。以下を満たさないと panic する。
    /// * second 以下が 0 である。(分単位)
    /// * 昇順ソート済みである。
    ///
    /// F: [Control] を引数に、T を返す関数。
    /// T: Future<Output = anyhow::Result<()> かつスレッド間移動可能。
    ///
    /// つまり、F は [Control] を引数に、anyhow::Result<()> を返す async function。
    pub fn spawn_periodic_task<F, T>(&self, name: &str, wakeup_list: &[NaiveTime], f: F)
    where
        F: Fn(Control) -> T + Send + 'static,
        T: Future<Output = Result<()>> + Send + 'static,
    {
        // move するデータを準備する
        let name = name.to_string();
        let ctrl = self.clone();

        // 空でなくソート済み、秒以下がゼロなのを確認後
        // 今日のその時刻からなる Local DateTime に変換する
        // TODO: is_sorted() がまだ unstable
        assert!(!wakeup_list.is_empty(), "time list is empty");
        let sorted = wakeup_list.windows(2).all(|t| t[0] <= t[1]);
        assert!(sorted, "time list is not sorted");
        let today = Local::today();
        let mut dt_list: Vec<_> = wakeup_list
            .iter()
            .map(|time| {
                assert_eq!(time.second(), 0);
                assert_eq!(time.nanosecond(), 0);
                today.and_time(*time).unwrap()
            })
            .collect();

        // wakeup time list を最初の LOG_LIMIT 個までログに出力する
        const LOG_LIMIT: usize = 5;
        let log_iter = wakeup_list.iter().take(LOG_LIMIT);
        let mut str = log_iter.enumerate().fold(String::new(), |sum, (i, v)| {
            let str = if i == 0 {
                format!("{}", v)
            } else {
                format!(", {}", v)
            };
            sum + &str
        });
        if wakeup_list.len() > LOG_LIMIT {
            str += &format!(", ... ({} items)", wakeup_list.len());
        }
        info!("[{}] registered as a periodic task", name);
        info!("[{}] wakeup time: {}", name, str);

        self.internal.rt.spawn(async move {
            type CDuration = chrono::Duration;
            type TDuration = tokio::time::Duration;

            loop {
                // 現在時刻を取得して分までに切り捨てる
                let now = Local::now();
                let now_hmd = now.date().and_hms(now.hour(), now.minute(), 0);
                let next_min = now_hmd + CDuration::minutes(1);
                trace!("[{}] periodic task check: {}", name, now_hmd);

                // 起動時刻リスト内で二分探索
                match dt_list.binary_search(&now_hmd) {
                    Ok(ind) => {
                        // 一致するものを発見したので続行
                        trace!("[{}] hit in time list: {}", name, now_hmd);
                    }
                    Err(ind) => {
                        // ind = insertion point
                        trace!("[{}] not found in time list: {}", name, now_hmd);
                        // 起きるべき時刻は dt_list[ind]
                        if ind < dt_list.len() {
                            let target_dt = dt_list[ind] + CDuration::seconds(1);
                            let sleep_duration = target_dt - Local::now();
                            let sleep_sec = sleep_duration.num_seconds().clamp(0, i64::MAX) as u64;
                            trace!("[{}] target: {}, sleep_sec: {}", name, target_dt, sleep_sec);
                            tokio::time::sleep(TDuration::from_secs(sleep_sec)).await;
                            trace!("[{}] wake up", name);
                        } else {
                            // 一番後ろよりも現在時刻が後
                            // 起動時刻リストをすべて1日ずつ後ろにずらす
                            for dt in dt_list.iter_mut() {
                                let tomorrow = dt.date() + CDuration::days(1);
                                let time = dt.time();
                                *dt = tomorrow.and_time(time).unwrap();
                                trace!("[{}] advance time list by 1 day", name);
                            }
                        }
                        continue;
                    }
                }

                let future = f(ctrl.clone());
                info!("[{}] start (periodic)", name);
                let result = future.await;
                if let Err(e) = result {
                    error!("[{}] finish (error): {:?}", name, e);
                } else {
                    info!("[{}] finish (success)", name);
                }

                // 次の "分" を狙って sleep する
                // 目標は安全のため hh:mm:05 を狙う
                let target_dt = next_min + CDuration::seconds(5);
                // タスクの実行に1分以上かかると負になるが、
                // chrono::Duration は負数を許している
                // その場合は 0 に補正する
                let sleep_duration = target_dt - Local::now();
                let sleep_sec = sleep_duration.num_seconds().clamp(0, i64::MAX) as u64;
                trace!("[{}] target: {}, sleep_sec: {}", name, target_dt, sleep_sec);
                tokio::time::sleep(tokio::time::Duration::from_secs(sleep_sec)).await;
                trace!("[{}] wake up", name);
            }
        });
    }

    /// [crate::sysmod::SystemModule] リストへの参照を取得する。
    pub fn sysmods(&self) -> &SystemModules {
        &self.internal.sysmods
    }
}

impl TaskServer {
    /// タスクサーバを生成して初期化する。
    pub fn new(sysmods: SystemModules) -> Self {
        let rt = tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()
            .unwrap();

        let (cancel_tx, cancel_rx) = watch::channel(false);

        let internal = InternalControl { rt, sysmods };
        let ctrl = Control {
            internal: Arc::new(internal),
            cancel_rx: Some(cancel_rx),
        };

        TaskServer { ctrl, cancel_tx }
    }

    /// [Control::spawn_oneshot_task] と同じ。
    pub fn spawn_oneshot_task<F, T>(&self, name: &str, f: F)
    where
        F: FnOnce(Control) -> T,
        T: Future<Output = Result<()>> + Send + 'static,
    {
        self.ctrl.spawn_oneshot_task(name, f);
    }

    /// 実行を開始し、完了するまでブロックする。
    ///
    /// self の所有権はを consume するため、一度しか実行できない。
    pub fn run(mut self) -> RunResult {
        // async block へ move するためのコピーを作る
        let ctrl = self.ctrl.clone();
        // オリジナルの cancel_rx を self から奪って drop しておく
        drop(self.ctrl.cancel_rx.take().unwrap());

        self.ctrl.internal.rt.block_on(async move {
            // SystemModule 全体に on_start イベントを配送
            ctrl.internal.sysmods.on_start(&ctrl).await;

            // この async block をシグナル処理に使う
            let mut sigint = signal(SignalKind::interrupt()).unwrap();
            let mut sigterm = signal(SignalKind::terminate()).unwrap();
            let mut sighup = signal(SignalKind::hangup()).unwrap();

            let run_result;
            tokio::select! {
                _ = sigint.recv() => {
                    info!("[signal] SIGINT");
                    run_result = RunResult::Shutdown;
                },
                _ = sigterm.recv() => {
                    info!("[signal] SIGTERM");
                    run_result = RunResult::Shutdown;
                },
                _ = sighup.recv() => {
                    info!("[signal] SIGHUP");
                    run_result = RunResult::Reboot;
                },
            }

            // 値を true に設定して全タスクにキャンセルリクエストを通知する
            self.cancel_tx.send_replace(true);
            // 全タスク完了待ち
            // この async block で持っている分の cancel_rx を ctrl ごと drop する
            // オリジナルの self.ctrl.cancel_rx も drop 済みなので
            // 残りは各 async task に clone された Control
            drop(ctrl);
            // 全 cancel_rx が drop されるまで待つ
            info!("waiting for all tasks to be completed....");
            self.cancel_tx.closed().await;
            info!("OK: all tasks are completed");

            run_result
        })
    }
}
<|endoftext|>


<commit_before>use std::{fs, path::Path};

use bytes::Bytes;
use eyre::{bail, Result};
use oci_spec::image::{Descriptor, ImageManifest};
use tokio::{
    fs::File,
    io::{self, AsyncReadExt, AsyncWriteExt},
};

use http_utils::{reqwest_middleware::RequestBuilder, CLIENT};

/// A client that implements the [OCI Distribution Specification](https://github.com/opencontainers/distribution-spec/blob/main/spec.md)
/// for pulling and pushing images from a container registry
pub struct RegistryClient {
    /// Base URL of the image registry API
    url: String,

    /// Name of the image repository e.g. library/hello-world
    name: String,

    /// Token used to authenticate requests
    token: String,
}

impl RegistryClient {
    /// Create a new client
    fn new(url: &str, name: &str, token: &str) -> Self {
        Self {
            url: url.to_string(),
            name: name.to_string(),
            token: token.to_string(),
        }
    }

    /// Push an image
    ///
    /// Pushes an image from a local directory to the registry. The blobs in `<image_dir>/blobs/sha256` are uploaded first
    /// using `push_blob`. Once that is complete the manifest is uploaded using `push_manifest`.
    ///
    /// # Arguments
    ///
    /// - `reference`: a reference for the image (usually a tag)
    /// - `image_dir`: the image directory following the [OCI Image Layout](https://github.com/opencontainers/image-spec/blob/main/image-layout.md) spec
    pub async fn push_image<P: AsRef<Path>>(&self, reference: &str, image_dir: P) -> Result<()> {
        let image_dir = image_dir.as_ref();

        let manifest_path = image_dir.join("index.json");
        if !manifest_path.exists() {
            bail!("The image directory does not have an `index.json` manifest file")
        }

        let blobs_dir = image_dir.join("blobs").join("sha256");
        if !blobs_dir.exists() {
            bail!("The image directory does not have as `blobs/sha256` sub-directory")
        }

        for entry in blobs_dir.read_dir()?.flatten() {
            self.push_blob(entry.path()).await?;
        }

        self.push_manifest(reference, manifest_path).await?;

        Ok(())
    }

    /// Pull an image
    ///
    /// Pulls an image from the registry to a local directory. The inverse of `push_image`.
    pub async fn pull_image<P: AsRef<Path>>(&self, reference: &str, image_dir: P) -> Result<()> {
        let image_dir = image_dir.as_ref();

        let manifest = self.pull_manifest(reference, image_dir).await?;

        let mut futures = vec![self.pull_blob(manifest.config(), image_dir)];
        for layer in manifest.layers() {
            let future = self.pull_blob(layer, image_dir);
            futures.push(future);
        }

        futures::future::try_join_all(futures).await?;

        Ok(())
    }

    /// Pull a manifest from the registry to a local file and return it
    ///
    /// See https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pulling-manifests
    pub async fn pull_manifest<S: AsRef<str>, P: AsRef<Path>>(
        &self,
        reference: S,
        image_dir: P,
    ) -> Result<ImageManifest> {
        let reference = reference.as_ref();

        tracing::info!(
            "Pulling manifest from repository `{}/{}:{}`",
            self.url,
            self.name,
            reference
        );

        let response = self
            .get(&["/manifests/", reference].concat())
            .header(
                "Accept",
                "application/vnd.docker.distribution.manifest.v2+json",
            )
            .send()
            .await?
            .error_for_status()?;

        let manifest: ImageManifest = response.json().await?;

        fs::create_dir_all(&image_dir)?;
        let manifest_path = image_dir.as_ref().join("index.json");
        let mut manifest_file = fs::File::create(manifest_path)?;
        manifest.to_writer_pretty(&mut manifest_file)?;

        Ok(manifest)
    }

    /// Push a manifest from a local file to the registry
    ///
    /// See https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pushing-manifests
    pub async fn push_manifest<S: AsRef<str>, P: AsRef<Path>>(
        &self,
        reference: S,
        mainfest_path: P,
    ) -> Result<()> {
        let reference = reference.as_ref();

        tracing::info!(
            "Pushing manifest to repository `{}/{}:{}`",
            self.url,
            self.name,
            reference
        );

        let manifest = fs::read_to_string(mainfest_path)?;

        self.put(&["/manifests/", reference].concat())
            .header(
                "Content-Type",
                "application/vnd.docker.distribution.manifest.v2+json",
            )
            .body(manifest)
            .send()
            .await?
            .error_for_status()?;

        Ok(())
    }

    /// Pull a blob from the registry to a local file
    ///
    /// See https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pulling-blobs
    pub async fn pull_blob<P: AsRef<Path>>(
        &self,
        descriptor: &Descriptor,
        image_dir: P,
    ) -> Result<()> {
        let digest = descriptor.digest();
        let size = descriptor.size();

        tracing::info!(
            "Pulling blob `{}` from repository `{}/{}`",
            digest,
            self.url,
            self.name
        );

        let mut response = self
            .get(&["/blobs/", digest].concat())
            .header("Accept", "application/octet-stream")
            .send()
            .await?
            .error_for_status()?;

        if let Some(length) = response
            .headers()
            .get("Content-Length")
            .and_then(|length| length.to_str().ok())
            .and_then(|length| length.parse::<i64>().ok())
        {
            if length != size {
                bail!(
                    "Content-Length header is different from size in descriptor ({} != {})",
                    length,
                    size
                )
            }
        }

        let blobs_dir = image_dir.as_ref().join("blobs").join("sha256");
        fs::create_dir_all(&blobs_dir)?;

        let filename = match digest.strip_prefix("sha256:") {
            Some(sha256) => sha256,
            None => bail!(
                "Expected digest to be prefixed by 'sha256:' but got: {}",
                digest
            ),
        };
        let blob_path = blobs_dir.join(filename);
        let mut file = File::create(blob_path).await?;

        while let Some(chunk) = response.chunk().await? {
            file.write(chunk.as_ref()).await?;
        }

        Ok(())
    }

    /// Push a blob from a local file to the registry
    ///
    /// See https://docs.docker.com/registry/spec/api/#pushing-a-layer for a description of the flow.
    /// See also https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pushing-blobs.
    pub async fn push_blob<P: AsRef<Path>>(&self, blob_path: P) -> Result<()> {
        let filename = blob_path
            .as_ref()
            .file_name()
            .expect("Path to always have a filename");
        let digest = format!("sha256:{}", filename.to_string_lossy());

        // Check that the blob actually needs to be pushed
        let response = self.get(["/blobs/", &digest].concat()).send().await?;
        if response.status() == 200 {
            tracing::info!(
                "Blob `{}` already exists in repository `{}/{}`, will not push",
                digest,
                self.url,
                self.name
            );
            return Ok(());
        }

        // Get the UUID and URL for this upload
        let response = self
            .post("/blobs/uploads/")
            .header("Content-Length", "0")
            .send()
            .await?
            .error_for_status()?;
        let upload_uuid = match response.headers().get("Docker-Upload-UUID") {
            Some(url) => url.to_str()?,
            None => bail!("Did not receive upload UUID from registry"),
        };
        let upload_url = match response.headers().get("Location") {
            Some(url) => url.to_str()?,
            None => bail!("Did not receive upload URL from registry"),
        };

        tracing::info!(
            "Pushing blob `{}` to repository `{}/{}`",
            digest,
            self.url,
            self.name
        );

        const MAX_CHUNK_SIZE: usize = 1000; //5_048_576;

        let mut file = File::open(&blob_path).await?;
        let length = fs::metadata(&blob_path)?.len() as usize;

        // If the blob is less than the maximum chunk size then upload it "monolithically"
        // rather than in chunks (to reduce the number of requests)
        if length <= MAX_CHUNK_SIZE {
            let mut bytes = Vec::with_capacity(length);
            io::copy(&mut file, &mut bytes).await?;

            let upload_url = [
                upload_url,
                if upload_url.contains('?') { "&" } else { "?" },
                "digest=",
                &digest,
            ]
            .concat();

            let response = CLIENT
                .put(upload_url)
                .bearer_auth(self.token.clone())
                .header("Content-Type", "application/octet-stream")
                .body(bytes)
                .send()
                .await?;
            if response.status() != 201 {
                bail!("While uploading blob: {}", response.text().await?)
            }

            return Ok(());
        }

        // Read the blob in chunks and upload each
        let mut buffer = vec![0u8; MAX_CHUNK_SIZE];
        let mut start_byte = 0;
        while let Ok(chunk_length) = file.read(&mut buffer[..]).await {
            if chunk_length == 0 {
                break;
            }

            let chunk = Bytes::copy_from_slice(&buffer[0..chunk_length]);

            tracing::info!(
                "{} {} {} {}",
                start_byte,
                start_byte + chunk_length - 1,
                chunk_length,
                chunk.len()
            );

            let response = CLIENT
                .patch(upload_url)
                .bearer_auth(self.token.clone())
                .header("Content-Type", "application/octet-stream")
                .header(
                    "Content-Range",
                    format!("{}-{}", start_byte, start_byte + chunk_length - 1),
                )
                .header("Content-Length", chunk.len())
                .body(chunk)
                .send()
                .await?;
            if response.status() != 202 {
                bail!(
                    "While uploading chunk: {} {}",
                    response.status(),
                    response.text().await?
                )
            } else {
                let range = response
                    .headers()
                    .get("Range")
                    .and_then(|range| range.to_str().ok())
                    .unwrap_or_default();
                tracing::info!("Uploaded chunk range `{}`", range);
            }

            start_byte += chunk_length;
        }

        // Notify the registry that the upload is complete
        self.put(["/blobs/uploads/", upload_uuid, "?digest=", &digest].concat())
            .send()
            .await?
            .error_for_status()?;

        Ok(())
    }

    /// Make a GET request to the registry
    fn get<S: AsRef<str>>(&self, path: S) -> RequestBuilder {
        CLIENT
            .get([&self.url, "/v2/", &self.name, path.as_ref()].concat())
            .bearer_auth(self.token.clone())
    }

    /// Make a POST request to the registry
    fn post<S: AsRef<str>>(&self, path: S) -> RequestBuilder {
        CLIENT
            .post([&self.url, "/v2/", &self.name, path.as_ref()].concat())
            .bearer_auth(self.token.clone())
    }

    /// Make a PUT request to the registry
    fn put<S: AsRef<str>>(&self, path: S) -> RequestBuilder {
        CLIENT
            .put([&self.url, "/v2/", &self.name, path.as_ref()].concat())
            .bearer_auth(self.token.clone())
    }
}

#[cfg(test)]
mod tests {
    use test_utils::{print_logs_level, skip_ci, tempfile::tempdir};

    use super::*;

    /// Pull and push back the Docker `hello-world` image from the local registry
    ///
    /// To set up this test, run a registry container:
    ///
    ///    docker run -p5000:5000 registry
    ///
    /// Then push `library/hello-world` to that registry:
    ///
    ///    docker pull library/hello-world
    ///    docker tag hello-world localhost:5000/library/hello-world:latest
    ///    docker push localhost:5000/library/hello-world:latest
    #[ignore]
    #[tokio::test]
    async fn hello_world() -> Result<()> {
        skip_ci("Requires an image registry to be running locally");

        print_logs_level(tracing::Level::INFO);

        let image_dir = tempdir()?;

        let client = RegistryClient::new("http://localhost:5000", "library/hello-world", "");
        client.pull_image("latest", &image_dir).await?;
        client.push_image("latest", &image_dir).await?;

        Ok(())
    }
}
<commit_msg>refactor(Images): Refactor client methods and properties<commit_after>use std::{env, fs, path::Path};

use bytes::Bytes;
use eyre::{bail, Result};
use oci_spec::image::{Descriptor, ImageConfiguration, ImageManifest};
use tokio::{
    fs::File,
    io::{self, AsyncReadExt, AsyncWriteExt},
};

use http_utils::{
    reqwest::{Method, Response},
    reqwest_middleware::RequestBuilder,
    CLIENT,
};

/// A client that implements the [OCI Distribution Specification](https://github.com/opencontainers/distribution-spec/blob/main/spec.md)
/// for pulling and pushing images from a container registry
pub struct Client {
    /// URL of the image registry e.g. `registry.fly.io`, `localhost:5000`
    registry_host: String,

    /// Name of the image e.g. `library/hello-world`
    image_name: String,

    /// Token used to authenticate requests
    registry_token: Option<String>,
}

impl Client {
    /// Create a new client
    pub fn new(registry_host: &str, image_name: &str, registry_token: Option<&str>) -> Self {
        let registry_token =
            registry_token
                .map(|str| str.to_string())
                .or_else(|| match registry_host {
                    "registry.hub.docker.io" => env::var("DOCKER_TOKEN")
                        .or_else(|_| env::var("DOCKER_API_TOKEN"))
                        .ok(),
                    "registry.fly.io" => env::var("FLY_TOKEN")
                        .or_else(|_| env::var("FLY_API_TOKEN"))
                        .ok(),
                    _ => None,
                });

        Self {
            registry_host: registry_host.to_string(),
            image_name: image_name.to_string(),
            registry_token,
        }
    }

    /// Push an image
    ///
    /// Pushes an image from a local directory to the registry. The blobs in `<image_dir>/blobs/sha256` are uploaded first
    /// using `push_blob`. Once that is complete the manifest is uploaded using `push_manifest`.
    ///
    /// # Arguments
    ///
    /// - `reference`: a reference for the image (usually a tag)
    /// - `image_dir`: the image directory following the [OCI Image Layout](https://github.com/opencontainers/image-spec/blob/main/image-layout.md) spec
    pub async fn push_image<P: AsRef<Path>>(&self, reference: &str, image_dir: P) -> Result<()> {
        let image_dir = image_dir.as_ref();

        let manifest_path = image_dir.join("index.json");
        if !manifest_path.exists() {
            bail!("The image directory does not have an `index.json` manifest file")
        }

        let blobs_dir = image_dir.join("blobs").join("sha256");
        if !blobs_dir.exists() {
            bail!("The image directory does not have as `blobs/sha256` sub-directory")
        }

        for entry in blobs_dir.read_dir()?.flatten() {
            self.push_blob(entry.path()).await?;
        }

        self.push_manifest(reference, manifest_path).await?;

        Ok(())
    }

    /// Pull an image
    ///
    /// Pulls an image from the registry to a local directory. The inverse of `push_image`.
    pub async fn pull_image<P: AsRef<Path>>(&self, reference: &str, image_dir: P) -> Result<()> {
        let image_dir = image_dir.as_ref();

        let manifest = self.pull_manifest(reference, image_dir).await?;

        let mut futures = vec![self.pull_blob(manifest.config(), image_dir)];
        for layer in manifest.layers() {
            let future = self.pull_blob(layer, image_dir);
            futures.push(future);
        }

        futures::future::try_join_all(futures).await?;

        Ok(())
    }

    /// Get a manifest from the registry
    ///
    /// See https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pulling-manifests
    pub async fn get_manifest<S: AsRef<str>>(&self, reference: S) -> Result<ImageManifest> {
        let reference = reference.as_ref();

        tracing::info!(
            "Pulling manifest from repository `{}/{}:{}`",
            self.registry_host,
            self.image_name,
            reference
        );

        let response = self
            .get(&["/manifests/", reference].concat())
            .header(
                "Accept",
                "application/vnd.docker.distribution.manifest.v2+json",
            )
            .send()
            .await?
            .error_for_status()?;

        let manifest: ImageManifest = response.json().await?;
        Ok(manifest)
    }

    /// Pull a manifest from the registry to a local file and return it
    pub async fn pull_manifest<S: AsRef<str>, P: AsRef<Path>>(
        &self,
        reference: S,
        image_dir: P,
    ) -> Result<ImageManifest> {
        let manifest = self.get_manifest(reference).await?;

        fs::create_dir_all(&image_dir)?;
        let manifest_path = image_dir.as_ref().join("index.json");
        let mut manifest_file = fs::File::create(manifest_path)?;
        manifest.to_writer_pretty(&mut manifest_file)?;

        Ok(manifest)
    }

    /// Push a manifest from a local file to the registry
    ///
    /// See https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pushing-manifests
    pub async fn push_manifest<S: AsRef<str>, P: AsRef<Path>>(
        &self,
        reference: S,
        mainfest_path: P,
    ) -> Result<()> {
        let reference = reference.as_ref();

        tracing::info!(
            "Pushing manifest to repository `{}/{}:{}`",
            self.registry_host,
            self.image_name,
            reference
        );

        let manifest = fs::read_to_string(mainfest_path)?;

        self.put(&["/manifests/", reference].concat())
            .header(
                "Content-Type",
                "application/vnd.docker.distribution.manifest.v2+json",
            )
            .body(manifest)
            .send()
            .await?
            .error_for_status()?;

        Ok(())
    }

    /// Get an image config from the registry
    pub async fn get_config<S: AsRef<str>>(&self, reference: S) -> Result<ImageConfiguration> {
        let manifest = self.get_manifest(reference).await?;
        let descriptor = manifest.config();
        let response = self.get_blob(descriptor).await?;
        let config: ImageConfiguration = response.json().await?;
        Ok(config)
    }

    /// Get a blob from the registry
    ///
    /// See https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pulling-blobs
    pub async fn get_blob(&self, descriptor: &Descriptor) -> Result<Response> {
        let digest = descriptor.digest();
        let size = descriptor.size();

        tracing::info!(
            "Pulling blob `{}` from repository `{}/{}`",
            digest,
            self.registry_host,
            self.image_name
        );

        let response = self
            .get(&["/blobs/", digest].concat())
            .header("Accept", "application/octet-stream")
            .send()
            .await?
            .error_for_status()?;

        if let Some(length) = response
            .headers()
            .get("Content-Length")
            .and_then(|length| length.to_str().ok())
            .and_then(|length| length.parse::<i64>().ok())
        {
            if length != size {
                bail!(
                    "Content-Length header is different from size in descriptor ({} != {})",
                    length,
                    size
                )
            }
        }

        Ok(response)
    }

    /// Pull a blob from the registry to a local file
    pub async fn pull_blob<P: AsRef<Path>>(
        &self,
        descriptor: &Descriptor,
        image_dir: P,
    ) -> Result<()> {
        let mut response = self.get_blob(descriptor).await?;

        let blobs_dir = image_dir.as_ref().join("blobs").join("sha256");
        fs::create_dir_all(&blobs_dir)?;

        let digest = descriptor.digest();
        let filename = match digest.strip_prefix("sha256:") {
            Some(sha256) => sha256,
            None => bail!(
                "Expected digest to be prefixed by 'sha256:' but got: {}",
                digest
            ),
        };
        let blob_path = blobs_dir.join(filename);
        let mut file = File::create(blob_path).await?;

        while let Some(chunk) = response.chunk().await? {
            file.write(chunk.as_ref()).await?;
        }

        Ok(())
    }

    /// Push a blob from a local file to the registry
    ///
    /// See https://docs.docker.com/registry/spec/api/#pushing-a-layer for a description of the flow.
    /// See also https://github.com/opencontainers/distribution-spec/blob/main/spec.md#pushing-blobs.
    pub async fn push_blob<P: AsRef<Path>>(&self, blob_path: P) -> Result<()> {
        let filename = blob_path
            .as_ref()
            .file_name()
            .expect("Path to always have a filename");
        let digest = format!("sha256:{}", filename.to_string_lossy());

        // Check that the blob actually needs to be pushed
        let response = self.get(["/blobs/", &digest].concat()).send().await?;
        if response.status() == 200 {
            tracing::info!(
                "Blob `{}` already exists in repository `{}/{}`, will not push",
                digest,
                self.registry_host,
                self.image_name
            );
            return Ok(());
        }

        // Get the UUID and URL for this upload
        let response = self
            .post("/blobs/uploads/")
            .header("Content-Length", "0")
            .send()
            .await?
            .error_for_status()?;
        let upload_uuid = match response.headers().get("Docker-Upload-UUID") {
            Some(url) => url.to_str()?,
            None => bail!("Did not receive upload UUID from registry"),
        };
        let upload_url = match response.headers().get("Location") {
            Some(url) => url.to_str()?,
            None => bail!("Did not receive upload URL from registry"),
        };

        tracing::info!(
            "Pushing blob `{}` to repository `{}/{}`",
            digest,
            self.registry_host,
            self.image_name
        );

        const MAX_CHUNK_SIZE: usize = 1000; //5_048_576;

        let mut file = File::open(&blob_path).await?;
        let length = fs::metadata(&blob_path)?.len() as usize;

        // If the blob is less than the maximum chunk size then upload it "monolithically"
        // rather than in chunks (to reduce the number of requests)
        if length <= MAX_CHUNK_SIZE {
            let mut bytes = Vec::with_capacity(length);
            io::copy(&mut file, &mut bytes).await?;

            let upload_url = [
                upload_url,
                if upload_url.contains('?') { "&" } else { "?" },
                "digest=",
                &digest,
            ]
            .concat();

            let response = self
                .put(upload_url)
                .header("Content-Type", "application/octet-stream")
                .body(bytes)
                .send()
                .await?;
            if response.status() != 201 {
                bail!("While uploading blob: {}", response.text().await?)
            }

            return Ok(());
        }

        // Read the blob in chunks and upload each
        let mut buffer = vec![0u8; MAX_CHUNK_SIZE];
        let mut start_byte = 0;
        while let Ok(chunk_length) = file.read(&mut buffer[..]).await {
            if chunk_length == 0 {
                break;
            }

            let chunk = Bytes::copy_from_slice(&buffer[0..chunk_length]);

            tracing::info!(
                "{} {} {} {}",
                start_byte,
                start_byte + chunk_length - 1,
                chunk_length,
                chunk.len()
            );

            let response = self
                .patch(upload_url)
                .header("Content-Type", "application/octet-stream")
                .header(
                    "Content-Range",
                    format!("{}-{}", start_byte, start_byte + chunk_length - 1),
                )
                .header("Content-Length", chunk.len())
                .body(chunk)
                .send()
                .await?;
            if response.status() != 202 {
                bail!(
                    "While uploading chunk: {} {}",
                    response.status(),
                    response.text().await?
                )
            } else {
                let range = response
                    .headers()
                    .get("Range")
                    .and_then(|range| range.to_str().ok())
                    .unwrap_or_default();
                tracing::info!("Uploaded chunk range `{}`", range);
            }

            start_byte += chunk_length;
        }

        // Notify the registry that the upload is complete
        self.put(["/blobs/uploads/", upload_uuid, "?digest=", &digest].concat())
            .send()
            .await?
            .error_for_status()?;

        Ok(())
    }

    /// Make a request to the registry
    fn request<S: AsRef<str>>(&self, method: Method, path: S) -> RequestBuilder {
        let path = path.as_ref();

        let url = if path.starts_with("http") {
            path.to_string()
        } else {
            let registry_url = if self.registry_host.starts_with("https://") {
                self.registry_host.clone()
            } else if self.registry_host.starts_with("localhost") {
                ["http://", &self.registry_host].concat()
            } else {
                ["https://", &self.registry_host].concat()
            };
            [&registry_url, "/v2/", &self.image_name, path].concat()
        };

        let mut request = CLIENT.request(method, url);

        if let Some(token) = &self.registry_token {
            request = request.bearer_auth(token);
        }

        request
    }

    /// Make a GET request to the registry
    fn get<S: AsRef<str>>(&self, path: S) -> RequestBuilder {
        self.request(Method::GET, path)
    }

    /// Make a POST request to the registry
    fn post<S: AsRef<str>>(&self, path: S) -> RequestBuilder {
        self.request(Method::POST, path)
    }

    /// Make a PATCH request to the registry
    fn patch<S: AsRef<str>>(&self, path: S) -> RequestBuilder {
        self.request(Method::PATCH, path)
    }

    /// Make a PUT request to the registry
    fn put<S: AsRef<str>>(&self, path: S) -> RequestBuilder {
        self.request(Method::PUT, path)
    }
}

#[cfg(test)]
mod tests {
    use test_utils::{print_logs_level, skip_ci, tempfile::tempdir};

    use super::*;

    /// Pull and push back the Docker `hello-world` image from the local registry
    ///
    /// To set up this test, run a registry container:
    ///
    ///    docker run -p5000:5000 registry
    ///
    /// Then push `library/hello-world` to that registry:
    ///
    ///    docker pull library/hello-world
    ///    docker tag hello-world localhost:5000/library/hello-world:latest
    ///    docker push localhost:5000/library/hello-world:latest
    #[ignore]
    #[tokio::test]
    async fn hello_world() -> Result<()> {
        skip_ci("Requires an image registry to be running locally");

        print_logs_level(tracing::Level::INFO);

        let image_dir = tempdir()?;

        let client = Client::new("localhost:5000", "hello-world", None);
        client.pull_image("latest", &image_dir).await?;
        client.push_image("latest", &image_dir).await?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>
pub( crate ) mod private
{
  use crate::command::*;
  use crate::instruction::*;
  use wtools::meta::*;
  use wtools::error::*;
  use wtools::string::split;
  use wtools::former::Former;

  ///
  /// Commands aggregator.
  ///

  /*
     Dmytro : owned types are used because Former does not work with combination of links and
     containers
   */

  #[ derive( Debug, PartialEq ) ]
  #[ derive( Former ) ]
  #[ allow( missing_docs ) ]
  pub struct CommandsAggregator
  {
    pub base_path : Option<std::path::PathBuf>,
    #[ default( "".to_string() ) ]
    pub command_prefix : String,
    #[ default( vec![ ".".to_string(), " ".to_string() ] ) ]
    pub delimeter : Vec< String >,
    #[ default( ";".to_string() ) ]
    pub command_explicit_delimeter : String,
    #[ default( " ".to_string() ) ]
    pub command_implicit_delimeter : String,
    #[ default( true ) ]
    pub commands_explicit_delimiting : bool,
    #[ default( false ) ]
    pub commands_implicit_delimiting : bool,
    #[ default( false ) ]
    pub properties_map_parsing : bool,
    #[ default( true ) ]
    pub several_values : bool,
    #[ default( true ) ]
    pub with_help : bool,
    #[ default( true ) ]
    pub changing_exit_code : bool,
    // logger : Option<Logger>, /* qqq : implement */
    pub commands : std::collections::HashMap< String, Command >,
    // pub vocabulary : Option<vocabulary>, /* qqq : implement */
  }

  impl CommandsAggregator
  {
    /// Perform instructions queue as single program.
    pub fn program_perform( &self, program : impl AsRef< str > ) -> Result< (), BasicError >
    {
      let program = program.as_ref().trim();

      if !program.starts_with( '.' /* aggregator.vocabulary.default_delimeter */ )
        || program.starts_with( "./" /* `${aggregator.vocabulary.default_delimeter}/` */ )
        || program.starts_with( ".\\" /* `${aggregator.vocabulary.default_delimeter}\\` */ )
      {
        return self.on_syntax_error( program );
      }

      /* should use logger and condition */
      println!( "Command \"{}\"", program );

      let instructions = self.instructions_parse( program );

      for instruction in &instructions
      {
        match self._instruction_perform( instruction )
        {
          Ok( _ ) => {},
          Err( err ) => { return Err( err ) }
        }
      }

      Ok( () )
    }

    /// Perform instruction.
    pub fn instruction_perform( &self, instruction : impl AsRef< str > ) -> Result< (), BasicError >
    {
      let parsed : Instruction = instruction_parse()
      .instruction( instruction.as_ref() )
      .several_values( self.several_values )
      .properties_map_parsing( self.properties_map_parsing )
      .quoting( true )
      .unquoting( true )
      .perform();

      self._instruction_perform( &parsed )
    }

    //

    fn _instruction_perform( &self, instruction : &Instruction ) -> Result< (), BasicError >
    {
      let result = match self.command_resolve( instruction )
      {
        Some( command ) =>
        {
          command.perform( instruction )
        },
        None =>
        {
          if self.with_help
          {
            self.on_get_help().unwrap();
          }
          if self.changing_exit_code
          {
            std::process::exit( 1 );
          }
          Ok( () )
        },
      };

      result
    }

    /// Print help for command.
    fn command_help( &self, command : impl AsRef< str > )
    {
      /* qqq : command should print help, not CA */
      if command.as_ref() == ""
      {
        for ( command_name, command_descriptor ) in self.commands.iter()
        {
          println!( "{} - {}", command_name, command_descriptor.hint );
        }
      }
      else
      {
        if let Some( command_descriptor ) = self.commands.get( command.as_ref() )
        {
          println!( "{} - {}", command.as_ref(), command_descriptor.hint );
        }
      }
    }

    /// Find command in dictionary.
    fn command_resolve( &self, instruction : &Instruction ) -> Option<&Command>
    {
      self.commands.get( &instruction.command_name )
    }

    /// Parse multiple instructions.
    pub fn instructions_parse( &self, program : impl AsRef< str > ) -> Vec< Instruction >
    {
      let commands = split()
      .src( program.as_ref().trim() )
      .delimeter( self.command_explicit_delimeter.as_str() )
      .preserving_empty( false )
      .preserving_delimeters( false )
      .preserving_quoting( false )
      .perform();
      let commands = commands.map( | e | String::from( e ) ).collect::< Vec< _ > >();

      let mut string_commands = vec![];
      for command in commands
      {
        let splitted = split()
        .src( command.trim() )
        .delimeter( self.command_implicit_delimeter.as_str() )
        .preserving_empty( false )
        .preserving_delimeters( false )
        .preserving_quoting( false )
        .perform();
        let splitted = splitted.map( | e | String::from( e ) ).collect::< Vec< _ > >();

        if self.command_implicit_delimeter == " "
        {
          let start_index = if splitted[ 0 ].is_empty() { 1 } else { 0 };
          let mut string_command = String::from( &splitted[ start_index ] );

          for i in start_index + 1 .. splitted.len()
          {
            let part = splitted[ i ].trim();
            if part.starts_with( '.' )
            {
              string_commands.push( string_command );
              string_command = String::from( part );
            }
            else
            {
              string_command.push( ' ' );
              string_command.push_str( part );
            }
          }

          string_commands.push( string_command );
        }
        else
        {
          for command in splitted
          {
            string_commands.push( String::from( command.trim() ) );
          }
        }
      }

      let instructions = string_commands.iter().map( | instruction |
      {
        instruction_parse()
        .instruction( instruction.as_str() )
        .properties_map_parsing( true )
        .several_values( true )
        .perform()
      }).collect::< Vec< Instruction > >();

      instructions
    }
  }

  //

  ///
  /// On error helper.
  ///

  pub trait OnError
  {
    /// Handle error.
    fn on_error( &self, err : BasicError ) -> Result< (), BasicError >;
  }

  ///
  /// On syntax error helper.
  ///

  pub trait OnSyntaxError
  {
    /// Handle syntax error.
    fn on_syntax_error( &self, command : impl AsRef< str > ) -> Result< (), BasicError >;
  }

  ///
  /// On ambiguity helper.
  ///

  pub trait OnAmbiguity
  {
    /// Handle ambiguity.
    fn on_ambiguity( &self, command : impl AsRef< str > ) -> Result< (), BasicError >;
  }

  ///
  /// On unknown command error helper.
  ///

  pub trait OnUnknownCommandError
  {
    /// Handle unknown command error.
    fn on_unknown_command_error( &self, command : impl AsRef< str > ) -> Result< (), BasicError >;
  }

  ///
  /// Help helper.
  ///

  pub trait OnGetHelp
  {
    /// Get help.
    fn on_get_help( &self ) -> Result< (), BasicError >;
  }

  ///
  /// Printing commands helper.
  ///

  pub trait OnPrintCommands
  {
    /// Print all commands.
    fn on_print_commands( &self ) -> Result< (), BasicError >;
  }

  ///
  /// Super trait that checks that all helpers are implemented.
  ///

  pub trait CommandsAggregatorHandlers : OnError + OnSyntaxError + OnAmbiguity + OnUnknownCommandError + OnGetHelp + OnPrintCommands
  {
  }

  impl CommandsAggregatorHandlers for CommandsAggregator {}

  #[ cfg( feature = "on_error_default" ) ]
  impl OnError for CommandsAggregator
  {
    /// Handle error.
    fn on_error( &self, err : BasicError ) -> Result< (), BasicError >
    {
      if self.changing_exit_code
      {
        unimplemented!();
      }
      Err( err )
    }
  }

  #[ cfg( feature = "on_syntax_error_default" ) ]
  impl OnSyntaxError for CommandsAggregator
  {
    /// Handle syntax error.
    fn on_syntax_error( &self, command : impl AsRef< str > ) -> Result< (), BasicError >
    {
      let err_formatted = format!( "Illformed command \"{}\"", command.as_ref() );
      eprintln!( "{}", err_formatted );
      self.on_get_help().unwrap();

      let err = BasicError::new( err_formatted );
      return self.on_error( err );
    }
  }

  #[ cfg( feature = "on_ambiguity_default" ) ]
  impl OnAmbiguity for CommandsAggregator
  {
    /// Handle ambiguity.
    fn on_ambiguity( &self, command : impl AsRef< str > ) -> Result< (), BasicError >
    {
      eprintln!( "Ambiguity. Did you mean?" );
      self.command_help( command.as_ref() );
      println!( "" );

      let err_formatted = format!( "Ambiguity \"{}\"", command.as_ref() );
      let err = BasicError::new( err_formatted );
      return self.on_error( err );
    }
  }

  #[ cfg( feature = "on_unknown_command_error_default" ) ]
  impl OnUnknownCommandError for CommandsAggregator
  {
    /// Handle unknown command error.
    fn on_unknown_command_error( &self, command : impl AsRef< str > ) -> Result< (), BasicError >
    {
      let mut err_formatted = format!( "Unknown command \"{}\"", command.as_ref() );

      let instruction = instruction_parse()
      .instruction( ".help" )
      .perform();
      if self.command_resolve( &instruction ).is_some()
      {
        err_formatted.push_str( "\nTry \".help\"" );
      }
      let err = BasicError::new( err_formatted );
      return self.on_error( err );
    }
  }

  #[ cfg( feature = "on_get_help_default" ) ]
  impl OnGetHelp for CommandsAggregator
  {
    /// Get help.
    fn on_get_help( &self ) -> Result< (), BasicError >
    {
      let instruction = instruction_parse()
      .instruction( ".help" )
      .perform();
      if let Some( command ) = self.command_resolve( &instruction )
      {
        let instruction = instruction_parse()
        .instruction( "" )
        .perform();
        return command.perform( &instruction );
      }
      else
      {
        self.command_help( "" );
        return Ok( () );
      }
    }
  }

  #[ cfg( feature = "on_print_commands_default" ) ]
  impl OnPrintCommands for CommandsAggregator
  {
    /// Print all commands.
    fn on_print_commands( &self ) -> Result< (), BasicError >
    {
      println!( "" );
      self.command_help( "" );
      println!( "" );
      Ok( () )
    }
  }

  //

  ///
  /// Get instruction parser builder.
  ///

  pub fn commands_aggregator() -> CommandsAggregatorFormer
  {
    CommandsAggregator::former()
  }

}

/// Protected namespace of the module.
pub mod protected
{
  pub use super::private::CommandsAggregator;
  pub use super::private::OnError;
  pub use super::private::OnSyntaxError;
  pub use super::private::OnAmbiguity;
  pub use super::private::OnUnknownCommandError;
  pub use super::private::OnGetHelp;
  pub use super::private::OnPrintCommands;
  pub use super::private::commands_aggregator;
}

pub use protected::*;

/// Exposed namespace of the module.
pub mod exposed
{
  pub use super::prelude::*;
}

/// Namespace of the module to include with `use module::*`.
pub mod prelude
{
  pub use super::private::CommandsAggregator;
  pub use super::private::OnError;
  pub use super::private::OnSyntaxError;
  pub use super::private::OnAmbiguity;
  pub use super::private::OnUnknownCommandError;
  pub use super::private::OnGetHelp;
  pub use super::private::OnPrintCommands;
  pub use super::private::commands_aggregator;
}

<commit_msg>Extend routine `command_help`, use methods of substructures<commit_after>
pub( crate ) mod private
{
  use crate::command::*;
  use crate::instruction::*;
  use wtools::meta::*;
  use wtools::error::*;
  use wtools::string::split;
  use wtools::former::Former;

  ///
  /// Commands aggregator.
  ///

  /*
     Dmytro : owned types are used because Former does not work with combination of links and
     containers
   */

  #[ derive( Debug, PartialEq ) ]
  #[ derive( Former ) ]
  #[ allow( missing_docs ) ]
  pub struct CommandsAggregator
  {
    pub base_path : Option<std::path::PathBuf>,
    #[ default( "".to_string() ) ]
    pub command_prefix : String,
    #[ default( vec![ ".".to_string(), " ".to_string() ] ) ]
    pub delimeter : Vec< String >,
    #[ default( ";".to_string() ) ]
    pub command_explicit_delimeter : String,
    #[ default( " ".to_string() ) ]
    pub command_implicit_delimeter : String,
    #[ default( true ) ]
    pub commands_explicit_delimiting : bool,
    #[ default( false ) ]
    pub commands_implicit_delimiting : bool,
    #[ default( false ) ]
    pub properties_map_parsing : bool,
    #[ default( true ) ]
    pub several_values : bool,
    #[ default( true ) ]
    pub with_help : bool,
    #[ default( true ) ]
    pub changing_exit_code : bool,
    // logger : Option<Logger>, /* qqq : implement */
    pub commands : std::collections::HashMap< String, Command >,
    // pub vocabulary : Option<vocabulary>, /* qqq : implement */
  }

  impl CommandsAggregator
  {
    /// Perform instructions queue as single program.
    pub fn program_perform( &self, program : impl AsRef< str > ) -> Result< (), BasicError >
    {
      let program = program.as_ref().trim();

      if !program.starts_with( '.' /* aggregator.vocabulary.default_delimeter */ )
        || program.starts_with( "./" /* `${aggregator.vocabulary.default_delimeter}/` */ )
        || program.starts_with( ".\\" /* `${aggregator.vocabulary.default_delimeter}\\` */ )
      {
        return self.on_syntax_error( program );
      }

      /* should use logger and condition */
      println!( "Command \"{}\"", program );

      let instructions = self.instructions_parse( program );

      for instruction in &instructions
      {
        match self._instruction_perform( instruction )
        {
          Ok( _ ) => {},
          Err( err ) => { return Err( err ) }
        }
      }

      Ok( () )
    }

    /// Perform instruction.
    pub fn instruction_perform( &self, instruction : impl AsRef< str > ) -> Result< (), BasicError >
    {
      let parsed : Instruction = instruction_parse()
      .instruction( instruction.as_ref() )
      .several_values( self.several_values )
      .properties_map_parsing( self.properties_map_parsing )
      .quoting( true )
      .unquoting( true )
      .perform();

      self._instruction_perform( &parsed )
    }

    //

    fn _instruction_perform( &self, instruction : &Instruction ) -> Result< (), BasicError >
    {
      let result = match self.command_resolve( instruction )
      {
        Some( command ) =>
        {
          command.perform( instruction )
        },
        None =>
        {
          if self.with_help
          {
            match self.on_ambiguity( instruction.command_name.as_str() )
            {
              _ => (),
            }
          }
          if self.changing_exit_code
          {
            std::process::exit( 1 );
          }
          Ok( () )
        },
      };

      result
    }

    /// Print help for command.
    fn command_help( &self, command : impl AsRef< str > )
    {
      if command.as_ref() == ""
      {
        for ( _name, command_descriptor ) in self.commands.iter()
        {
          println!( "{}", command_descriptor.help_short() );
        }
      }
      else
      {
        if let Some( command_descriptor ) = self.commands.get( command.as_ref() )
        {
          println!( "{}", command_descriptor.help_long() );
        }
        else
        {
          match self.on_unknown_command_error( command.as_ref() )
          {
            _ => ()
          };
        }
      }
    }

    /// Find command in dictionary.
    fn command_resolve( &self, instruction : &Instruction ) -> Option<&Command>
    {
      self.commands.get( &instruction.command_name )
    }

    /// Parse multiple instructions.
    pub fn instructions_parse( &self, program : impl AsRef< str > ) -> Vec< Instruction >
    {
      let commands = split()
      .src( program.as_ref().trim() )
      .delimeter( self.command_explicit_delimeter.as_str() )
      .preserving_empty( false )
      .preserving_delimeters( false )
      .preserving_quoting( false )
      .perform();
      let commands = commands.map( | e | String::from( e ) ).collect::< Vec< _ > >();

      let mut string_commands = vec![];
      for command in commands
      {
        let splitted = split()
        .src( command.trim() )
        .delimeter( self.command_implicit_delimeter.as_str() )
        .preserving_empty( false )
        .preserving_delimeters( false )
        .preserving_quoting( false )
        .perform();
        let splitted = splitted.map( | e | String::from( e ) ).collect::< Vec< _ > >();

        if self.command_implicit_delimeter == " "
        {
          let start_index = if splitted[ 0 ].is_empty() { 1 } else { 0 };
          let mut string_command = String::from( &splitted[ start_index ] );

          for i in start_index + 1 .. splitted.len()
          {
            let part = splitted[ i ].trim();
            if part.starts_with( '.' )
            {
              string_commands.push( string_command );
              string_command = String::from( part );
            }
            else
            {
              string_command.push( ' ' );
              string_command.push_str( part );
            }
          }

          string_commands.push( string_command );
        }
        else
        {
          for command in splitted
          {
            string_commands.push( String::from( command.trim() ) );
          }
        }
      }

      let instructions = string_commands.iter().map( | instruction |
      {
        instruction_parse()
        .instruction( instruction.as_str() )
        .properties_map_parsing( true )
        .several_values( true )
        .perform()
      }).collect::< Vec< Instruction > >();

      instructions
    }
  }

  //

  ///
  /// On error helper.
  ///

  pub trait OnError
  {
    /// Handle error.
    fn on_error( &self, err : BasicError ) -> Result< (), BasicError >;
  }

  ///
  /// On syntax error helper.
  ///

  pub trait OnSyntaxError
  {
    /// Handle syntax error.
    fn on_syntax_error( &self, command : impl AsRef< str > ) -> Result< (), BasicError >;
  }

  ///
  /// On ambiguity helper.
  ///

  pub trait OnAmbiguity
  {
    /// Handle ambiguity.
    fn on_ambiguity( &self, command : impl AsRef< str > ) -> Result< (), BasicError >;
  }

  ///
  /// On unknown command error helper.
  ///

  pub trait OnUnknownCommandError
  {
    /// Handle unknown command error.
    fn on_unknown_command_error( &self, command : impl AsRef< str > ) -> Result< (), BasicError >;
  }

  ///
  /// Help helper.
  ///

  pub trait OnGetHelp
  {
    /// Get help.
    fn on_get_help( &self ) -> Result< (), BasicError >;
  }

  ///
  /// Printing commands helper.
  ///

  pub trait OnPrintCommands
  {
    /// Print all commands.
    fn on_print_commands( &self ) -> Result< (), BasicError >;
  }

  ///
  /// Super trait that checks that all helpers are implemented.
  ///

  pub trait CommandsAggregatorHandlers : OnError + OnSyntaxError + OnAmbiguity + OnUnknownCommandError + OnGetHelp + OnPrintCommands
  {
  }

  impl CommandsAggregatorHandlers for CommandsAggregator {}

  #[ cfg( feature = "on_error_default" ) ]
  impl OnError for CommandsAggregator
  {
    /// Handle error.
    fn on_error( &self, err : BasicError ) -> Result< (), BasicError >
    {
      if self.changing_exit_code
      {
        unimplemented!();
      }
      Err( err )
    }
  }

  #[ cfg( feature = "on_syntax_error_default" ) ]
  impl OnSyntaxError for CommandsAggregator
  {
    /// Handle syntax error.
    fn on_syntax_error( &self, command : impl AsRef< str > ) -> Result< (), BasicError >
    {
      let err_formatted = format!( "Illformed command \"{}\"", command.as_ref() );
      eprintln!( "{}", err_formatted );
      self.on_get_help().unwrap();

      let err = BasicError::new( err_formatted );
      return self.on_error( err );
    }
  }

  #[ cfg( feature = "on_ambiguity_default" ) ]
  impl OnAmbiguity for CommandsAggregator
  {
    /// Handle ambiguity.
    fn on_ambiguity( &self, command : impl AsRef< str > ) -> Result< (), BasicError >
    {
      eprintln!( "Ambiguity. Did you mean?" );
      self.command_help( command.as_ref() );
      println!( "" );

      let err_formatted = format!( "Ambiguity \"{}\"", command.as_ref() );
      let err = BasicError::new( err_formatted );
      return self.on_error( err );
    }
  }

  #[ cfg( feature = "on_unknown_command_error_default" ) ]
  impl OnUnknownCommandError for CommandsAggregator
  {
    /// Handle unknown command error.
    fn on_unknown_command_error( &self, command : impl AsRef< str > ) -> Result< (), BasicError >
    {
      let mut err_formatted = format!( "Unknown command \"{}\"", command.as_ref() );

      let instruction = instruction_parse()
      .instruction( ".help" )
      .perform();
      if self.command_resolve( &instruction ).is_some()
      {
        err_formatted.push_str( "\nTry \".help\"" );
      }
      let err = BasicError::new( err_formatted );
      return self.on_error( err );
    }
  }

  #[ cfg( feature = "on_get_help_default" ) ]
  impl OnGetHelp for CommandsAggregator
  {
    /// Get help.
    fn on_get_help( &self ) -> Result< (), BasicError >
    {
      let instruction = instruction_parse()
      .instruction( ".help" )
      .perform();
      if let Some( command ) = self.command_resolve( &instruction )
      {
        let instruction = instruction_parse()
        .instruction( "" )
        .perform();
        return command.perform( &instruction );
      }
      else
      {
        self.command_help( "" );
        return Ok( () );
      }
    }
  }

  #[ cfg( feature = "on_print_commands_default" ) ]
  impl OnPrintCommands for CommandsAggregator
  {
    /// Print all commands.
    fn on_print_commands( &self ) -> Result< (), BasicError >
    {
      println!( "" );
      self.command_help( "" );
      println!( "" );
      Ok( () )
    }
  }

  //

  ///
  /// Get instruction parser builder.
  ///

  pub fn commands_aggregator() -> CommandsAggregatorFormer
  {
    CommandsAggregator::former()
  }

}

/// Protected namespace of the module.
pub mod protected
{
  pub use super::private::CommandsAggregator;
  pub use super::private::OnError;
  pub use super::private::OnSyntaxError;
  pub use super::private::OnAmbiguity;
  pub use super::private::OnUnknownCommandError;
  pub use super::private::OnGetHelp;
  pub use super::private::OnPrintCommands;
  pub use super::private::commands_aggregator;
}

pub use protected::*;

/// Exposed namespace of the module.
pub mod exposed
{
  pub use super::prelude::*;
}

/// Namespace of the module to include with `use module::*`.
pub mod prelude
{
  pub use super::private::CommandsAggregator;
  pub use super::private::OnError;
  pub use super::private::OnSyntaxError;
  pub use super::private::OnAmbiguity;
  pub use super::private::OnUnknownCommandError;
  pub use super::private::OnGetHelp;
  pub use super::private::OnPrintCommands;
  pub use super::private::commands_aggregator;
}

<|endoftext|>