<commit_before>// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.

// #[PerformanceCriticalPath]
use super::encoded::RawEncodeSnapshot;

use crate::storage::kv::{Cursor, Iterator, ScanMode, Snapshot};
use crate::storage::Statistics;
use crate::storage::{Error, Result};

use engine_traits::{CfName, IterOptions, CF_DEFAULT, DATA_KEY_PREFIX_LEN};
use kvproto::kvrpcpb::{ApiVersion, KeyRange};
use std::time::Duration;
use tikv_util::time::Instant;
use txn_types::{Key, KvPair};
use yatp::task::future::reschedule;

const MAX_TIME_SLICE: Duration = Duration::from_millis(2);
const MAX_BATCH_SIZE: usize = 1024;

pub enum RawStore<S: Snapshot> {
    Vanilla(RawStoreInner<S>),
    Encoded(RawStoreInner<RawEncodeSnapshot<S>>),
}

impl<'a, S: Snapshot> RawStore<S> {
    pub fn new(snapshot: S, api_version: ApiVersion) -> Self {
        if api_version == ApiVersion::V1 {
            RawStore::Vanilla(RawStoreInner::new(snapshot))
        } else {
            RawStore::Encoded(RawStoreInner::new(RawEncodeSnapshot::from_snapshot(
                snapshot,
                api_version,
            )))
        }
    }

    pub fn raw_get_key_value(
        &self,
        cf: CfName,
        key: &Key,
        stats: &mut Statistics,
    ) -> Result<Option<Vec<u8>>> {
        match self {
            RawStore::Vanilla(inner) => inner.raw_get_key_value(cf, key, stats),
            RawStore::Encoded(inner) => inner.raw_get_key_value(cf, key, stats),
        }
    }

    pub fn raw_get_key_ttl(
        &self,
        cf: CfName,
        key: &'a Key,
        stats: &'a mut Statistics,
    ) -> Result<Option<u64>> {
        match self {
            RawStore::Vanilla(_) => panic!("get ttl on non-ttl store"),
            RawStore::Encoded(inner) => inner
                .snapshot
                .get_key_ttl_cf(cf, key, stats)
                .map_err(Error::from),
        }
    }

    pub async fn forward_raw_scan(
        &'a self,
        cf: CfName,
        start_key: &'a Key,
        end_key: Option<&'a Key>,
        limit: usize,
        statistics: &'a mut Statistics,
        key_only: bool,
    ) -> Result<Vec<Result<KvPair>>> {
        let mut option = IterOptions::default();
        if let Some(end) = end_key {
            option.set_upper_bound(end.as_encoded(), DATA_KEY_PREFIX_LEN);
        }
        match self {
            RawStore::Vanilla(inner) => {
                if key_only {
                    option.set_key_only(key_only);
                }
                inner
                    .forward_raw_scan(cf, start_key, limit, statistics, option, key_only)
                    .await
            }
            RawStore::Encoded(inner) => {
                inner
                    .forward_raw_scan(cf, start_key, limit, statistics, option, key_only)
                    .await
            }
        }
    }

    pub async fn reverse_raw_scan(
        &'a self,
        cf: CfName,
        start_key: &'a Key,
        end_key: Option<&'a Key>,
        limit: usize,
        statistics: &'a mut Statistics,
        key_only: bool,
    ) -> Result<Vec<Result<KvPair>>> {
        let mut option = IterOptions::default();
        if let Some(end) = end_key {
            option.set_lower_bound(end.as_encoded(), DATA_KEY_PREFIX_LEN);
        }
        match self {
            RawStore::Vanilla(inner) => {
                if key_only {
                    option.set_key_only(key_only);
                }
                inner
                    .reverse_raw_scan(cf, start_key, limit, statistics, option, key_only)
                    .await
            }
            RawStore::Encoded(inner) => {
                inner
                    .reverse_raw_scan(cf, start_key, limit, statistics, option, key_only)
                    .await
            }
        }
    }
}

pub struct RawStoreInner<S: Snapshot> {
    snapshot: S,
}

impl<'a, S: Snapshot> RawStoreInner<S> {
    pub fn new(snapshot: S) -> Self {
        RawStoreInner { snapshot }
    }

    pub fn raw_get_key_value(
        &self,
        cf: CfName,
        key: &Key,
        stats: &mut Statistics,
    ) -> Result<Option<Vec<u8>>> {
        // no scan_count for this kind of op.
        let key_len = key.as_encoded().len();
        self.snapshot
            .get_cf(cf, key)
            .map(|value| {
                stats.data.flow_stats.read_keys = 1;
                stats.data.flow_stats.read_bytes =
                    key_len + value.as_ref().map(|v| v.len()).unwrap_or(0);
                value
            })
            .map_err(Error::from)
    }

    /// Scan raw keys in [`start_key`, `end_key`), returns at most `limit` keys. If `end_key` is
    /// `None`, it means unbounded.
    ///
    /// If `key_only` is true, the value corresponding to the key will not be read. Only scanned
    /// keys will be returned.
    pub async fn forward_raw_scan(
        &'a self,
        cf: CfName,
        start_key: &'a Key,
        limit: usize,
        statistics: &'a mut Statistics,
        option: IterOptions,
        key_only: bool,
    ) -> Result<Vec<Result<KvPair>>> {
        let mut cursor = Cursor::new(self.snapshot.iter_cf(cf, option)?, ScanMode::Forward, false);
        let statistics = statistics.mut_cf_statistics(cf);
        if !cursor.seek(start_key, statistics)? {
            return Ok(vec![]);
        }
        let mut pairs = vec![];
        let mut row_count = 0;
        let mut time_slice_start = Instant::now();
        while cursor.valid()? && pairs.len() < limit {
            row_count += 1;
            if row_count >= MAX_BATCH_SIZE {
                if time_slice_start.saturating_elapsed() > MAX_TIME_SLICE {
                    reschedule().await;
                    time_slice_start = Instant::now();
                }
                row_count = 0;
            }
            pairs.push(Ok((
                cursor.key(statistics).to_owned(),
                if key_only {
                    vec![]
                } else {
                    cursor.value(statistics).to_owned()
                },
            )));
            cursor.next(statistics);
        }
        Ok(pairs)
    }

    /// Scan raw keys in [`end_key`, `start_key`) in reverse order, returns at most `limit` keys. If
    /// `start_key` is `None`, it means it's unbounded.
    ///
    /// If `key_only` is true, the value
    /// corresponding to the key will not be read out. Only scanned keys will be returned.
    pub async fn reverse_raw_scan(
        &'a self,
        cf: CfName,
        start_key: &'a Key,
        limit: usize,
        statistics: &'a mut Statistics,
        option: IterOptions,
        key_only: bool,
    ) -> Result<Vec<Result<KvPair>>> {
        let mut cursor = Cursor::new(
            self.snapshot.iter_cf(cf, option)?,
            ScanMode::Backward,
            false,
        );
        let statistics = statistics.mut_cf_statistics(cf);
        if !cursor.reverse_seek(start_key, statistics)? {
            return Ok(vec![]);
        }
        let mut pairs = vec![];
        let mut row_count = 0;
        let mut time_slice_start = Instant::now();
        while cursor.valid()? && pairs.len() < limit {
            row_count += 1;
            if row_count >= MAX_BATCH_SIZE {
                if time_slice_start.saturating_elapsed() > MAX_TIME_SLICE {
                    reschedule().await;
                    time_slice_start = Instant::now();
                }
                row_count = 0;
            }
            pairs.push(Ok((
                cursor.key(statistics).to_owned(),
                if key_only {
                    vec![]
                } else {
                    cursor.value(statistics).to_owned()
                },
            )));
            cursor.prev(statistics);
        }
        Ok(pairs)
    }
}

pub async fn raw_checksum_ranges<S: Snapshot>(
    snapshot: S,
    ranges: Vec<KeyRange>,
) -> Result<(u64, u64, u64)> {
    let mut total_bytes = 0;
    let mut total_kvs = 0;
    let mut digest = crc64fast::Digest::new();
    let mut row_count = 0;
    let mut time_slice_start = Instant::now();
    for r in ranges {
        let mut opts = IterOptions::new(None, None, false);
        opts.set_upper_bound(r.get_end_key(), DATA_KEY_PREFIX_LEN);
        let mut iter = snapshot.iter_cf(CF_DEFAULT, opts)?;
        iter.seek(&Key::from_encoded(r.get_start_key().to_vec()))?;
        while iter.valid()? {
            row_count += 1;
            if row_count >= MAX_BATCH_SIZE {
                if time_slice_start.saturating_elapsed() > MAX_TIME_SLICE {
                    reschedule().await;
                    time_slice_start = Instant::now();
                }
                row_count = 0;
            }
            let k = iter.key();
            let v = iter.value();
            digest.write(k);
            digest.write(v);
            total_kvs += 1;
            total_bytes += k.len() + v.len();
            iter.next()?;
        }
    }
    Ok((digest.sum64(), total_kvs, total_bytes as u64))
}
<commit_msg>storage: reduce extra next/pre in row_scan (#11625)<commit_after>// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.

// #[PerformanceCriticalPath]
use super::encoded::RawEncodeSnapshot;

use crate::storage::kv::{Cursor, Iterator, ScanMode, Snapshot};
use crate::storage::Statistics;
use crate::storage::{Error, Result};

use engine_traits::{CfName, IterOptions, CF_DEFAULT, DATA_KEY_PREFIX_LEN};
use kvproto::kvrpcpb::{ApiVersion, KeyRange};
use std::time::Duration;
use tikv_util::time::Instant;
use txn_types::{Key, KvPair};
use yatp::task::future::reschedule;

const MAX_TIME_SLICE: Duration = Duration::from_millis(2);
const MAX_BATCH_SIZE: usize = 1024;

pub enum RawStore<S: Snapshot> {
    Vanilla(RawStoreInner<S>),
    Encoded(RawStoreInner<RawEncodeSnapshot<S>>),
}

impl<'a, S: Snapshot> RawStore<S> {
    pub fn new(snapshot: S, api_version: ApiVersion) -> Self {
        if api_version == ApiVersion::V1 {
            RawStore::Vanilla(RawStoreInner::new(snapshot))
        } else {
            RawStore::Encoded(RawStoreInner::new(RawEncodeSnapshot::from_snapshot(
                snapshot,
                api_version,
            )))
        }
    }

    pub fn raw_get_key_value(
        &self,
        cf: CfName,
        key: &Key,
        stats: &mut Statistics,
    ) -> Result<Option<Vec<u8>>> {
        match self {
            RawStore::Vanilla(inner) => inner.raw_get_key_value(cf, key, stats),
            RawStore::Encoded(inner) => inner.raw_get_key_value(cf, key, stats),
        }
    }

    pub fn raw_get_key_ttl(
        &self,
        cf: CfName,
        key: &'a Key,
        stats: &'a mut Statistics,
    ) -> Result<Option<u64>> {
        match self {
            RawStore::Vanilla(_) => panic!("get ttl on non-ttl store"),
            RawStore::Encoded(inner) => inner
                .snapshot
                .get_key_ttl_cf(cf, key, stats)
                .map_err(Error::from),
        }
    }

    pub async fn forward_raw_scan(
        &'a self,
        cf: CfName,
        start_key: &'a Key,
        end_key: Option<&'a Key>,
        limit: usize,
        statistics: &'a mut Statistics,
        key_only: bool,
    ) -> Result<Vec<Result<KvPair>>> {
        let mut option = IterOptions::default();
        if let Some(end) = end_key {
            option.set_upper_bound(end.as_encoded(), DATA_KEY_PREFIX_LEN);
        }
        match self {
            RawStore::Vanilla(inner) => {
                if key_only {
                    option.set_key_only(key_only);
                }
                inner
                    .forward_raw_scan(cf, start_key, limit, statistics, option, key_only)
                    .await
            }
            RawStore::Encoded(inner) => {
                inner
                    .forward_raw_scan(cf, start_key, limit, statistics, option, key_only)
                    .await
            }
        }
    }

    pub async fn reverse_raw_scan(
        &'a self,
        cf: CfName,
        start_key: &'a Key,
        end_key: Option<&'a Key>,
        limit: usize,
        statistics: &'a mut Statistics,
        key_only: bool,
    ) -> Result<Vec<Result<KvPair>>> {
        let mut option = IterOptions::default();
        if let Some(end) = end_key {
            option.set_lower_bound(end.as_encoded(), DATA_KEY_PREFIX_LEN);
        }
        match self {
            RawStore::Vanilla(inner) => {
                if key_only {
                    option.set_key_only(key_only);
                }
                inner
                    .reverse_raw_scan(cf, start_key, limit, statistics, option, key_only)
                    .await
            }
            RawStore::Encoded(inner) => {
                inner
                    .reverse_raw_scan(cf, start_key, limit, statistics, option, key_only)
                    .await
            }
        }
    }
}

pub struct RawStoreInner<S: Snapshot> {
    snapshot: S,
}

impl<'a, S: Snapshot> RawStoreInner<S> {
    pub fn new(snapshot: S) -> Self {
        RawStoreInner { snapshot }
    }

    pub fn raw_get_key_value(
        &self,
        cf: CfName,
        key: &Key,
        stats: &mut Statistics,
    ) -> Result<Option<Vec<u8>>> {
        // no scan_count for this kind of op.
        let key_len = key.as_encoded().len();
        self.snapshot
            .get_cf(cf, key)
            .map(|value| {
                stats.data.flow_stats.read_keys = 1;
                stats.data.flow_stats.read_bytes =
                    key_len + value.as_ref().map(|v| v.len()).unwrap_or(0);
                value
            })
            .map_err(Error::from)
    }

    /// Scan raw keys in [`start_key`, `end_key`), returns at most `limit` keys. If `end_key` is
    /// `None`, it means unbounded.
    ///
    /// If `key_only` is true, the value corresponding to the key will not be read. Only scanned
    /// keys will be returned.
    pub async fn forward_raw_scan(
        &'a self,
        cf: CfName,
        start_key: &'a Key,
        limit: usize,
        statistics: &'a mut Statistics,
        option: IterOptions,
        key_only: bool,
    ) -> Result<Vec<Result<KvPair>>> {
        if limit == 0 {
            return Ok(vec![]);
        }
        let mut cursor = Cursor::new(self.snapshot.iter_cf(cf, option)?, ScanMode::Forward, false);
        let statistics = statistics.mut_cf_statistics(cf);
        if !cursor.seek(start_key, statistics)? {
            return Ok(vec![]);
        }
        let mut pairs = vec![];
        let mut row_count = 0;
        let mut time_slice_start = Instant::now();
        while cursor.valid()? {
            row_count += 1;
            if row_count >= MAX_BATCH_SIZE {
                if time_slice_start.saturating_elapsed() > MAX_TIME_SLICE {
                    reschedule().await;
                    time_slice_start = Instant::now();
                }
                row_count = 0;
            }
            pairs.push(Ok((
                cursor.key(statistics).to_owned(),
                if key_only {
                    vec![]
                } else {
                    cursor.value(statistics).to_owned()
                },
            )));
            if pairs.len() < limit {
                cursor.next(statistics);
            } else {
                break;
            }
        }
        Ok(pairs)
    }

    /// Scan raw keys in [`end_key`, `start_key`) in reverse order, returns at most `limit` keys. If
    /// `start_key` is `None`, it means it's unbounded.
    ///
    /// If `key_only` is true, the value
    /// corresponding to the key will not be read out. Only scanned keys will be returned.
    pub async fn reverse_raw_scan(
        &'a self,
        cf: CfName,
        start_key: &'a Key,
        limit: usize,
        statistics: &'a mut Statistics,
        option: IterOptions,
        key_only: bool,
    ) -> Result<Vec<Result<KvPair>>> {
        if limit == 0 {
            return Ok(vec![]);
        }
        let mut cursor = Cursor::new(
            self.snapshot.iter_cf(cf, option)?,
            ScanMode::Backward,
            false,
        );
        let statistics = statistics.mut_cf_statistics(cf);
        if !cursor.reverse_seek(start_key, statistics)? {
            return Ok(vec![]);
        }
        let mut pairs = vec![];
        let mut row_count = 0;
        let mut time_slice_start = Instant::now();
        while cursor.valid()? {
            row_count += 1;
            if row_count >= MAX_BATCH_SIZE {
                if time_slice_start.saturating_elapsed() > MAX_TIME_SLICE {
                    reschedule().await;
                    time_slice_start = Instant::now();
                }
                row_count = 0;
            }
            pairs.push(Ok((
                cursor.key(statistics).to_owned(),
                if key_only {
                    vec![]
                } else {
                    cursor.value(statistics).to_owned()
                },
            )));
            if pairs.len() < limit {
                cursor.prev(statistics);
            } else {
                break;
            }
        }
        Ok(pairs)
    }
}

pub async fn raw_checksum_ranges<S: Snapshot>(
    snapshot: S,
    ranges: Vec<KeyRange>,
) -> Result<(u64, u64, u64)> {
    let mut total_bytes = 0;
    let mut total_kvs = 0;
    let mut digest = crc64fast::Digest::new();
    let mut row_count = 0;
    let mut time_slice_start = Instant::now();
    for r in ranges {
        let mut opts = IterOptions::new(None, None, false);
        opts.set_upper_bound(r.get_end_key(), DATA_KEY_PREFIX_LEN);
        let mut iter = snapshot.iter_cf(CF_DEFAULT, opts)?;
        iter.seek(&Key::from_encoded(r.get_start_key().to_vec()))?;
        while iter.valid()? {
            row_count += 1;
            if row_count >= MAX_BATCH_SIZE {
                if time_slice_start.saturating_elapsed() > MAX_TIME_SLICE {
                    reschedule().await;
                    time_slice_start = Instant::now();
                }
                row_count = 0;
            }
            let k = iter.key();
            let v = iter.value();
            digest.write(k);
            digest.write(v);
            total_kvs += 1;
            total_bytes += k.len() + v.len();
            iter.next()?;
        }
    }
    Ok((digest.sum64(), total_kvs, total_bytes as u64))
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    cmp::min,
    collections::{
        BTreeMap,
        Bound::{Included, Unbounded},
    },
};

use devicemapper::Sectors;

use crate::{
    engine::strat_engine::backstore::metadata::BlockdevSize,
    stratis::{ErrorEnum, StratisError, StratisResult},
};

#[derive(Debug)]
pub struct RangeAllocator {
    limit: Sectors,
    used: BTreeMap<Sectors, Sectors>,
}

impl RangeAllocator {
    /// Create a new RangeAllocator with the specified (offset, length)
    /// ranges marked as used.
    #[allow(clippy::new_ret_no_self)]
    pub fn new(
        limit: BlockdevSize,
        initial_used: &[(Sectors, Sectors)],
    ) -> StratisResult<RangeAllocator> {
        let mut allocator = RangeAllocator {
            limit: limit.sectors(),
            used: BTreeMap::new(),
        };
        allocator.insert_ranges(initial_used)?;
        Ok(allocator)
    }

    /// The maximum allocation from this manager
    pub fn size(&self) -> BlockdevSize {
        BlockdevSize::new(self.limit)
    }

    fn check_for_overflow(&self, off: Sectors, len: Sectors) -> StratisResult<()> {
        if let Some(sum) = off.checked_add(len) {
            if sum > self.limit {
                let err_msg = format!(
                    "elements in range ({}, {}) exceed limit {}",
                    off, len, self.limit
                );
                return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
            }
        } else {
            let err_msg = format!(
                "elements in range ({}, {}) inexpressible in this format",
                off, len
            );
            return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
        }
        Ok(())
    }

    /// Mark ranges previously marked as unused as now used.
    /// Return an error if ranges overlap with each other or with previously
    /// inserted ranges.
    /// TODO: Make this operation atomic.
    /// TODO: Consider using a different algorithmic that first sorts ranges
    /// and then merges used and ranges by traversing them in parallel, for
    /// efficiency.
    fn insert_ranges(&mut self, ranges: &[(Sectors, Sectors)]) -> StratisResult<()> {
        for &(off, len) in ranges {
            self.check_for_overflow(off, len)?;

            let prev = self.used.range(..off).rev().next().map(|(k, v)| (*k, *v));

            let mut contig_prev = None;
            if let Some((prev_off, prev_len)) = prev {
                if prev_off + prev_len > off {
                    let err_msg = format!(
                        "range starting at {} overlaps previous range ({}, {})",
                        off, prev_off, prev_len
                    );
                    return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
                }
                if prev_off + prev_len == off {
                    contig_prev = Some((prev_off, prev_len))
                }
            }

            let next = self.used.range(off..).next().map(|(k, v)| (*k, *v));

            let mut contig_next = None;
            if let Some((next_off, next_len)) = next {
                if off + len > next_off {
                    let err_msg = format!(
                        "range ({}, {}) overlaps subsequent range starting at {}",
                        off, len, next_off
                    );
                    return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
                }
                if off + len == next_off {
                    contig_next = Some((next_off, next_len))
                }
            }

            match (contig_prev, contig_next) {
                (None, None) => {
                    self.used.insert(off, len);
                }
                (None, Some((next_off, next_len))) => {
                    // Contig with next, make new entry
                    self.used.insert(off, len + next_len);
                    self.used
                        .remove(&next_off)
                        .expect("matched Some((next_off, ...");
                }
                (Some((prev_off, prev_len)), None) => {
                    // Contig with prev, just extend prev
                    *self
                        .used
                        .get_mut(&prev_off)
                        .expect("matched Some((prev_off, ...") = prev_len + len;
                }
                (Some((prev_off, prev_len)), Some((next_off, next_len))) => {
                    // Contig with both, remove next and extend prev
                    self.used.remove(&next_off);
                    *self
                        .used
                        .get_mut(&prev_off)
                        .expect("matched Some((prev_off, ...") = prev_len + len + next_len;
                }
            }
        }
        Ok(())
    }

    #[allow(dead_code)]
    /// Mark ranges previously marked as used as now unused.
    fn remove_ranges(&mut self, to_free: &[(Sectors, Sectors)]) -> StratisResult<()> {
        for &(off, len) in to_free {
            self.check_for_overflow(off, len)?;

            let maybe_prev = self
                .used
                .range((Unbounded, Included(off)))
                .rev()
                .next()
                .map(|(k, v)| (*k, *v));

            let (prev_off, prev_len) = match maybe_prev {
                Some(range) => range,
                None => panic!("Existing matching allocated range not found"),
            };

            assert!(
                prev_off + prev_len >= off + len,
                "must not extend past existing range"
            );

            // switch based on if the to-remove range starts or ends
            // at the same point as the existing range
            match (prev_off == off, prev_off + prev_len == off + len) {
                (true, true) => {
                    // Exactly matches existing range
                    self.used.remove(&prev_off);
                }
                (true, false) => {
                    // A tail segment remains
                    self.used.remove(&prev_off).expect("must exist");
                    self.used.insert(prev_off + len, prev_len - len);
                }
                (false, true) => {
                    // Head segment remains
                    *self.used.get_mut(&prev_off).expect("must exist") = prev_len - len;
                }
                (false, false) => {
                    // Head and tail segments both remain
                    let prev_end = prev_off + prev_len;
                    let tail_off = off + len;
                    let tail_len = prev_end - tail_off;
                    let head_len = prev_len - len - tail_len;
                    *self.used.get_mut(&prev_off).expect("must exist") = head_len;
                    self.used.insert(tail_off, tail_len);
                }
            }
        }
        Ok(())
    }

    /// Available sectors
    pub fn available(&self) -> Sectors {
        self.limit - self.used()
    }

    /// Allocated sectors
    pub fn used(&self) -> Sectors {
        self.used.values().cloned().sum()
    }

    /// Get a list of (offset, length) segments that are in use
    fn used_ranges(&self) -> Vec<(Sectors, Sectors)> {
        self.used.iter().map(|(k, v)| (*k, *v)).collect()
    }

    /// Get a list of (offset, length) segments that are not in use
    fn avail_ranges(&self) -> Vec<(Sectors, Sectors)> {
        let mut free = Vec::new();

        // Insert an entry to mark the end so the fold works correctly
        let mut used = self.used_ranges();
        used.push((self.limit, Sectors(0)));

        used.into_iter().fold(Sectors(0), |prev_end, (start, len)| {
            if prev_end < start {
                free.push((prev_end, start - prev_end))
            }
            start + len
        });

        free
    }

    /// Attempt to allocate. Returns number of sectors allocated (may
    /// be less than request, including zero) and a Vec<(offset,
    /// length)> of sectors successfully allocated.
    /// If all available sectors are desired, use available() method to
    /// discover that amount.
    pub fn request(&mut self, amount: Sectors) -> (Sectors, Vec<(Sectors, Sectors)>) {
        let mut segs = Vec::new();
        let mut needed = amount;

        for (start, len) in self.avail_ranges() {
            if needed == Sectors(0) {
                break;
            }

            let to_use = min(needed, len);

            let used_range = (start, to_use);
            segs.push(used_range);
            self.insert_ranges(&[used_range])
                .expect("available ranges must be insertable");

            needed -= to_use;
        }

        (amount - needed, segs)
    }
}

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

    #[test]
    /// Test proper operation of RangeAllocator.
    /// 1. Instantiate a RangeAllocator.
    /// 2. Verify that no sectors are used (all are available).
    /// 3. Insert range (10, 100) into the allocator.
    /// 4. Verify that 100 sectors are taken and 28 remain.
    /// 5. Request 50 sectors from the allocator.
    /// 6. Verify that the maximum available, 28, were returned in two ranges.
    /// 7. Remove two adjacent ranges of total length 60 sectors.
    /// 8. Verify that number of available sectors is 60, used is 68.
    /// 9. Request all available, then verify that nothing is left.
    fn test_allocator_allocations() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        assert_eq!(allocator.used(), Sectors(0));
        assert_eq!(allocator.available(), Sectors(128));

        allocator
            .insert_ranges(&[(Sectors(10), Sectors(100))])
            .unwrap();

        assert_eq!(allocator.used(), Sectors(100));
        assert_eq!(allocator.available(), Sectors(28));

        let request = allocator.request(Sectors(50));
        assert_eq!(request.0, Sectors(28));
        assert_eq!(allocator.used(), Sectors(128));
        assert_eq!(allocator.available(), Sectors(0));
        assert_eq!(request.1.len(), 2);

        let good_remove_ranges = [(Sectors(21), Sectors(20)), (Sectors(41), Sectors(40))];
        allocator.remove_ranges(&good_remove_ranges).unwrap();
        assert_eq!(allocator.used(), Sectors(68));
        assert_eq!(allocator.available(), Sectors(60));

        let available = allocator.available();
        allocator.request(available);
        assert_eq!(allocator.available(), Sectors(0));
    }

    #[test]
    // Verify some proper functioning when allocator initialized with ranges.
    fn test_allocator_initialized_with_range() {
        let ranges = [
            (Sectors(20), Sectors(10)),
            (Sectors(10), Sectors(10)),
            (Sectors(30), Sectors(10)),
        ];
        let allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &ranges).unwrap();
        let used = allocator.used_ranges();
        assert_eq!(used.len(), 1);
        assert_eq!(used[0], (Sectors(10), Sectors(30)));
    }

    #[test]
    /// Verify insert_ranges properly coalesces adjacent allocations.
    fn test_allocator_insert_ranges_contig() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        allocator
            .insert_ranges(&[(Sectors(20), Sectors(10))])
            .unwrap();
        allocator
            .insert_ranges(&[(Sectors(10), Sectors(10))])
            .unwrap();
        allocator
            .insert_ranges(&[(Sectors(30), Sectors(10))])
            .unwrap();

        let used = allocator.used_ranges();
        assert_eq!(used.len(), 1);
        assert_eq!(used[0], (Sectors(10), Sectors(30)));
    }

    #[test]
    /// Verify remove_ranges properly handles different cases.
    /// 1. Removing a range from the start of an existing range
    /// 2. Removing a range from the end of an existing range
    /// 3. Removing a range from the middle of an existing range
    /// 4. Removing an entire range
    fn test_allocator_remove_ranges_contig() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        allocator
            .insert_ranges(&[(Sectors(20), Sectors(20))])
            .unwrap();

        allocator
            .remove_ranges(&[(Sectors(20), Sectors(3))])
            .unwrap();
        let used = allocator.used_ranges();
        assert_eq!(used.len(), 1);
        assert_eq!(used[0], (Sectors(23), Sectors(17)));

        allocator
            .remove_ranges(&[(Sectors(36), Sectors(4))])
            .unwrap();
        let used = allocator.used_ranges();
        assert_eq!(used.len(), 1);
        assert_eq!(used[0], (Sectors(23), Sectors(13)));

        allocator
            .remove_ranges(&[(Sectors(24), Sectors(2))])
            .unwrap();
        let used = allocator.used_ranges();
        assert_eq!(used.len(), 2);
        assert_eq!(used[0], (Sectors(23), Sectors(1)));
        assert_eq!(used[1], (Sectors(26), Sectors(10)));

        allocator
            .remove_ranges(&[(Sectors(26), Sectors(10))])
            .unwrap();
        let used = allocator.used_ranges();
        assert_eq!(used.len(), 1);
        assert_eq!(used[0], (Sectors(23), Sectors(1)));
    }

    #[test]
    /// Verify that the largest possible limit may be used for the
    /// allocator.
    fn test_max_allocator_range() {
        use std::u64::MAX;

        RangeAllocator::new(BlockdevSize::new(Sectors(MAX)), &[]).unwrap();
    }

    #[test]
    fn test_allocator_insert_prev_overlap() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        let bad_insert_ranges = [(Sectors(21), Sectors(20)), (Sectors(40), Sectors(40))];
        assert_matches!(allocator.insert_ranges(&bad_insert_ranges), Err(_))
    }

    #[test]
    fn test_allocator_insert_next_overlap() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        let bad_insert_ranges = [(Sectors(40), Sectors(1)), (Sectors(39), Sectors(2))];
        assert_matches!(allocator.insert_ranges(&bad_insert_ranges), Err(_))
    }

    #[test]
    #[should_panic]
    /// Verify that remove_ranges() panics if ranges to be removed share
    /// elements.
    fn test_allocator_failures_alloc_overlap() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        let _request = allocator.request(Sectors(128));

        let bad_remove_ranges = [(Sectors(21), Sectors(20)), (Sectors(40), Sectors(40))];
        allocator.remove_ranges(&bad_remove_ranges).unwrap();
    }

    #[test]
    /// Verify that insert_ranges() errors when all sectors have already been
    /// allocated.
    fn test_allocator_failures_range_overwrite() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        let request = allocator.request(Sectors(128));
        assert_eq!(request.0, Sectors(128));
        assert_eq!(request.1, &[(Sectors(0), Sectors(128))]);

        assert_matches!(allocator.insert_ranges(&[(Sectors(1), Sectors(1))]), Err(_));
    }

    #[test]
    #[should_panic]
    /// Verify that remove_ranges() panics when an element at the
    /// beginning of the specified range is not in use.
    fn test_allocator_failures_removing_unused_beginning() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        allocator
            .insert_ranges(&[(Sectors(20), Sectors(20))])
            .unwrap();
        allocator
            .remove_ranges(&[(Sectors(19), Sectors(2))])
            .unwrap();
    }

    #[test]
    #[should_panic]
    /// Verify that remove_ranges() panics when an element at the
    /// end of the specified range is not in use.
    fn test_allocator_failures_removing_unused_end() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        allocator
            .insert_ranges(&[(Sectors(20), Sectors(20))])
            .unwrap();
        allocator
            .remove_ranges(&[(Sectors(39), Sectors(2))])
            .unwrap();
    }

    #[test]
    #[should_panic]
    /// Verify that remove_ranges() panics when the entire specified
    /// range is not in use.
    fn test_allocator_failures_removing_unused() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        allocator
            .remove_ranges(&[(Sectors(39), Sectors(2))])
            .unwrap();
    }

    #[test]
    /// Verify that insert_ranges() errors when an element outside the range
    /// limit is requested.
    fn test_allocator_failures_overflow_limit() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        // overflow limit range
        assert_matches!(
            allocator.insert_ranges(&[(Sectors(1), Sectors(128))]),
            Err(_)
        );
    }

    #[test]
    /// Verify that insert_ranges() errors when an element in a requested range
    /// exceeds u64::MAX.
    fn test_allocator_failures_overflow_max() {
        use std::u64::MAX;

        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(MAX)), &[]).unwrap();

        // overflow max u64
        assert_matches!(
            allocator.insert_ranges(&[(Sectors(MAX), Sectors(1))]),
            Err(_)
        );
    }
}
<commit_msg>Remove remove_ranges method<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{cmp::min, collections::BTreeMap};

use devicemapper::Sectors;

use crate::{
    engine::strat_engine::backstore::metadata::BlockdevSize,
    stratis::{ErrorEnum, StratisError, StratisResult},
};

#[derive(Debug)]
pub struct RangeAllocator {
    limit: Sectors,
    used: BTreeMap<Sectors, Sectors>,
}

impl RangeAllocator {
    /// Create a new RangeAllocator with the specified (offset, length)
    /// ranges marked as used.
    #[allow(clippy::new_ret_no_self)]
    pub fn new(
        limit: BlockdevSize,
        initial_used: &[(Sectors, Sectors)],
    ) -> StratisResult<RangeAllocator> {
        let mut allocator = RangeAllocator {
            limit: limit.sectors(),
            used: BTreeMap::new(),
        };
        allocator.insert_ranges(initial_used)?;
        Ok(allocator)
    }

    /// The maximum allocation from this manager
    pub fn size(&self) -> BlockdevSize {
        BlockdevSize::new(self.limit)
    }

    fn check_for_overflow(&self, off: Sectors, len: Sectors) -> StratisResult<()> {
        if let Some(sum) = off.checked_add(len) {
            if sum > self.limit {
                let err_msg = format!(
                    "elements in range ({}, {}) exceed limit {}",
                    off, len, self.limit
                );
                return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
            }
        } else {
            let err_msg = format!(
                "elements in range ({}, {}) inexpressible in this format",
                off, len
            );
            return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
        }
        Ok(())
    }

    /// Mark ranges previously marked as unused as now used.
    /// Return an error if ranges overlap with each other or with previously
    /// inserted ranges.
    /// TODO: Make this operation atomic.
    /// TODO: Consider using a different algorithmic that first sorts ranges
    /// and then merges used and ranges by traversing them in parallel, for
    /// efficiency.
    fn insert_ranges(&mut self, ranges: &[(Sectors, Sectors)]) -> StratisResult<()> {
        for &(off, len) in ranges {
            self.check_for_overflow(off, len)?;

            let prev = self.used.range(..off).rev().next().map(|(k, v)| (*k, *v));

            let mut contig_prev = None;
            if let Some((prev_off, prev_len)) = prev {
                if prev_off + prev_len > off {
                    let err_msg = format!(
                        "range starting at {} overlaps previous range ({}, {})",
                        off, prev_off, prev_len
                    );
                    return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
                }
                if prev_off + prev_len == off {
                    contig_prev = Some((prev_off, prev_len))
                }
            }

            let next = self.used.range(off..).next().map(|(k, v)| (*k, *v));

            let mut contig_next = None;
            if let Some((next_off, next_len)) = next {
                if off + len > next_off {
                    let err_msg = format!(
                        "range ({}, {}) overlaps subsequent range starting at {}",
                        off, len, next_off
                    );
                    return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
                }
                if off + len == next_off {
                    contig_next = Some((next_off, next_len))
                }
            }

            match (contig_prev, contig_next) {
                (None, None) => {
                    self.used.insert(off, len);
                }
                (None, Some((next_off, next_len))) => {
                    // Contig with next, make new entry
                    self.used.insert(off, len + next_len);
                    self.used
                        .remove(&next_off)
                        .expect("matched Some((next_off, ...");
                }
                (Some((prev_off, prev_len)), None) => {
                    // Contig with prev, just extend prev
                    *self
                        .used
                        .get_mut(&prev_off)
                        .expect("matched Some((prev_off, ...") = prev_len + len;
                }
                (Some((prev_off, prev_len)), Some((next_off, next_len))) => {
                    // Contig with both, remove next and extend prev
                    self.used.remove(&next_off);
                    *self
                        .used
                        .get_mut(&prev_off)
                        .expect("matched Some((prev_off, ...") = prev_len + len + next_len;
                }
            }
        }
        Ok(())
    }

    /// Available sectors
    pub fn available(&self) -> Sectors {
        self.limit - self.used()
    }

    /// Allocated sectors
    pub fn used(&self) -> Sectors {
        self.used.values().cloned().sum()
    }

    /// Get a list of (offset, length) segments that are in use
    fn used_ranges(&self) -> Vec<(Sectors, Sectors)> {
        self.used.iter().map(|(k, v)| (*k, *v)).collect()
    }

    /// Get a list of (offset, length) segments that are not in use
    fn avail_ranges(&self) -> Vec<(Sectors, Sectors)> {
        let mut free = Vec::new();

        // Insert an entry to mark the end so the fold works correctly
        let mut used = self.used_ranges();
        used.push((self.limit, Sectors(0)));

        used.into_iter().fold(Sectors(0), |prev_end, (start, len)| {
            if prev_end < start {
                free.push((prev_end, start - prev_end))
            }
            start + len
        });

        free
    }

    /// Attempt to allocate. Returns number of sectors allocated (may
    /// be less than request, including zero) and a Vec<(offset,
    /// length)> of sectors successfully allocated.
    /// If all available sectors are desired, use available() method to
    /// discover that amount.
    pub fn request(&mut self, amount: Sectors) -> (Sectors, Vec<(Sectors, Sectors)>) {
        let mut segs = Vec::new();
        let mut needed = amount;

        for (start, len) in self.avail_ranges() {
            if needed == Sectors(0) {
                break;
            }

            let to_use = min(needed, len);

            let used_range = (start, to_use);
            segs.push(used_range);
            self.insert_ranges(&[used_range])
                .expect("available ranges must be insertable");

            needed -= to_use;
        }

        (amount - needed, segs)
    }
}

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

    #[test]
    /// Test proper operation of RangeAllocator.
    /// 1. Instantiate a RangeAllocator.
    /// 2. Verify that no sectors are used (all are available).
    /// 3. Insert range (10, 100) into the allocator.
    /// 4. Verify that 100 sectors are taken and 28 remain.
    /// 5. Request 50 sectors from the allocator.
    /// 6. Verify that the maximum available, 28, were returned in two ranges.
    /// 7. Remove two adjacent ranges of total length 60 sectors.
    /// 8. Verify that number of available sectors is 60, used is 68.
    /// 9. Request all available, then verify that nothing is left.
    fn test_allocator_allocations() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        assert_eq!(allocator.used(), Sectors(0));
        assert_eq!(allocator.available(), Sectors(128));

        allocator
            .insert_ranges(&[(Sectors(10), Sectors(100))])
            .unwrap();

        assert_eq!(allocator.used(), Sectors(100));
        assert_eq!(allocator.available(), Sectors(28));

        let request = allocator.request(Sectors(50));
        assert_eq!(request.0, Sectors(28));
        assert_eq!(allocator.used(), Sectors(128));
        assert_eq!(allocator.available(), Sectors(0));
        assert_eq!(request.1.len(), 2);

        let available = allocator.available();
        allocator.request(available);
        assert_eq!(allocator.available(), Sectors(0));
    }

    #[test]
    // Verify some proper functioning when allocator initialized with ranges.
    fn test_allocator_initialized_with_range() {
        let ranges = [
            (Sectors(20), Sectors(10)),
            (Sectors(10), Sectors(10)),
            (Sectors(30), Sectors(10)),
        ];
        let allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &ranges).unwrap();
        let used = allocator.used_ranges();
        assert_eq!(used.len(), 1);
        assert_eq!(used[0], (Sectors(10), Sectors(30)));
    }

    #[test]
    /// Verify insert_ranges properly coalesces adjacent allocations.
    fn test_allocator_insert_ranges_contig() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        allocator
            .insert_ranges(&[(Sectors(20), Sectors(10))])
            .unwrap();
        allocator
            .insert_ranges(&[(Sectors(10), Sectors(10))])
            .unwrap();
        allocator
            .insert_ranges(&[(Sectors(30), Sectors(10))])
            .unwrap();

        let used = allocator.used_ranges();
        assert_eq!(used.len(), 1);
        assert_eq!(used[0], (Sectors(10), Sectors(30)));
    }

    #[test]
    /// Verify that the largest possible limit may be used for the
    /// allocator.
    fn test_max_allocator_range() {
        use std::u64::MAX;

        RangeAllocator::new(BlockdevSize::new(Sectors(MAX)), &[]).unwrap();
    }

    #[test]
    fn test_allocator_insert_prev_overlap() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        let bad_insert_ranges = [(Sectors(21), Sectors(20)), (Sectors(40), Sectors(40))];
        assert_matches!(allocator.insert_ranges(&bad_insert_ranges), Err(_))
    }

    #[test]
    fn test_allocator_insert_next_overlap() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        let bad_insert_ranges = [(Sectors(40), Sectors(1)), (Sectors(39), Sectors(2))];
        assert_matches!(allocator.insert_ranges(&bad_insert_ranges), Err(_))
    }

    #[test]
    /// Verify that insert_ranges() errors when all sectors have already been
    /// allocated.
    fn test_allocator_failures_range_overwrite() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        let request = allocator.request(Sectors(128));
        assert_eq!(request.0, Sectors(128));
        assert_eq!(request.1, &[(Sectors(0), Sectors(128))]);

        assert_matches!(allocator.insert_ranges(&[(Sectors(1), Sectors(1))]), Err(_));
    }

    #[test]
    /// Verify that insert_ranges() errors when an element outside the range
    /// limit is requested.
    fn test_allocator_failures_overflow_limit() {
        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(128)), &[]).unwrap();

        // overflow limit range
        assert_matches!(
            allocator.insert_ranges(&[(Sectors(1), Sectors(128))]),
            Err(_)
        );
    }

    #[test]
    /// Verify that insert_ranges() errors when an element in a requested range
    /// exceeds u64::MAX.
    fn test_allocator_failures_overflow_max() {
        use std::u64::MAX;

        let mut allocator = RangeAllocator::new(BlockdevSize::new(Sectors(MAX)), &[]).unwrap();

        // overflow max u64
        assert_matches!(
            allocator.insert_ranges(&[(Sectors(MAX), Sectors(1))]),
            Err(_)
        );
    }
}
<|endoftext|>


<commit_before><commit_msg>remove some clones<commit_after><|endoftext|>


<commit_before>use common::util::*;

extern crate uu_stat;
pub use self::uu_stat::*;

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

    #[test]
    fn test_access() {
        assert_eq!("drwxr-xr-x", pretty_access(S_IFDIR | 0o755));
        assert_eq!("-rw-r--r--", pretty_access(S_IFREG | 0o644));
        assert_eq!("srw-r-----", pretty_access(S_IFSOCK | 0o640));
        assert_eq!("lrw-r-xr-x", pretty_access(S_IFLNK | 0o655));
        assert_eq!("?rw-r-xr-x", pretty_access(0o655));

        assert_eq!("brwSr-xr-x",
                   pretty_access(S_IFBLK | S_ISUID as mode_t | 0o655));
        assert_eq!("brwsr-xr-x",
                   pretty_access(S_IFBLK | S_ISUID as mode_t | 0o755));

        assert_eq!("prw---sr--",
                   pretty_access(S_IFIFO | S_ISGID as mode_t | 0o614));
        assert_eq!("prw---Sr--",
                   pretty_access(S_IFIFO | S_ISGID as mode_t | 0o604));

        assert_eq!("c---r-xr-t",
                   pretty_access(S_IFCHR | S_ISVTX as mode_t | 0o055));
        assert_eq!("c---r-xr-T",
                   pretty_access(S_IFCHR | S_ISVTX as mode_t | 0o054));
    }

    #[test]
    fn test_file_type() {
        assert_eq!("block special file", pretty_filetype(S_IFBLK, 0));
        assert_eq!("character special file", pretty_filetype(S_IFCHR, 0));
        assert_eq!("regular file", pretty_filetype(S_IFREG, 1));
        assert_eq!("regular empty file", pretty_filetype(S_IFREG, 0));
        assert_eq!("weird file", pretty_filetype(0, 0));
    }

    #[test]
    fn test_fs_type() {
        assert_eq!("ext2/ext3", pretty_fstype(0xEF53));
        assert_eq!("tmpfs", pretty_fstype(0x01021994));
        assert_eq!("nfs", pretty_fstype(0x6969));
        assert_eq!("btrfs", pretty_fstype(0x9123683e));
        assert_eq!("xfs", pretty_fstype(0x58465342));
        assert_eq!("zfs", pretty_fstype(0x2FC12FC1));
        assert_eq!("ntfs", pretty_fstype(0x5346544e));
        assert_eq!("fat", pretty_fstype(0x4006));
        assert_eq!("UNKNOWN (0x1234)", pretty_fstype(0x1234));
    }
}

#[test]
fn test_scanutil() {
    assert_eq!(Some((-5, 2)), "-5zxc".scan_num::<i32>());
    assert_eq!(Some((51, 2)), "51zxc".scan_num::<u32>());
    assert_eq!(Some((192, 4)), "+192zxc".scan_num::<i32>());
    assert_eq!(None, "z192zxc".scan_num::<i32>());

    assert_eq!(Some(('a', 3)), "141zxc".scan_char(8));
    assert_eq!(Some(('\n', 2)), "12qzxc".scan_char(8));
    assert_eq!(Some(('\r', 1)), "dqzxc".scan_char(16));
    assert_eq!(None, "z2qzxc".scan_char(8));
}

#[test]
fn test_groupnum() {
    assert_eq!("12,379,821,234", group_num("12379821234"));
    assert_eq!("21,234", group_num("21234"));
    assert_eq!("821,234", group_num("821234"));
    assert_eq!("1,821,234", group_num("1821234"));
    assert_eq!("1,234", group_num("1234"));
    assert_eq!("234", group_num("234"));
    assert_eq!("24", group_num("24"));
    assert_eq!("4", group_num("4"));
    assert_eq!("", group_num(""));
}

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

    #[test]
    fn normal_format() {
        let s = "%'010.2ac%-#5.w\n";
        let expected = vec![Token::Directive {
                                flag: F_GROUP | F_ZERO,
                                width: 10,
                                precision: 2,
                                format: 'a',
                            },
                            Token::Char('c'),
                            Token::Directive {
                                flag: F_LEFT | F_ALTER,
                                width: 5,
                                precision: 0,
                                format: 'w',
                            },
                            Token::Char('\n')];
        assert_eq!(&expected, &Stater::generate_tokens(s, false).unwrap());
    }

    #[test]
    fn printf_format() {
        let s = "%-# 15a\\r\\\"\\\\\\a\\b\\e\\f\\v%+020.-23w\\x12\\167\\132\\112\\n";
        let expected = vec![Token::Directive {
                                flag: F_LEFT | F_ALTER | F_SPACE,
                                width: 15,
                                precision: -1,
                                format: 'a',
                            },
                            Token::Char('\r'),
                            Token::Char('"'),
                            Token::Char('\\'),
                            Token::Char('\x07'),
                            Token::Char('\x08'),
                            Token::Char('\x1B'),
                            Token::Char('\x0C'),
                            Token::Char('\x0B'),
                            Token::Directive {
                                flag: F_SIGN | F_ZERO,
                                width: 20,
                                precision: -1,
                                format: 'w',
                            },
                            Token::Char('\x12'),
                            Token::Char('w'),
                            Token::Char('Z'),
                            Token::Char('J'),
                            Token::Char('\n')];
        assert_eq!(&expected, &Stater::generate_tokens(s, true).unwrap());
    }
}

#[test]
fn test_invalid_option() {
    new_ucmd!()
        .arg("-w").arg("-q").arg("/").fails();
}

#[cfg(target_os = "linux")]
const NORMAL_FMTSTR: &'static str = "%a %A %b %B %d %D %f %F %g %G %h %i %m %n %o %s %u %U %w %W %x %X %y %Y %z %Z";
#[cfg(target_os = "linux")]
const DEV_FMTSTR: &'static str = "%a %A %b %B %d %D %f %F %g %G %h %i %m %n %o %s (%t/%T) %u %U %w %W %x %X %y %Y %z %Z";
#[cfg(target_os = "linux")]
const FS_FMTSTR: &'static str = "%a %b %c %d %f %i %l %n %s %S %t %T";

#[test]
#[cfg(target_os = "linux")]
fn test_terse_fs_format() {
    let args = ["-f", "-t", "/proc"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_fs_format() {
    let args = ["-f", "-c", FS_FMTSTR, "/dev/shm"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_terse_normal_format() {
    let args = ["-t", "/"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_normal_format() {
    let args = ["-c", NORMAL_FMTSTR, "/boot"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_follow_symlink() {
    let args = ["-L", "-c", DEV_FMTSTR, "/dev/cdrom"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_symlink() {
    let args = ["-c", DEV_FMTSTR, "/dev/cdrom"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_char() {
    let args = ["-c", DEV_FMTSTR, "/dev/pts/ptmx"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_multi_files() {
    let args = ["-c", NORMAL_FMTSTR, "/dev", "/usr/lib", "/etc/fstab", "/var"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_printf() {
    let args = ["--printf=123%-# 15q\\r\\\"\\\\\\a\\b\\e\\f\\v%+020.23m\\x12\\167\\132\\112\\n", "/"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[cfg(target_os = "linux")]
fn expected_result(args: &[&str]) -> String {
    TestScenario::new(util_name!()).cmd_keepenv(util_name!()).env("LANGUAGE", "C").args(args).run().stdout
}
<commit_msg>tests/fix ~ (stat) refactor '%w' and '%W' tests due to `stat` limitations<commit_after>extern crate regex;

use common::util::*;

extern crate uu_stat;
pub use self::uu_stat::*;

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

    #[test]
    fn test_access() {
        assert_eq!("drwxr-xr-x", pretty_access(S_IFDIR | 0o755));
        assert_eq!("-rw-r--r--", pretty_access(S_IFREG | 0o644));
        assert_eq!("srw-r-----", pretty_access(S_IFSOCK | 0o640));
        assert_eq!("lrw-r-xr-x", pretty_access(S_IFLNK | 0o655));
        assert_eq!("?rw-r-xr-x", pretty_access(0o655));

        assert_eq!("brwSr-xr-x",
                   pretty_access(S_IFBLK | S_ISUID as mode_t | 0o655));
        assert_eq!("brwsr-xr-x",
                   pretty_access(S_IFBLK | S_ISUID as mode_t | 0o755));

        assert_eq!("prw---sr--",
                   pretty_access(S_IFIFO | S_ISGID as mode_t | 0o614));
        assert_eq!("prw---Sr--",
                   pretty_access(S_IFIFO | S_ISGID as mode_t | 0o604));

        assert_eq!("c---r-xr-t",
                   pretty_access(S_IFCHR | S_ISVTX as mode_t | 0o055));
        assert_eq!("c---r-xr-T",
                   pretty_access(S_IFCHR | S_ISVTX as mode_t | 0o054));
    }

    #[test]
    fn test_file_type() {
        assert_eq!("block special file", pretty_filetype(S_IFBLK, 0));
        assert_eq!("character special file", pretty_filetype(S_IFCHR, 0));
        assert_eq!("regular file", pretty_filetype(S_IFREG, 1));
        assert_eq!("regular empty file", pretty_filetype(S_IFREG, 0));
        assert_eq!("weird file", pretty_filetype(0, 0));
    }

    #[test]
    fn test_fs_type() {
        assert_eq!("ext2/ext3", pretty_fstype(0xEF53));
        assert_eq!("tmpfs", pretty_fstype(0x01021994));
        assert_eq!("nfs", pretty_fstype(0x6969));
        assert_eq!("btrfs", pretty_fstype(0x9123683e));
        assert_eq!("xfs", pretty_fstype(0x58465342));
        assert_eq!("zfs", pretty_fstype(0x2FC12FC1));
        assert_eq!("ntfs", pretty_fstype(0x5346544e));
        assert_eq!("fat", pretty_fstype(0x4006));
        assert_eq!("UNKNOWN (0x1234)", pretty_fstype(0x1234));
    }
}

#[test]
fn test_scanutil() {
    assert_eq!(Some((-5, 2)), "-5zxc".scan_num::<i32>());
    assert_eq!(Some((51, 2)), "51zxc".scan_num::<u32>());
    assert_eq!(Some((192, 4)), "+192zxc".scan_num::<i32>());
    assert_eq!(None, "z192zxc".scan_num::<i32>());

    assert_eq!(Some(('a', 3)), "141zxc".scan_char(8));
    assert_eq!(Some(('\n', 2)), "12qzxc".scan_char(8));
    assert_eq!(Some(('\r', 1)), "dqzxc".scan_char(16));
    assert_eq!(None, "z2qzxc".scan_char(8));
}

#[test]
fn test_groupnum() {
    assert_eq!("12,379,821,234", group_num("12379821234"));
    assert_eq!("21,234", group_num("21234"));
    assert_eq!("821,234", group_num("821234"));
    assert_eq!("1,821,234", group_num("1821234"));
    assert_eq!("1,234", group_num("1234"));
    assert_eq!("234", group_num("234"));
    assert_eq!("24", group_num("24"));
    assert_eq!("4", group_num("4"));
    assert_eq!("", group_num(""));
}

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

    #[test]
    fn normal_format() {
        let s = "%'010.2ac%-#5.w\n";
        let expected = vec![Token::Directive {
                                flag: F_GROUP | F_ZERO,
                                width: 10,
                                precision: 2,
                                format: 'a',
                            },
                            Token::Char('c'),
                            Token::Directive {
                                flag: F_LEFT | F_ALTER,
                                width: 5,
                                precision: 0,
                                format: 'w',
                            },
                            Token::Char('\n')];
        assert_eq!(&expected, &Stater::generate_tokens(s, false).unwrap());
    }

    #[test]
    fn printf_format() {
        let s = "%-# 15a\\r\\\"\\\\\\a\\b\\e\\f\\v%+020.-23w\\x12\\167\\132\\112\\n";
        let expected = vec![Token::Directive {
                                flag: F_LEFT | F_ALTER | F_SPACE,
                                width: 15,
                                precision: -1,
                                format: 'a',
                            },
                            Token::Char('\r'),
                            Token::Char('"'),
                            Token::Char('\\'),
                            Token::Char('\x07'),
                            Token::Char('\x08'),
                            Token::Char('\x1B'),
                            Token::Char('\x0C'),
                            Token::Char('\x0B'),
                            Token::Directive {
                                flag: F_SIGN | F_ZERO,
                                width: 20,
                                precision: -1,
                                format: 'w',
                            },
                            Token::Char('\x12'),
                            Token::Char('w'),
                            Token::Char('Z'),
                            Token::Char('J'),
                            Token::Char('\n')];
        assert_eq!(&expected, &Stater::generate_tokens(s, true).unwrap());
    }
}

#[test]
fn test_invalid_option() {
    new_ucmd!()
        .arg("-w").arg("-q").arg("/").fails();
}

#[cfg(target_os = "linux")]
const NORMAL_FMTSTR: &'static str = "%a %A %b %B %d %D %f %F %g %G %h %i %m %n %o %s %u %U %x %X %y %Y %z %Z"; // avoid "%w %W" (birth/creation) due to `stat` limitations and linux kernel & rust version capability variations
#[cfg(target_os = "linux")]
const DEV_FMTSTR: &'static str = "%a %A %b %B %d %D %f %F %g %G %h %i %m %n %o %s (%t/%T) %u %U %w %W %x %X %y %Y %z %Z";
#[cfg(target_os = "linux")]
const FS_FMTSTR: &'static str = "%a %b %c %d %f %i %l %n %s %S %t %T";

#[test]
#[cfg(target_os = "linux")]
fn test_terse_fs_format() {
    let args = ["-f", "-t", "/proc"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_fs_format() {
    let args = ["-f", "-c", FS_FMTSTR, "/dev/shm"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_terse_normal_format() {
    // note: contains birth/creation date which increases test fragility
    // * results may vary due to built-in `stat` limitations as well as linux kernel and rust version capability variations
    let args = ["-t", "/"];
    let actual = new_ucmd!().args(&args).run().stdout;
    let expect = expected_result(&args);
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    let v_actual: Vec<&str> = actual.split(' ').collect();
    let v_expect: Vec<&str> = expect.split(' ').collect();
    // * allow for inequality if `stat` (aka, expect) returns "0" (unknown value)
    assert!(v_actual.iter().zip(v_expect.iter()).all(|(a,e)| a == e || *e == "0"));
}

#[test]
#[cfg(target_os = "linux")]
fn test_format_created_time() {
    let args = ["-c", "%w", "/boot"];
    let actual = new_ucmd!().args(&args).run().stdout;
    let expect = expected_result(&args);
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    // note: using a regex instead of `split_whitespace()` in order to detect whitespace differences
    let re = regex::Regex::new(r"\s").unwrap();
    let v_actual: Vec<&str> = re.split(&actual).collect();
    let v_expect: Vec<&str> = re.split(&expect).collect();
    // * allow for inequality if `stat` (aka, expect) returns "-" (unknown value)
    assert!(v_actual.iter().zip(v_expect.iter()).all(|(a,e)| a == e || *e == "-"));
}

#[test]
#[cfg(target_os = "linux")]
fn test_format_created_seconds() {
    let args = ["-c", "%W", "/boot"];
    let actual = new_ucmd!().args(&args).run().stdout;
    let expect = expected_result(&args);
    println!("actual: {:?}", actual);
    println!("expect: {:?}", expect);
    // note: using a regex instead of `split_whitespace()` in order to detect whitespace differences
    let re = regex::Regex::new(r"\s").unwrap();
    let v_actual: Vec<&str> = re.split(&actual).collect();
    let v_expect: Vec<&str> = re.split(&expect).collect();
    // * allow for inequality if `stat` (aka, expect) returns "0" (unknown value)
    assert!(v_actual.iter().zip(v_expect.iter()).all(|(a,e)| a == e || *e == "0"));
}

#[test]
#[cfg(target_os = "linux")]
fn test_normal_format() {
    let args = ["-c", NORMAL_FMTSTR, "/boot"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_follow_symlink() {
    let args = ["-L", "-c", DEV_FMTSTR, "/dev/cdrom"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_symlink() {
    let args = ["-c", DEV_FMTSTR, "/dev/cdrom"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_char() {
    let args = ["-c", DEV_FMTSTR, "/dev/pts/ptmx"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_multi_files() {
    let args = ["-c", NORMAL_FMTSTR, "/dev", "/usr/lib", "/etc/fstab", "/var"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[test]
#[cfg(target_os = "linux")]
fn test_printf() {
    let args = ["--printf=123%-# 15q\\r\\\"\\\\\\a\\b\\e\\f\\v%+020.23m\\x12\\167\\132\\112\\n", "/"];
    new_ucmd!().args(&args)
        .run()
        .stdout_is(expected_result(&args));
}

#[cfg(target_os = "linux")]
fn expected_result(args: &[&str]) -> String {
    TestScenario::new(util_name!()).cmd_keepenv(util_name!()).env("LANGUAGE", "C").args(args).run().stdout
}
<|endoftext|>


<commit_before>use {primitives, Parcel, Error};

use std::io::prelude::*;
use std;

/// The integer type that we will use to send length prefixes.
pub type SizeType = u32;

/// An dynamic array type with a custom size prefix type.
#[derive(Clone, Debug, PartialEq)]
pub struct DynArray<S: primitives::Integer, T: Parcel>
{
    pub elements: Vec<T>,
    _a: std::marker::PhantomData<S>,
}

impl<S: primitives::Integer, T: Parcel> DynArray<S,T>
{
    pub fn new(elements: Vec<T>) -> Self {
        DynArray { elements: elements, _a: std::marker::PhantomData }
    }
}

impl<S: primitives::Integer, T: Parcel> Parcel for DynArray<S, T>
{
    fn read(read: &mut Read) -> Result<Self, Error> {
        let size = S::read(read)?;
        let size: usize = size.try_into()?;

        let mut elements = Vec::new();

        for _ in 0..size {
            elements.push(T::read(read)?);
        }

        Ok(Self::new(elements))
    }

    fn write(&self, write: &mut Write) -> Result<(), Error> {
        S::try_from(self.elements.len())?.write(write)?;

        for element in self.elements.iter() {
            element.write(write)?;
        }

        Ok(())
    }
}

impl<T: Parcel> Parcel for Vec<T>
{
    fn read(read: &mut Read) -> Result<Self, Error> {
        Ok(DynArray::<SizeType, T>::read(read)?.elements)
    }

    fn write(&self, write: &mut Write) -> Result<(), Error> {
        DynArray::<SizeType, T>::new(self.clone()).write(write)
    }
}

<commit_msg>Implement Parcel for array types<commit_after>use {primitives, Parcel, Error};

use std::io::prelude::*;
use std;

/// The integer type that we will use to send length prefixes.
pub type SizeType = u32;

/// An dynamic array type with a custom size prefix type.
#[derive(Clone, Debug, PartialEq)]
pub struct DynArray<S: primitives::Integer, T: Parcel>
{
    pub elements: Vec<T>,
    _a: std::marker::PhantomData<S>,
}

impl<S: primitives::Integer, T: Parcel> DynArray<S,T>
{
    pub fn new(elements: Vec<T>) -> Self {
        DynArray { elements: elements, _a: std::marker::PhantomData }
    }
}

impl<S: primitives::Integer, T: Parcel> Parcel for DynArray<S, T>
{
    fn read(read: &mut Read) -> Result<Self, Error> {
        let size = S::read(read)?;
        let size: usize = size.try_into()?;

        let mut elements = Vec::new();

        for _ in 0..size {
            elements.push(T::read(read)?);
        }

        Ok(Self::new(elements))
    }

    fn write(&self, write: &mut Write) -> Result<(), Error> {
        S::try_from(self.elements.len())?.write(write)?;

        for element in self.elements.iter() {
            element.write(write)?;
        }

        Ok(())
    }
}

macro_rules! impl_parcel_for_array {
    ($n:expr) => {
        impl<T: Parcel> Parcel for [T; $n] where T: Copy {
            fn read(read: &mut Read) -> Result<Self, Error> {
                use std::mem;

                let mut elements: Vec<T> = Vec::with_capacity($n);

                for _ in 0..$n {
                    let elem = T::read(read)?;
                    elements.push(elem);
                }

                let mut array: [T; $n] = unsafe { mem::uninitialized() };
                array.clone_from_slice(&elements[..]);

                Ok(array)
            }

            fn write(&self, write: &mut Write) -> Result<(), Error> {
                for elem in self.iter() {
                    elem.write(write)?;
                }

                Ok(())
            }
        }
    }
}

impl_parcel_for_array!(1);
impl_parcel_for_array!(2);
impl_parcel_for_array!(3);
impl_parcel_for_array!(4);
impl_parcel_for_array!(5);
impl_parcel_for_array!(6);
impl_parcel_for_array!(7);
impl_parcel_for_array!(8);
impl_parcel_for_array!(9);
impl_parcel_for_array!(10);
impl_parcel_for_array!(11);
impl_parcel_for_array!(12);
impl_parcel_for_array!(13);
impl_parcel_for_array!(14);
impl_parcel_for_array!(15);
impl_parcel_for_array!(16);
impl_parcel_for_array!(17);
impl_parcel_for_array!(18);
impl_parcel_for_array!(19);
impl_parcel_for_array!(20);
impl_parcel_for_array!(21);
impl_parcel_for_array!(22);
impl_parcel_for_array!(23);
impl_parcel_for_array!(24);
impl_parcel_for_array!(25);
impl_parcel_for_array!(26);
impl_parcel_for_array!(27);
impl_parcel_for_array!(28);
impl_parcel_for_array!(29);
impl_parcel_for_array!(30);
impl_parcel_for_array!(31);
impl_parcel_for_array!(32);

impl<T: Parcel> Parcel for Vec<T>
{
    fn read(read: &mut Read) -> Result<Self, Error> {
        Ok(DynArray::<SizeType, T>::read(read)?.elements)
    }

    fn write(&self, write: &mut Write) -> Result<(), Error> {
        DynArray::<SizeType, T>::new(self.clone()).write(write)
    }
}

#[cfg(test)]
mod test {
    use Parcel;
    use std::io::Cursor;

    #[test]
    fn can_read_array() {
        let mut data = Cursor::new([0u8, 1, 2, 3]);
        let read_back: [u8; 4] = Parcel::read(&mut data).unwrap();
        assert_eq!(read_back, [0, 1, 2, 3]);
    }

    #[test]
    fn can_write_array() {
        let mut buffer = Cursor::new(Vec::new());
        [5u8, 7, 9, 11].write(&mut buffer).unwrap();
        assert_eq!(buffer.into_inner(), vec![5, 7, 9, 11]);
    }
}
<|endoftext|>


<commit_before>extern crate websocket;

use super::super::*;
use self::websocket::{Client, Message, Receiver, Sender};
use self::websocket::client::request::Url;

pub static PROVIDER: &'static Provider = &WonziuProvider;

struct WonziuProvider;

lazy_static!{
    static ref STATUS: Mutex<WonziuData> = Mutex::new(WonziuData{stream: None, topic: None});
}

impl Provider for WonziuProvider {
    fn start(&self, config: &Json) -> Option<thread::JoinHandle<()>>{
        use self::websocket::message::Type::*;
        let config = config.clone();
        Some(thread::spawn(move || {
            let url = Url::parse("wss://api.pancernik.info/notifier").unwrap();
            let request = Client::connect(url).unwrap();
            let response = request.send().unwrap();
            response.validate().unwrap();
            let client = response.begin();
            let (mut sender, mut receiver) = client.split();

            for message in receiver.incoming_messages() {
                let mut message: Message = message.unwrap();
                match message.opcode {
                    Text | Binary => {
                        println!("{}", String::from_utf8_lossy(&message.payload));
                        if let Some(msg) = handle_message(serde_json::from_slice(&message.payload).unwrap()) {
                            sender.send_message(&msg).unwrap();
                        }
                    },
                    Ping => {
                        message.into_pong().unwrap();
                        sender.send_message(&message).unwrap();
                    },
                    Pong => {},
                    Close => {PROVIDER.start(&config); panic!("Wonziu websocket closed!")}
                }
            }
        }))
    }

    fn load_feed(&self, _data: &Json) -> Result<Feed, Box<Error>> {
        let mut feed = Feed::new();

        let stream;
        let topic;
        {
            let status = STATUS.lock().unwrap();
            if status.stream.is_none() {
                return Ok(feed);
            }
            stream = status.stream.clone().unwrap();
            topic = status.topic.clone().unwrap();
        }

        if stream.status {
            feed.status.push(Entry::new(&topic.text, &hash(&(&stream.online_at, "wonziu")))
                    .set_timestamp(time::strptime(&stream.online_at, "%Y-%m-%dT%H:%M:%S")
                            .ok()
                            .map(|tm| to_timestamp(tm)))
                    .link("http://jadisco.pl/"))
        }

        Ok(feed)
    }
}

#[derive(Deserialize, Serialize)]
struct WonziuMessage {
    #[serde(rename="type")]
    kind: String,
    data: Option<WonziuData>,
}

#[derive(Deserialize, Serialize)]
struct WonziuData {
    stream: Option<WonziuStream>,
    topic: Option<WonziuTopic>,
}

#[derive(Deserialize, Serialize, Clone)]
struct WonziuStream {
    status: bool,
    online_at: String,
}

#[derive(Deserialize, Serialize, Clone)]
struct WonziuTopic {
    text: String,
}

fn handle_message(message: WonziuMessage) -> Option<Message<'static>> {
    match &message.kind[..] {
        "status" => {
            *STATUS.lock().unwrap() = message.data.unwrap();
            None
        }
        "update" => {
            let status = message.data.unwrap();
            if let Some(stream) = status.stream {
                STATUS.lock().unwrap().stream = Some(stream);
            }
            if let Some(topic) = status.topic {
                STATUS.lock().unwrap().topic = Some(topic);
            }
            None
        }
        "ping" => Some(Message::text(serde_json::to_string(&WonziuMessage{kind: "pong".to_string(), data:None}).unwrap())),
        _ => unreachable!(),
    }
}
<commit_msg>Fixed incorect merging update in wonziu.<commit_after>extern crate websocket;

use super::super::*;
use self::websocket::{Client, Message, Receiver, Sender};
use self::websocket::client::request::Url;

pub static PROVIDER: &'static Provider = &WonziuProvider;

struct WonziuProvider;

lazy_static!{
    static ref STATUS: Mutex<WonziuData> = Mutex::new(WonziuData{stream: None, topic: None});
}

impl Provider for WonziuProvider {
    fn start(&self, config: &Json) -> Option<thread::JoinHandle<()>>{
        use self::websocket::message::Type::*;
        let config = config.clone();
        Some(thread::spawn(move || {
            let url = Url::parse("wss://api.pancernik.info/notifier").unwrap();
            let request = Client::connect(url).unwrap();
            let response = request.send().unwrap();
            response.validate().unwrap();
            let client = response.begin();
            let (mut sender, mut receiver) = client.split();

            for message in receiver.incoming_messages() {
                let mut message: Message = message.unwrap();
                match message.opcode {
                    Text | Binary => {
                        if let Some(msg) = handle_message(serde_json::from_slice(&message.payload).unwrap()) {
                            sender.send_message(&msg).unwrap();
                        }
                    },
                    Ping => {
                        message.into_pong().unwrap();
                        sender.send_message(&message).unwrap();
                    },
                    Pong => {},
                    Close => {PROVIDER.start(&config); panic!("Wonziu websocket closed!")}
                }
            }
        }))
    }

    fn load_feed(&self, _data: &Json) -> Result<Feed, Box<Error>> {
        let mut feed = Feed::new();

        let stream;
        let topic;
        {
            let status = STATUS.lock().unwrap();
            if status.stream.is_none() {
                return Ok(feed);
            }
            stream = status.stream.clone().unwrap();
            topic = status.topic.clone().unwrap();
        }

        if stream.status {
            feed.status.push(Entry::new(&topic.text, &hash(&(&stream.online_at, "wonziu")))
                    .set_timestamp(time::strptime(&stream.online_at, "%Y-%m-%dT%H:%M:%S")
                            .ok()
                            .map(|tm| to_timestamp(tm)))
                    .link("http://jadisco.pl/"))
        }

        Ok(feed)
    }
}

#[derive(Deserialize, Serialize)]
struct WonziuMessage {
    #[serde(rename="type")]
    kind: String,
    data: Option<Json>,
}

#[derive(Deserialize, Serialize)]
struct WonziuData {
    stream: Option<WonziuStream>,
    topic: Option<WonziuTopic>,
}

#[derive(Deserialize, Serialize, Clone)]
struct WonziuStream {
    status: bool,
    online_at: String,
}

#[derive(Deserialize, Serialize, Clone)]
struct WonziuTopic {
    text: String,
}

fn handle_message(message: WonziuMessage) -> Option<Message<'static>> {
    match &message.kind[..] {
        "status" => {
            *STATUS.lock().unwrap() = serde_json::from_value(message.data.unwrap()).unwrap();
            None
        }
        "update" => {
            let status = serde_json::to_value(&*STATUS.lock().unwrap());
            let status_new = message.data.unwrap();
            *STATUS.lock().unwrap() = serde_json::from_value(json_merge(status, status_new)).unwrap();
            None
        }
        "ping" => Some(Message::text(serde_json::to_string(&WonziuMessage{kind: "pong".to_string(), data:None}).unwrap())),
        _ => unreachable!(),
    }
}

fn json_merge(obj: Json, diff: Json) -> Json {
    match diff {
        Json::Array(diff) => {
                if let Json::Array(mut obj) = obj {
                    for (n, el) in diff.into_iter().enumerate() {
                        let merged = json_merge(obj[n].clone(), el);
                        obj.insert(n, merged);
                    }
                    Json::Array(obj)
                } else {
                    Json::Array(diff)
                }
            },
        Json::Object(diff) => {
                if let Json::Object(mut obj) = obj {
                    for (k, v) in diff.into_iter() {
                        let merged = json_merge(obj.remove(&k).unwrap_or(Json::Null), v);
                        obj.insert(k, merged);
                    }
                    Json::Object(obj)
                } else {
                    Json::Object(diff)
                }
            },
        _ => diff,
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 Serde Developers
//
// 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 ast::{Data, Container, Style};
use attr::Identifier;
use Ctxt;

/// Cross-cutting checks that require looking at more than a single attrs
/// object. Simpler checks should happen when parsing and building the attrs.
pub fn check(cx: &Ctxt, cont: &Container) {
    check_getter(cx, cont);
    check_identifier(cx, cont);
    check_variant_skip_attrs(cx, cont);
}

/// Getters are only allowed inside structs (not enums) with the `remote`
/// attribute.
fn check_getter(cx: &Ctxt, cont: &Container) {
    match cont.data {
        Data::Enum(_) => {
            if cont.data.has_getter() {
                cx.error("#[serde(getter = \"...\")] is not allowed in an enum");
            }
        }
        Data::Struct(_, _) => {
            if cont.data.has_getter() && cont.attrs.remote().is_none() {
                cx.error(
                    "#[serde(getter = \"...\")] can only be used in structs \
                     that have #[serde(remote = \"...\")]",
                );
            }
        }
    }
}

/// The `other` attribute must be used at most once and it must be the last
/// variant of an enum that has the `field_identifier` attribute.
///
/// Inside a `variant_identifier` all variants must be unit variants. Inside a
/// `field_identifier` all but possibly one variant must be unit variants. The
/// last variant may be a newtype variant which is an implicit "other" case.
fn check_identifier(cx: &Ctxt, cont: &Container) {
    let variants = match cont.data {
        Data::Enum(ref variants) => variants,
        Data::Struct(_, _) => {
            return;
        }
    };

    for (i, variant) in variants.iter().enumerate() {
        match (
            variant.style,
            cont.attrs.identifier(),
            variant.attrs.other(),
        ) {
            // The `other` attribute may only be used in a field_identifier.
            (_, Identifier::Variant, true) | (_, Identifier::No, true) => {
                cx.error("#[serde(other)] may only be used inside a field_identifier");
            }

            // Variant with `other` attribute must be the last one.
            (Style::Unit, Identifier::Field, true) => {
                if i < variants.len() - 1 {
                    cx.error("#[serde(other)] must be the last variant");
                }
            }

            // Variant with `other` attribute must be a unit variant.
            (_, Identifier::Field, true) => {
                cx.error("#[serde(other)] must be on a unit variant");
            }

            // Any sort of variant is allowed if this is not an identifier.
            (_, Identifier::No, false) => {}

            // Unit variant without `other` attribute is always fine.
            (Style::Unit, _, false) => {}

            // The last field is allowed to be a newtype catch-all.
            (Style::Newtype, Identifier::Field, false) => {
                if i < variants.len() - 1 {
                    cx.error(format!("`{}` must be the last variant", variant.ident));
                }
            }

            (_, Identifier::Field, false) => {
                cx.error("field_identifier may only contain unit variants");
            }

            (_, Identifier::Variant, false) => {
                cx.error("variant_identifier may only contain unit variants");
            }
        }
    }
}

/// Skip-(de)serializing attributes are not allowed on variants marked
/// (de)serialize_with.
fn check_variant_skip_attrs(cx: &Ctxt, cont: &Container) {
    let variants = match cont.data {
        Data::Enum(ref variants) => variants,
        Data::Struct(_, _) => {
            return;
        }
    };

    for variant in variants.iter() {
        if variant.attrs.serialize_with().is_some() {
            if variant.attrs.skip_serializing() {
                cx.error(format!(
                    "variant `{}` cannot have both #[serde(serialize_with)] and \
                     #[serde(skip_serializing)]",
                    variant.ident
                ));
            }

            for (i, field) in variant.fields.iter().enumerate() {
                let ident = field
                    .ident
                    .as_ref()
                    .map_or_else(|| format!("{}", i), |ident| format!("`{}`", ident));

                if field.attrs.skip_serializing() {
                    cx.error(format!(
                        "variant `{}` cannot have both #[serde(serialize_with)] and \
                         a field {} marked with #[serde(skip_serializing)]",
                        variant.ident, ident
                    ));
                }

                if field.attrs.skip_serializing_if().is_some() {
                    cx.error(format!(
                        "variant `{}` cannot have both #[serde(serialize_with)] and \
                         a field {} marked with #[serde(skip_serializing_if)]",
                        variant.ident, ident
                    ));
                }
            }
        }

        if variant.attrs.deserialize_with().is_some() {
            if variant.attrs.skip_deserializing() {
                cx.error(format!(
                    "variant `{}` cannot have both #[serde(deserialize_with)] and \
                     #[serde(skip_deserializing)]",
                    variant.ident
                ));
            }

            for (i, field) in variant.fields.iter().enumerate() {
                if field.attrs.skip_deserializing() {
                    let ident = field
                        .ident
                        .as_ref()
                        .map_or_else(|| format!("{}", i), |ident| format!("`{}`", ident));

                    cx.error(format!(
                        "variant `{}` cannot have both #[serde(deserialize_with)] \
                         and a field {} marked with #[serde(skip_deserializing)]",
                        variant.ident, ident
                    ));
                }
            }
        }
    }
}
<commit_msg>Disallow variant field names to conflict with tag of internally-tagged enum<commit_after>// Copyright 2017 Serde Developers
//
// 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 ast::{Data, Container, Style};
use attr::{Identifier, EnumTag};
use Ctxt;

/// Cross-cutting checks that require looking at more than a single attrs
/// object. Simpler checks should happen when parsing and building the attrs.
pub fn check(cx: &Ctxt, cont: &Container) {
    check_getter(cx, cont);
    check_identifier(cx, cont);
    check_variant_skip_attrs(cx, cont);
    check_internally_tagged_variant_name_conflict(cx, cont);
}

/// Getters are only allowed inside structs (not enums) with the `remote`
/// attribute.
fn check_getter(cx: &Ctxt, cont: &Container) {
    match cont.data {
        Data::Enum(_) => {
            if cont.data.has_getter() {
                cx.error("#[serde(getter = \"...\")] is not allowed in an enum");
            }
        }
        Data::Struct(_, _) => {
            if cont.data.has_getter() && cont.attrs.remote().is_none() {
                cx.error(
                    "#[serde(getter = \"...\")] can only be used in structs \
                     that have #[serde(remote = \"...\")]",
                );
            }
        }
    }
}

/// The `other` attribute must be used at most once and it must be the last
/// variant of an enum that has the `field_identifier` attribute.
///
/// Inside a `variant_identifier` all variants must be unit variants. Inside a
/// `field_identifier` all but possibly one variant must be unit variants. The
/// last variant may be a newtype variant which is an implicit "other" case.
fn check_identifier(cx: &Ctxt, cont: &Container) {
    let variants = match cont.data {
        Data::Enum(ref variants) => variants,
        Data::Struct(_, _) => {
            return;
        }
    };

    for (i, variant) in variants.iter().enumerate() {
        match (
            variant.style,
            cont.attrs.identifier(),
            variant.attrs.other(),
        ) {
            // The `other` attribute may only be used in a field_identifier.
            (_, Identifier::Variant, true) | (_, Identifier::No, true) => {
                cx.error("#[serde(other)] may only be used inside a field_identifier");
            }

            // Variant with `other` attribute must be the last one.
            (Style::Unit, Identifier::Field, true) => {
                if i < variants.len() - 1 {
                    cx.error("#[serde(other)] must be the last variant");
                }
            }

            // Variant with `other` attribute must be a unit variant.
            (_, Identifier::Field, true) => {
                cx.error("#[serde(other)] must be on a unit variant");
            }

            // Any sort of variant is allowed if this is not an identifier.
            (_, Identifier::No, false) => {}

            // Unit variant without `other` attribute is always fine.
            (Style::Unit, _, false) => {}

            // The last field is allowed to be a newtype catch-all.
            (Style::Newtype, Identifier::Field, false) => {
                if i < variants.len() - 1 {
                    cx.error(format!("`{}` must be the last variant", variant.ident));
                }
            }

            (_, Identifier::Field, false) => {
                cx.error("field_identifier may only contain unit variants");
            }

            (_, Identifier::Variant, false) => {
                cx.error("variant_identifier may only contain unit variants");
            }
        }
    }
}

/// Skip-(de)serializing attributes are not allowed on variants marked
/// (de)serialize_with.
fn check_variant_skip_attrs(cx: &Ctxt, cont: &Container) {
    let variants = match cont.data {
        Data::Enum(ref variants) => variants,
        Data::Struct(_, _) => {
            return;
        }
    };

    for variant in variants.iter() {
        if variant.attrs.serialize_with().is_some() {
            if variant.attrs.skip_serializing() {
                cx.error(format!(
                    "variant `{}` cannot have both #[serde(serialize_with)] and \
                     #[serde(skip_serializing)]",
                    variant.ident
                ));
            }

            for (i, field) in variant.fields.iter().enumerate() {
                let ident = field
                    .ident
                    .as_ref()
                    .map_or_else(|| format!("{}", i), |ident| format!("`{}`", ident));

                if field.attrs.skip_serializing() {
                    cx.error(format!(
                        "variant `{}` cannot have both #[serde(serialize_with)] and \
                         a field {} marked with #[serde(skip_serializing)]",
                        variant.ident, ident
                    ));
                }

                if field.attrs.skip_serializing_if().is_some() {
                    cx.error(format!(
                        "variant `{}` cannot have both #[serde(serialize_with)] and \
                         a field {} marked with #[serde(skip_serializing_if)]",
                        variant.ident, ident
                    ));
                }
            }
        }

        if variant.attrs.deserialize_with().is_some() {
            if variant.attrs.skip_deserializing() {
                cx.error(format!(
                    "variant `{}` cannot have both #[serde(deserialize_with)] and \
                     #[serde(skip_deserializing)]",
                    variant.ident
                ));
            }

            for (i, field) in variant.fields.iter().enumerate() {
                if field.attrs.skip_deserializing() {
                    let ident = field
                        .ident
                        .as_ref()
                        .map_or_else(|| format!("{}", i), |ident| format!("`{}`", ident));

                    cx.error(format!(
                        "variant `{}` cannot have both #[serde(deserialize_with)] \
                         and a field {} marked with #[serde(skip_deserializing)]",
                        variant.ident, ident
                    ));
                }
            }
        }
    }
}

fn check_internally_tagged_variant_name_conflict(
    cx: &Ctxt,
    cont: &Container,
) {
    let variants = match cont.data {
        Data::Enum(ref variants) => variants,
        Data::Struct(_, _) => return,
    };

    let tag = match *cont.attrs.tag() {
        EnumTag::Internal { ref tag } => tag.as_str(),
        EnumTag::External | EnumTag::Adjacent { .. } | EnumTag::None => return,
    };

    let diagnose_conflict = || {
        let message = format!(
            "variant field name `{}` conflicts with internal tag",
            tag
        );
        cx.error(message);
    };

    for variant in variants {
        match variant.style {
            Style::Struct => {
                for field in &variant.fields {
                    let check_ser = !field.attrs.skip_serializing();
                    let check_de = !field.attrs.skip_deserializing();
                    let name = field.attrs.name();
                    let ser_name = name.serialize_name();
                    let de_name = name.deserialize_name();

                    if check_ser && ser_name == tag || check_de && de_name == tag {
                        diagnose_conflict();
                        return;
                    }
                }
            },
            _ => {},
        }
    }
}
<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2014 T. Jameson Little <t.jameson.little@gmail.com>
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

#[macro_use]
extern crate uucore;

use std::error::Error;
use std::fmt::Write as FmtWrite;
use std::io::{self, stdin, stdout, BufRead, Write};

mod factor;
use clap::{crate_version, App, Arg};
pub use factor::*;
use uucore::display::Quotable;

mod miller_rabin;
pub mod numeric;
mod rho;
pub mod table;

static SUMMARY: &str = "Print the prime factors of the given NUMBER(s).
If none are specified, read from standard input.";

mod options {
    pub static NUMBER: &str = "NUMBER";
}

fn print_factors_str(
    num_str: &str,
    w: &mut io::BufWriter<impl io::Write>,
    factors_buffer: &mut String,
) -> Result<(), Box<dyn Error>> {
    num_str.parse::<u64>().map_err(|e| e.into()).and_then(|x| {
        factors_buffer.clear();
        writeln!(factors_buffer, "{}:{}", x, factor(x))?;
        w.write_all(factors_buffer.as_bytes())?;
        Ok(())
    })
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = uu_app().get_matches_from(args);
    let stdout = stdout();
    // We use a smaller buffer here to pass a gnu test. 4KiB appears to be the default pipe size for bash.
    let mut w = io::BufWriter::with_capacity(4 * 1024, stdout.lock());
    let mut factors_buffer = String::new();

    if let Some(values) = matches.values_of(options::NUMBER) {
        for number in values {
            if let Err(e) = print_factors_str(number, &mut w, &mut factors_buffer) {
                show_warning!("{}: {}", number.maybe_quote(), e);
            }
        }
    } else {
        let stdin = stdin();

        for line in stdin.lock().lines() {
            for number in line.unwrap().split_whitespace() {
                if let Err(e) = print_factors_str(number, &mut w, &mut factors_buffer) {
                    show_warning!("{}: {}", number.maybe_quote(), e);
                }
            }
        }
    }

    if let Err(e) = w.flush() {
        show_error!("{}", e);
    }

    0
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(uucore::util_name())
        .version(crate_version!())
        .about(SUMMARY)
        .arg(Arg::with_name(options::NUMBER).multiple(true))
}
<commit_msg>factor: return UResult from uumain() function<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2014 T. Jameson Little <t.jameson.little@gmail.com>
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

#[macro_use]
extern crate uucore;

use std::error::Error;
use std::fmt::Write as FmtWrite;
use std::io::{self, stdin, stdout, BufRead, Write};

mod factor;
use clap::{crate_version, App, Arg};
pub use factor::*;
use uucore::display::Quotable;
use uucore::error::UResult;

mod miller_rabin;
pub mod numeric;
mod rho;
pub mod table;

static SUMMARY: &str = "Print the prime factors of the given NUMBER(s).
If none are specified, read from standard input.";

mod options {
    pub static NUMBER: &str = "NUMBER";
}

fn print_factors_str(
    num_str: &str,
    w: &mut io::BufWriter<impl io::Write>,
    factors_buffer: &mut String,
) -> Result<(), Box<dyn Error>> {
    num_str.parse::<u64>().map_err(|e| e.into()).and_then(|x| {
        factors_buffer.clear();
        writeln!(factors_buffer, "{}:{}", x, factor(x))?;
        w.write_all(factors_buffer.as_bytes())?;
        Ok(())
    })
}

#[uucore_procs::gen_uumain]
pub fn uumain(args: impl uucore::Args) -> UResult<()> {
    let matches = uu_app().get_matches_from(args);
    let stdout = stdout();
    // We use a smaller buffer here to pass a gnu test. 4KiB appears to be the default pipe size for bash.
    let mut w = io::BufWriter::with_capacity(4 * 1024, stdout.lock());
    let mut factors_buffer = String::new();

    if let Some(values) = matches.values_of(options::NUMBER) {
        for number in values {
            if let Err(e) = print_factors_str(number, &mut w, &mut factors_buffer) {
                show_warning!("{}: {}", number.maybe_quote(), e);
            }
        }
    } else {
        let stdin = stdin();

        for line in stdin.lock().lines() {
            for number in line.unwrap().split_whitespace() {
                if let Err(e) = print_factors_str(number, &mut w, &mut factors_buffer) {
                    show_warning!("{}: {}", number.maybe_quote(), e);
                }
            }
        }
    }

    if let Err(e) = w.flush() {
        show_error!("{}", e);
    }

    Ok(())
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(uucore::util_name())
        .version(crate_version!())
        .about(SUMMARY)
        .arg(Arg::with_name(options::NUMBER).multiple(true))
}
<|endoftext|>