<commit_before>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::cmp;
use std::collections::HashMap;

use crate::{Alloc, Error, Result};

/// Manages allocating address ranges.
/// Use `AddressAllocator` whenever an address range needs to be allocated to different users.
/// Allocations must be uniquely tagged with an Alloc enum, which can be used for lookup.
/// An human-readable tag String must also be provided for debugging / reference.
///
/// # Examples
///
/// ```
/// // Anon is used for brevity. Don't manually instantiate Anon allocs!
/// # use resources::{Alloc, AddressAllocator};
///   AddressAllocator::new(0x1000, 0x10000, Some(0x100)).map(|mut pool| {
///       assert_eq!(pool.allocate(0x110, Alloc::Anon(0), "caps".to_string()), Ok(0x1000));
///       assert_eq!(pool.allocate(0x100, Alloc::Anon(1), "cache".to_string()), Ok(0x1200));
///       assert_eq!(pool.allocate(0x100, Alloc::Anon(2), "etc".to_string()), Ok(0x1300));
///       assert_eq!(pool.get(&Alloc::Anon(1)), Some(&(0x1200, 0x100, "cache".to_string())));
///   });
/// ```
#[derive(Debug, Eq, PartialEq)]
pub struct AddressAllocator {
    pool_base: u64,
    pool_end: u64,
    alignment: u64,
    next_addr: u64,
    allocs: HashMap<Alloc, (u64, u64, String)>,
}

impl AddressAllocator {
    /// Creates a new `AddressAllocator` for managing a range of addresses.
    /// Can return `None` if `pool_base` + `pool_size` overflows a u64 or if alignment isn't a power
    /// of two.
    ///
    /// * `pool_base` - The starting address of the range to manage.
    /// * `pool_size` - The size of the address range in bytes.
    /// * `align_size` - The minimum size of an address region to align to, defaults to four.
    pub fn new(pool_base: u64, pool_size: u64, align_size: Option<u64>) -> Result<Self> {
        if pool_size == 0 {
            return Err(Error::PoolSizeZero);
        }
        let pool_end = pool_base
            .checked_add(pool_size - 1)
            .ok_or(Error::PoolOverflow {
                base: pool_base,
                size: pool_size,
            })?;
        let alignment = align_size.unwrap_or(4);
        if !alignment.is_power_of_two() || alignment == 0 {
            return Err(Error::BadAlignment);
        }
        Ok(AddressAllocator {
            pool_base,
            pool_end,
            alignment,
            next_addr: pool_base,
            allocs: HashMap::new(),
        })
    }

    /// Allocates a range of addresses from the managed region with an optional tag
    /// and minimal alignment. Returns allocated_address. (allocated_address, size, tag)
    /// can be retrieved through the `get` method.
    pub fn allocate_with_align(
        &mut self,
        size: u64,
        alloc: Alloc,
        tag: String,
        alignment: u64,
    ) -> Result<u64> {
        let alignment = cmp::max(self.alignment, alignment);

        if self.allocs.contains_key(&alloc) {
            return Err(Error::ExistingAlloc(alloc));
        }
        if size == 0 {
            return Err(Error::AllocSizeZero);
        }
        if !alignment.is_power_of_two() {
            return Err(Error::BadAlignment);
        }
        let align_adjust = if self.next_addr % alignment != 0 {
            alignment - (self.next_addr % alignment)
        } else {
            0
        };
        let addr = self
            .next_addr
            .checked_add(align_adjust)
            .ok_or(Error::OutOfSpace)?;
        let end_addr = addr.checked_add(size - 1).ok_or(Error::OutOfSpace)?;
        if end_addr > self.pool_end {
            return Err(Error::OutOfSpace);
        }

        // TODO(dgreid): Use a smarter allocation strategy. The current strategy is just
        // bumping this pointer, meaning it will eventually exhaust available addresses.
        self.next_addr = end_addr.saturating_add(1);

        self.allocs.insert(alloc, (addr, size, tag));
        Ok(addr)
    }

    pub fn allocate(&mut self, size: u64, alloc: Alloc, tag: String) -> Result<u64> {
        self.allocate_with_align(size, alloc, tag, self.alignment)
    }

    /// Returns allocation associated with `alloc`, or None if no such allocation exists.
    pub fn get(&self, alloc: &Alloc) -> Option<&(u64, u64, String)> {
        self.allocs.get(alloc)
    }
}

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

    #[test]
    fn new_fails_overflow() {
        assert!(AddressAllocator::new(u64::max_value(), 0x100, None).is_err());
    }

    #[test]
    fn new_fails_size_zero() {
        assert!(AddressAllocator::new(0x1000, 0, None).is_err());
    }

    #[test]
    fn new_fails_alignment_zero() {
        assert!(AddressAllocator::new(0x1000, 0x10000, Some(0)).is_err());
    }

    #[test]
    fn new_fails_alignment_non_power_of_two() {
        assert!(AddressAllocator::new(0x1000, 0x10000, Some(200)).is_err());
    }

    #[test]
    fn allocate_fails_exising_alloc() {
        let mut pool = AddressAllocator::new(0x1000, 0x1000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate(0x800, Alloc::Anon(0), String::from("bar0")),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate(0x800, Alloc::Anon(0), String::from("bar0")),
            Err(Error::ExistingAlloc(Alloc::Anon(0)))
        );
    }

    #[test]
    fn allocate_fails_not_enough_space() {
        let mut pool = AddressAllocator::new(0x1000, 0x1000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate(0x800, Alloc::Anon(0), String::from("bar0")),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate(0x900, Alloc::Anon(1), String::from("bar1")),
            Err(Error::OutOfSpace)
        );
        assert_eq!(
            pool.allocate(0x800, Alloc::Anon(2), String::from("bar2")),
            Ok(0x1800)
        );
    }

    #[test]
    fn allocate_alignment() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate(0x110, Alloc::Anon(0), String::from("bar0")),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate(0x100, Alloc::Anon(1), String::from("bar1")),
            Ok(0x1200)
        );
    }

    #[test]
    fn allocate_retrieve_alloc() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate(0x110, Alloc::Anon(0), String::from("bar0")),
            Ok(0x1000)
        );
        assert_eq!(
            pool.get(&Alloc::Anon(0)),
            Some(&(0x1000, 0x110, String::from("bar0")))
        );
    }

    #[test]
    fn allocate_with_alignment_allocator_alignment() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 0x1),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate_with_align(0x100, Alloc::Anon(1), String::from("bar1"), 0x1),
            Ok(0x1200)
        );
    }

    #[test]
    fn allocate_with_alignment_custom_alignment() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, Some(0x4)).unwrap();
        assert_eq!(
            pool.allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 0x100),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate_with_align(0x100, Alloc::Anon(1), String::from("bar1"), 0x100),
            Ok(0x1200)
        );
    }

    #[test]
    fn allocate_with_alignment_no_allocator_alignment() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, None).unwrap();
        assert_eq!(
            pool.allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 0x100),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate_with_align(0x100, Alloc::Anon(1), String::from("bar1"), 0x100),
            Ok(0x1200)
        );
    }

    #[test]
    fn allocate_with_alignment_alignment_non_power_of_two() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, None).unwrap();
        assert!(pool
            .allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 200)
            .is_err());
    }
}
<commit_msg>resources: allocate_at support for address_allocator<commit_after>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::cmp;
use std::collections::{BTreeSet, HashMap};

use crate::{Alloc, Error, Result};

/// Manages allocating address ranges.
/// Use `AddressAllocator` whenever an address range needs to be allocated to different users.
/// Allocations must be uniquely tagged with an Alloc enum, which can be used for lookup.
/// An human-readable tag String must also be provided for debugging / reference.
///
/// # Examples
///
/// ```
/// // Anon is used for brevity. Don't manually instantiate Anon allocs!
/// # use resources::{Alloc, AddressAllocator};
///   AddressAllocator::new(0x1000, 0x10000, Some(0x100)).map(|mut pool| {
///       assert_eq!(pool.allocate(0x110, Alloc::Anon(0), "caps".to_string()), Ok(0x1000));
///       assert_eq!(pool.allocate(0x100, Alloc::Anon(1), "cache".to_string()), Ok(0x1200));
///       assert_eq!(pool.allocate(0x100, Alloc::Anon(2), "etc".to_string()), Ok(0x1300));
///       assert_eq!(pool.get(&Alloc::Anon(1)), Some(&(0x1200, 0x100, "cache".to_string())));
///   });
/// ```
#[derive(Debug, Eq, PartialEq)]
pub struct AddressAllocator {
    alignment: u64,
    allocs: HashMap<Alloc, (u64, u64, String)>,
    regions: BTreeSet<(u64, u64)>,
}

impl AddressAllocator {
    /// Creates a new `AddressAllocator` for managing a range of addresses.
    /// Can return `None` if `pool_base` + `pool_size` overflows a u64 or if alignment isn't a power
    /// of two.
    ///
    /// * `pool_base` - The starting address of the range to manage.
    /// * `pool_size` - The size of the address range in bytes.
    /// * `align_size` - The minimum size of an address region to align to, defaults to four.
    pub fn new(pool_base: u64, pool_size: u64, align_size: Option<u64>) -> Result<Self> {
        if pool_size == 0 {
            return Err(Error::PoolSizeZero);
        }
        let pool_end = pool_base
            .checked_add(pool_size - 1)
            .ok_or(Error::PoolOverflow {
                base: pool_base,
                size: pool_size,
            })?;
        let alignment = align_size.unwrap_or(4);
        if !alignment.is_power_of_two() || alignment == 0 {
            return Err(Error::BadAlignment);
        }
        let mut regions = BTreeSet::new();
        regions.insert((pool_base, pool_end));
        Ok(AddressAllocator {
            alignment,
            allocs: HashMap::new(),
            regions,
        })
    }

    /// Allocates a range of addresses from the managed region with an optional tag
    /// and minimal alignment. Returns allocated_address. (allocated_address, size, tag)
    /// can be retrieved through the `get` method.
    pub fn allocate_with_align(
        &mut self,
        size: u64,
        alloc: Alloc,
        tag: String,
        alignment: u64,
    ) -> Result<u64> {
        let alignment = cmp::max(self.alignment, alignment);

        if self.allocs.contains_key(&alloc) {
            return Err(Error::ExistingAlloc(alloc));
        }
        if size == 0 {
            return Err(Error::AllocSizeZero);
        }
        if !alignment.is_power_of_two() {
            return Err(Error::BadAlignment);
        }

        // finds first region matching alignment and size.
        match self
            .regions
            .iter()
            .find(|range| {
                match range.0 % alignment {
                    0 => range.0.checked_add(size - 1),
                    r => range.0.checked_add(size - 1 + alignment - r),
                }
                .map_or(false, |end| end <= range.1)
            })
            .cloned()
        {
            Some(slot) => {
                self.regions.remove(&slot);
                let start = match slot.0 % alignment {
                    0 => slot.0,
                    r => slot.0 + alignment - r,
                };
                let end = start + size - 1;
                if slot.0 < start {
                    self.regions.insert((slot.0, start - 1));
                }
                if slot.1 > end {
                    self.regions.insert((end + 1, slot.1));
                }
                self.allocs.insert(alloc, (start, size, tag));

                Ok(start)
            }
            None => Err(Error::OutOfSpace),
        }
    }

    pub fn allocate(&mut self, size: u64, alloc: Alloc, tag: String) -> Result<u64> {
        self.allocate_with_align(size, alloc, tag, self.alignment)
    }

    /// Allocates a range of addresses from the managed region with an optional tag
    /// and required location. Allocation alignment is not enforced.
    /// Returns OutOfSpace if requested range is not available (e.g. already allocated
    /// with a different alloc tag).
    pub fn allocate_at(&mut self, start: u64, size: u64, alloc: Alloc, tag: String) -> Result<()> {
        if self.allocs.contains_key(&alloc) {
            return Err(Error::ExistingAlloc(alloc));
        }
        if size == 0 {
            return Err(Error::AllocSizeZero);
        }

        let end = start.checked_add(size - 1).ok_or(Error::OutOfSpace)?;
        match self
            .regions
            .iter()
            .find(|range| range.0 <= start && range.1 >= end)
            .cloned()
        {
            Some(slot) => {
                self.regions.remove(&slot);
                if slot.0 < start {
                    self.regions.insert((slot.0, start - 1));
                }
                if slot.1 > end {
                    self.regions.insert((end + 1, slot.1));
                }
                self.allocs.insert(alloc, (start, size, tag));

                Ok(())
            }
            None => Err(Error::OutOfSpace),
        }
    }

    /// Returns allocation associated with `alloc`, or None if no such allocation exists.
    pub fn get(&self, alloc: &Alloc) -> Option<&(u64, u64, String)> {
        self.allocs.get(alloc)
    }
}

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

    #[test]
    fn new_fails_overflow() {
        assert!(AddressAllocator::new(u64::max_value(), 0x100, None).is_err());
    }

    #[test]
    fn new_fails_size_zero() {
        assert!(AddressAllocator::new(0x1000, 0, None).is_err());
    }

    #[test]
    fn new_fails_alignment_zero() {
        assert!(AddressAllocator::new(0x1000, 0x10000, Some(0)).is_err());
    }

    #[test]
    fn new_fails_alignment_non_power_of_two() {
        assert!(AddressAllocator::new(0x1000, 0x10000, Some(200)).is_err());
    }

    #[test]
    fn allocate_fails_exising_alloc() {
        let mut pool = AddressAllocator::new(0x1000, 0x1000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate(0x800, Alloc::Anon(0), String::from("bar0")),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate(0x800, Alloc::Anon(0), String::from("bar0")),
            Err(Error::ExistingAlloc(Alloc::Anon(0)))
        );
    }

    #[test]
    fn allocate_fails_not_enough_space() {
        let mut pool = AddressAllocator::new(0x1000, 0x1000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate(0x800, Alloc::Anon(0), String::from("bar0")),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate(0x900, Alloc::Anon(1), String::from("bar1")),
            Err(Error::OutOfSpace)
        );
        assert_eq!(
            pool.allocate(0x800, Alloc::Anon(2), String::from("bar2")),
            Ok(0x1800)
        );
    }

    #[test]
    fn allocate_with_special_alignment() {
        let mut pool = AddressAllocator::new(0x1000, 0x1000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate(0x10, Alloc::Anon(0), String::from("bar0")),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate_at(0x1200, 0x100, Alloc::Anon(1), String::from("bar1")),
            Ok(())
        );
        assert_eq!(
            pool.allocate_with_align(0x800, Alloc::Anon(2), String::from("bar2"), 0x800),
            Ok(0x1800)
        );
    }

    #[test]
    fn allocate_and_split_allocate_at() {
        let mut pool = AddressAllocator::new(0x1000, 0x1000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate_at(0x1200, 0x800, Alloc::Anon(0), String::from("bar0")),
            Ok(())
        );
        assert_eq!(
            pool.allocate(0x800, Alloc::Anon(1), String::from("bar1")),
            Err(Error::OutOfSpace)
        );
        assert_eq!(
            pool.allocate(0x600, Alloc::Anon(2), String::from("bar2")),
            Ok(0x1a00)
        );
        assert_eq!(
            pool.allocate(0x200, Alloc::Anon(3), String::from("bar3")),
            Ok(0x1000)
        );
    }

    #[test]
    fn allocate_alignment() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate(0x110, Alloc::Anon(0), String::from("bar0")),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate(0x100, Alloc::Anon(1), String::from("bar1")),
            Ok(0x1200)
        );
    }

    #[test]
    fn allocate_retrieve_alloc() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate(0x110, Alloc::Anon(0), String::from("bar0")),
            Ok(0x1000)
        );
        assert_eq!(
            pool.get(&Alloc::Anon(0)),
            Some(&(0x1000, 0x110, String::from("bar0")))
        );
    }

    #[test]
    fn allocate_with_alignment_allocator_alignment() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, Some(0x100)).unwrap();
        assert_eq!(
            pool.allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 0x1),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate_with_align(0x100, Alloc::Anon(1), String::from("bar1"), 0x1),
            Ok(0x1200)
        );
    }

    #[test]
    fn allocate_with_alignment_custom_alignment() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, Some(0x4)).unwrap();
        assert_eq!(
            pool.allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 0x100),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate_with_align(0x100, Alloc::Anon(1), String::from("bar1"), 0x100),
            Ok(0x1200)
        );
    }

    #[test]
    fn allocate_with_alignment_no_allocator_alignment() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, None).unwrap();
        assert_eq!(
            pool.allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 0x100),
            Ok(0x1000)
        );
        assert_eq!(
            pool.allocate_with_align(0x100, Alloc::Anon(1), String::from("bar1"), 0x100),
            Ok(0x1200)
        );
    }

    #[test]
    fn allocate_with_alignment_alignment_non_power_of_two() {
        let mut pool = AddressAllocator::new(0x1000, 0x10000, None).unwrap();
        assert!(pool
            .allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 200)
            .is_err());
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Delix Project Authors. See the AUTHORS file at the top level directory.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

extern crate delix;

mod helper;

use std::io;
use std::iter;
use std::sync::{Arc, Mutex, RwLock, mpsc};
use std::thread;

use delix::node::request;

#[test]
fn single_echo_from_local_without_timeout() {
    helper::set_up();

    let (node, metric) = helper::build_node("localhost:3001", &[], None);
    node.register("echo", Box::new(|request| Ok(request)))
            .unwrap();

    helper::wait_for_services(&[&metric], 1);

    assert_eq!("test message", String::from_utf8_lossy(&node.request_bytes("echo", b"test message").unwrap()));
}

#[test]
fn single_large_echo_from_local_without_timeout() {
    helper::set_up();

    let (node, metric) = helper::build_node("localhost:3011", &[], None);
    node.register("echo", Box::new(|request| Ok(request)))
            .unwrap();

    helper::wait_for_services(&[&metric], 1);

    let request_bytes = iter::repeat(0u8).take(70000).collect::<Vec<_>>();
    let request = Box::new(io::Cursor::new(request_bytes.clone()));
    node.request("echo", request, Box::new(move |mut reader| {
        assert_eq!(Some(70000), io::copy(&mut reader, &mut io::sink()).ok());
    })).unwrap();
}

#[test]
fn single_echo_from_local_with_timeout() {
    helper::set_up();

    let (node, metric) = helper::build_node("localhost:3021", &[], Some(10));
    node.register("echo", Box::new(|request| {
        thread::sleep(::std::time::Duration::from_millis(20));
        Ok(request)
    })).unwrap();

    helper::wait_for_services(&[&metric], 1);

    assert_eq!(Err(request::Error::Timeout), node.request_bytes("echo", b""));
}

#[test]
fn single_echo_from_remote_without_timeout() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3031", &[], None);
    node_one.register("echo", Box::new(|request| Ok(request)))
            .unwrap();

    let (node_two, metric_two) = helper::build_node("localhost:3032", &["localhost:3031"], None);

    helper::wait_for_joined(&[&metric_one, &metric_two]);
    helper::wait_for_services(&[&metric_one, &metric_two], 1);

    assert_eq!("test message", String::from_utf8_lossy(&node_two.request_bytes("echo", b"test message").unwrap()));
}

#[test]
fn single_echo_from_remote_with_timeout() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3041", &[], None);
    node_one.register("echo", Box::new(|request| {
        thread::sleep(::std::time::Duration::from_millis(20));
        Ok(request)
    })).unwrap();

    let (node_two, metric_two) = helper::build_node("localhost:3042", &["localhost:3041"], Some(10));

    helper::wait_for_joined(&[&metric_one, &metric_two]);
    helper::wait_for_services(&[&metric_one, &metric_two], 1);

    assert_eq!(Err(request::Error::Timeout), node_two.request_bytes("echo", b""));
}

#[test]
fn multiple_echos_from_remote() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3051", &[], None);
    node_one.register("echo", Box::new(|request| Ok(request)))
            .unwrap();

    let (node_two, metric_two) = helper::build_node("localhost:3052", &["localhost:3051"], None);

    helper::wait_for_joined(&[&metric_one, &metric_two]);
    helper::wait_for_services(&[&metric_one, &metric_two], 1);

    assert_eq!(b"test message one".to_vec(), node_two.request_bytes("echo", b"test message one").unwrap());
    assert_eq!(b"test message two".to_vec(), node_two.request_bytes("echo", b"test message two").unwrap());
}

#[test]
fn balanced_echos_from_two_remotes() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3061", &[], None);

    let (tx, rx) = mpsc::channel();

    let (node_two, metric_two) = helper::build_node("localhost:3062", &["localhost:3061"], None);
    let tx_clone = Mutex::new(tx.clone());
    node_two.register("echo", Box::new(move |request| {
        tx_clone.lock().unwrap().send("two").unwrap();
        Ok(request)
    })).unwrap();

    let (node_three, metric_three) = helper::build_node("localhost:3063", &["localhost:3061"], None);
    let tx_clone = Mutex::new(tx.clone());
    node_three.register("echo", Box::new(move |request| {
        tx_clone.lock().unwrap().send("three").unwrap();
        Ok(request)
    })).unwrap();

    helper::wait_for_joined(&[&metric_one, &metric_two, &metric_three]);
    helper::wait_for_services(&[&metric_one, &metric_two, &metric_three], 1);

    assert_eq!("test", String::from_utf8_lossy(&node_one.request_bytes("echo", b"test").unwrap()));
    assert_eq!("test", String::from_utf8_lossy(&node_one.request_bytes("echo", b"test").unwrap()));

    helper::assert_contains_all(&["two", "three"], &helper::recv_all(&rx));

    node_three.deregister("echo").unwrap();

    assert_eq!("test", String::from_utf8_lossy(&node_one.request_bytes("echo", b"test").unwrap()));
    assert_eq!("test", String::from_utf8_lossy(&node_one.request_bytes("echo", b"test").unwrap()));

    helper::assert_contains_all(&["two", "two"], &helper::recv_all(&rx));
}

#[test]
fn parallel_requests_while_a_node_joins_and_leaves() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3071", &[], None);
    node_one.register("echo", Box::new(move |request| {
        Ok(request)
    })).unwrap();

    let running = Arc::new(RwLock::new(true));

    let node_one_clone = node_one.clone();
    let running_clone = running.clone();
    let jh_one = thread::spawn(move || {
        while *running_clone.read().unwrap() {
            match node_one_clone.request_bytes("echo", b"test") {
                Ok(response) => assert_eq!("test", String::from_utf8_lossy(&response)),
                Err(error) => println!("got error: {:?}", error),
            }
            thread::sleep(::std::time::Duration::from_millis(10));
        }
    });

    let node_one_clone = node_one.clone();
    let running_clone = running.clone();
    let jh_two = thread::spawn(move || {
        while *running_clone.read().unwrap() {
            match node_one_clone.request_bytes("echo", b"test") {
                Ok(response) => assert_eq!("test", String::from_utf8_lossy(&response)),
                Err(error) => println!("got error: {:?}", error),
            }
            thread::sleep(::std::time::Duration::from_millis(10));
        }
    });

    helper::wait_for_requests(&[&metric_one], 10);

    {
        let (node_two, metric_two) = helper::build_node("localhost:3072", &["localhost:3071"], None);
        node_two.register("echo", Box::new(move |request| {
            Ok(request)
        })).unwrap();

        helper::wait_for_joined(&[&metric_one, &metric_two]);
        helper::wait_for_services(&[&metric_one, &metric_two], 1);
        helper::wait_for_endpoints(&[&metric_one, &metric_two], 2);

        helper::wait_for_requests(&[&metric_one], 100);
    }

    helper::wait_for_requests(&[&metric_one], 200);

    *running.write().unwrap() = false;
    jh_one.join().unwrap();
    jh_two.join().unwrap();
}
<commit_msg>adjusted test timeout<commit_after>// Copyright 2015 The Delix Project Authors. See the AUTHORS file at the top level directory.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

extern crate delix;

mod helper;

use std::io;
use std::iter;
use std::sync::{Arc, Mutex, RwLock, mpsc};
use std::thread;

use delix::node::request;

#[test]
fn single_echo_from_local_without_timeout() {
    helper::set_up();

    let (node, metric) = helper::build_node("localhost:3001", &[], None);
    node.register("echo", Box::new(|request| Ok(request)))
            .unwrap();

    helper::wait_for_services(&[&metric], 1);

    assert_eq!("test message", String::from_utf8_lossy(&node.request_bytes("echo", b"test message").unwrap()));
}

#[test]
fn single_large_echo_from_local_without_timeout() {
    helper::set_up();

    let (node, metric) = helper::build_node("localhost:3011", &[], None);
    node.register("echo", Box::new(|request| Ok(request)))
            .unwrap();

    helper::wait_for_services(&[&metric], 1);

    let request_bytes = iter::repeat(0u8).take(70000).collect::<Vec<_>>();
    let request = Box::new(io::Cursor::new(request_bytes.clone()));
    node.request("echo", request, Box::new(move |mut reader| {
        assert_eq!(Some(70000), io::copy(&mut reader, &mut io::sink()).ok());
    })).unwrap();
}

#[test]
fn single_echo_from_local_with_timeout() {
    helper::set_up();

    let (node, metric) = helper::build_node("localhost:3021", &[], Some(10));
    node.register("echo", Box::new(|request| {
        thread::sleep(::std::time::Duration::from_millis(100));
        Ok(request)
    })).unwrap();

    helper::wait_for_services(&[&metric], 1);

    assert_eq!(Err(request::Error::Timeout), node.request_bytes("echo", b""));
}

#[test]
fn single_echo_from_remote_without_timeout() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3031", &[], None);
    node_one.register("echo", Box::new(|request| Ok(request)))
            .unwrap();

    let (node_two, metric_two) = helper::build_node("localhost:3032", &["localhost:3031"], None);

    helper::wait_for_joined(&[&metric_one, &metric_two]);
    helper::wait_for_services(&[&metric_one, &metric_two], 1);

    assert_eq!("test message", String::from_utf8_lossy(&node_two.request_bytes("echo", b"test message").unwrap()));
}

#[test]
fn single_echo_from_remote_with_timeout() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3041", &[], None);
    node_one.register("echo", Box::new(|request| {
        thread::sleep(::std::time::Duration::from_millis(20));
        Ok(request)
    })).unwrap();

    let (node_two, metric_two) = helper::build_node("localhost:3042", &["localhost:3041"], Some(10));

    helper::wait_for_joined(&[&metric_one, &metric_two]);
    helper::wait_for_services(&[&metric_one, &metric_two], 1);

    assert_eq!(Err(request::Error::Timeout), node_two.request_bytes("echo", b""));
}

#[test]
fn multiple_echos_from_remote() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3051", &[], None);
    node_one.register("echo", Box::new(|request| Ok(request)))
            .unwrap();

    let (node_two, metric_two) = helper::build_node("localhost:3052", &["localhost:3051"], None);

    helper::wait_for_joined(&[&metric_one, &metric_two]);
    helper::wait_for_services(&[&metric_one, &metric_two], 1);

    assert_eq!(b"test message one".to_vec(), node_two.request_bytes("echo", b"test message one").unwrap());
    assert_eq!(b"test message two".to_vec(), node_two.request_bytes("echo", b"test message two").unwrap());
}

#[test]
fn balanced_echos_from_two_remotes() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3061", &[], None);

    let (tx, rx) = mpsc::channel();

    let (node_two, metric_two) = helper::build_node("localhost:3062", &["localhost:3061"], None);
    let tx_clone = Mutex::new(tx.clone());
    node_two.register("echo", Box::new(move |request| {
        tx_clone.lock().unwrap().send("two").unwrap();
        Ok(request)
    })).unwrap();

    let (node_three, metric_three) = helper::build_node("localhost:3063", &["localhost:3061"], None);
    let tx_clone = Mutex::new(tx.clone());
    node_three.register("echo", Box::new(move |request| {
        tx_clone.lock().unwrap().send("three").unwrap();
        Ok(request)
    })).unwrap();

    helper::wait_for_joined(&[&metric_one, &metric_two, &metric_three]);
    helper::wait_for_services(&[&metric_one, &metric_two, &metric_three], 1);

    assert_eq!("test", String::from_utf8_lossy(&node_one.request_bytes("echo", b"test").unwrap()));
    assert_eq!("test", String::from_utf8_lossy(&node_one.request_bytes("echo", b"test").unwrap()));

    helper::assert_contains_all(&["two", "three"], &helper::recv_all(&rx));

    node_three.deregister("echo").unwrap();

    assert_eq!("test", String::from_utf8_lossy(&node_one.request_bytes("echo", b"test").unwrap()));
    assert_eq!("test", String::from_utf8_lossy(&node_one.request_bytes("echo", b"test").unwrap()));

    helper::assert_contains_all(&["two", "two"], &helper::recv_all(&rx));
}

#[test]
fn parallel_requests_while_a_node_joins_and_leaves() {
    helper::set_up();

    let (node_one, metric_one) = helper::build_node("localhost:3071", &[], None);
    node_one.register("echo", Box::new(move |request| {
        Ok(request)
    })).unwrap();

    let running = Arc::new(RwLock::new(true));

    let node_one_clone = node_one.clone();
    let running_clone = running.clone();
    let jh_one = thread::spawn(move || {
        while *running_clone.read().unwrap() {
            match node_one_clone.request_bytes("echo", b"test") {
                Ok(response) => assert_eq!("test", String::from_utf8_lossy(&response)),
                Err(error) => println!("got error: {:?}", error),
            }
            thread::sleep(::std::time::Duration::from_millis(10));
        }
    });

    let node_one_clone = node_one.clone();
    let running_clone = running.clone();
    let jh_two = thread::spawn(move || {
        while *running_clone.read().unwrap() {
            match node_one_clone.request_bytes("echo", b"test") {
                Ok(response) => assert_eq!("test", String::from_utf8_lossy(&response)),
                Err(error) => println!("got error: {:?}", error),
            }
            thread::sleep(::std::time::Duration::from_millis(10));
        }
    });

    helper::wait_for_requests(&[&metric_one], 10);

    {
        let (node_two, metric_two) = helper::build_node("localhost:3072", &["localhost:3071"], None);
        node_two.register("echo", Box::new(move |request| {
            Ok(request)
        })).unwrap();

        helper::wait_for_joined(&[&metric_one, &metric_two]);
        helper::wait_for_services(&[&metric_one, &metric_two], 1);
        helper::wait_for_endpoints(&[&metric_one, &metric_two], 2);

        helper::wait_for_requests(&[&metric_one], 100);
    }

    helper::wait_for_requests(&[&metric_one], 200);

    *running.write().unwrap() = false;
    jh_one.join().unwrap();
    jh_two.join().unwrap();
}
<|endoftext|>


<commit_before>mod helpers;

use helpers::{roundtrip, wayc, ways, TestClient, TestServer};

use ways::protocol::wl_compositor::WlCompositor as ServerCompositor;
use ways::protocol::wl_output::WlOutput as ServerOutput;
use ways::protocol::wl_shell::WlShell as ServerShell;

use std::sync::{Arc, Mutex};

#[test]
fn simple_global() {
    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display);

    roundtrip(&mut client, &mut server).unwrap();
    let globals = manager.list();
    assert!(globals.len() == 1);
    assert_eq!(globals[0], (1, "wl_compositor".into(), 1));
}

#[test]
fn multi_versions() {
    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 4, |_, _| {});
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 3, |_, _| {});
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 2, |_, _| {});
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display);

    roundtrip(&mut client, &mut server).unwrap();
    let globals = manager.list();
    assert!(globals.len() == 4);
    let mut seen = [false; 4];
    for &(_, ref interface, version) in &globals {
        assert!(interface == "wl_compositor");
        seen[version as usize - 1] = true;
    }
    assert_eq!(seen, [true, true, true, true]);
}

#[test]
fn dynamic_global() {
    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display);

    roundtrip(&mut client, &mut server).unwrap();
    assert!(manager.list().len() == 1);

    server
        .display
        .create_global::<ServerShell, _>(&loop_token, 1, |_, _| {});

    roundtrip(&mut client, &mut server).unwrap();
    assert!(manager.list().len() == 2);

    let output = server
        .display
        .create_global::<ServerOutput, _>(&loop_token, 1, |_, _| {});

    roundtrip(&mut client, &mut server).unwrap();
    assert!(manager.list().len() == 3);

    output.destroy();

    roundtrip(&mut client, &mut server).unwrap();
    assert!(manager.list().len() == 2);
}

#[test]
fn global_manager_cb() {
    use wayc::GlobalEvent;

    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let counter = Arc::new(Mutex::new(0));
    let counter2 = counter.clone();

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new_with_cb(&client.display, move |event, _| {
        if let GlobalEvent::New { .. } = event {
            *(counter2.lock().unwrap()) += 1;
        }
    });

    roundtrip(&mut client, &mut server).unwrap();

    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    roundtrip(&mut client, &mut server).unwrap();

    assert!(*counter.lock().unwrap() == 4);
}

#[test]
fn auto_instanciate() {
    use wayc::protocol::wl_compositor::WlCompositor;
    use wayc::protocol::wl_output::WlOutput;
    use wayc::protocol::wl_shell::WlShell;
    use wayc::GlobalError;

    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 4, |_, _| {});
    server
        .display
        .create_global::<ServerShell, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display);

    roundtrip(&mut client, &mut server).unwrap();

    let compositor = manager
        .instantiate_auto::<WlCompositor, _>(|newp| newp.implement(|_, _| {}))
        .unwrap();
    assert!(compositor.version() == 4);
    let shell = manager
        .instantiate_auto::<WlShell, _>(|newp| newp.implement(|_, _| {}))
        .unwrap();
    assert!(shell.version() == 1);

    assert!(
        manager.instantiate_exact::<WlCompositor, _>(5, |newp| newp.implement(|_, _| {}))
            == Err(GlobalError::VersionTooLow(4))
    );
    assert!(
        manager.instantiate_exact::<WlOutput, _>(5, |newp| newp.implement(|_, _| {}))
            == Err(GlobalError::Missing)
    );
    assert!(
        manager.instantiate_auto::<WlOutput, _>(|newp| newp.implement(|_, _| {}))
            == Err(GlobalError::Missing)
    );
}

#[test]
#[should_panic]
fn wrong_version_create_global() {
    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 42, |_, _| {});
}

#[test]
fn wrong_global() {
    use wayc::protocol::wl_display::RequestsTrait as DisplayRequests;
    use wayc::protocol::wl_output::WlOutput;
    use wayc::protocol::wl_registry::RequestsTrait as RegistryRequests;

    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let registry = client
        .display
        .get_registry(|newp| newp.implement(|_, _| {}))
        .unwrap();

    // instanciate a wrong global, this shoudl kill the client

    registry.bind::<WlOutput, _>(2, 1, |newp| newp.implement(|_, _| {}));

    assert!(roundtrip(&mut client, &mut server).is_err());
}
<commit_msg>tests: more detailed test of client global manager<commit_after>mod helpers;

use helpers::{roundtrip, wayc, ways, TestClient, TestServer};

use ways::protocol::wl_compositor::WlCompositor as ServerCompositor;
use ways::protocol::wl_output::WlOutput as ServerOutput;
use ways::protocol::wl_shell::WlShell as ServerShell;

use std::sync::{Arc, Mutex};

#[test]
fn simple_global() {
    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display);

    roundtrip(&mut client, &mut server).unwrap();
    let globals = manager.list();
    assert!(globals.len() == 1);
    assert_eq!(globals[0], (1, "wl_compositor".into(), 1));
}

#[test]
fn multi_versions() {
    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 4, |_, _| {});
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 3, |_, _| {});
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 2, |_, _| {});
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display);

    roundtrip(&mut client, &mut server).unwrap();
    let globals = manager.list();
    assert!(globals.len() == 4);
    let mut seen = [false; 4];
    for &(_, ref interface, version) in &globals {
        assert!(interface == "wl_compositor");
        seen[version as usize - 1] = true;
    }
    assert_eq!(seen, [true, true, true, true]);
}

#[test]
fn dynamic_global() {
    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display);

    roundtrip(&mut client, &mut server).unwrap();
    assert!(manager.list().len() == 1);

    server
        .display
        .create_global::<ServerShell, _>(&loop_token, 1, |_, _| {});

    roundtrip(&mut client, &mut server).unwrap();
    assert!(manager.list().len() == 2);

    let output = server
        .display
        .create_global::<ServerOutput, _>(&loop_token, 1, |_, _| {});

    roundtrip(&mut client, &mut server).unwrap();
    assert!(manager.list().len() == 3);

    output.destroy();

    roundtrip(&mut client, &mut server).unwrap();
    assert!(manager.list().len() == 2);
}

#[test]
fn global_manager_cb() {
    use wayc::GlobalEvent;

    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let counter = Arc::new(Mutex::new(0));
    let counter2 = counter.clone();

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new_with_cb(&client.display, move |event, _| match event {
        GlobalEvent::New { .. } => *(counter2.lock().unwrap()) += 1,
        GlobalEvent::Removed { .. } => *(counter2.lock().unwrap()) -= 1,
    });

    roundtrip(&mut client, &mut server).unwrap();

    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});
    let comp = server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    roundtrip(&mut client, &mut server).unwrap();

    assert!(manager.list().len() == 4);
    assert!(*counter.lock().unwrap() == 4);

    comp.destroy();

    roundtrip(&mut client, &mut server).unwrap();

    assert!(manager.list().len() == 3);
    assert!(*counter.lock().unwrap() == 3);
}

#[test]
fn auto_instanciate() {
    use wayc::protocol::wl_compositor::WlCompositor;
    use wayc::protocol::wl_output::WlOutput;
    use wayc::protocol::wl_shell::WlShell;
    use wayc::GlobalError;

    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 4, |_, _| {});
    server
        .display
        .create_global::<ServerShell, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let manager = wayc::GlobalManager::new(&client.display);

    roundtrip(&mut client, &mut server).unwrap();

    let compositor = manager
        .instantiate_auto::<WlCompositor, _>(|newp| newp.implement(|_, _| {}))
        .unwrap();
    assert!(compositor.version() == 4);
    let shell = manager
        .instantiate_auto::<WlShell, _>(|newp| newp.implement(|_, _| {}))
        .unwrap();
    assert!(shell.version() == 1);

    assert!(
        manager.instantiate_exact::<WlCompositor, _>(5, |newp| newp.implement(|_, _| {}))
            == Err(GlobalError::VersionTooLow(4))
    );
    assert!(
        manager.instantiate_exact::<WlOutput, _>(5, |newp| newp.implement(|_, _| {}))
            == Err(GlobalError::Missing)
    );
    assert!(
        manager.instantiate_auto::<WlOutput, _>(|newp| newp.implement(|_, _| {}))
            == Err(GlobalError::Missing)
    );
}

#[test]
#[should_panic]
fn wrong_version_create_global() {
    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 42, |_, _| {});
}

#[test]
fn wrong_global() {
    use wayc::protocol::wl_display::RequestsTrait as DisplayRequests;
    use wayc::protocol::wl_output::WlOutput;
    use wayc::protocol::wl_registry::RequestsTrait as RegistryRequests;

    let mut server = TestServer::new();
    let loop_token = server.event_loop.token();
    server
        .display
        .create_global::<ServerCompositor, _>(&loop_token, 1, |_, _| {});

    let mut client = TestClient::new(&server.socket_name);
    let registry = client
        .display
        .get_registry(|newp| newp.implement(|_, _| {}))
        .unwrap();

    // instanciate a wrong global, this shoudl kill the client

    registry.bind::<WlOutput, _>(2, 1, |newp| newp.implement(|_, _| {}));

    assert!(roundtrip(&mut client, &mut server).is_err());
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::env;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // The current working directory can vary depending on how the project is being
    // built or released so we build an absolute path to the proto file
    let mut path = env::current_dir()?;
    while !path.ends_with("arrow") {
        path.pop();
    }
    path.push("format");
    path.push("Flight.proto");
    tonic_build::compile_protos(path)?;
    Ok(())
}
<commit_msg>ARROW-8366: [Rust] Revert "ARROW-7794: [Rust] Support releasing arrow-flight"<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

fn main() -> Result<(), Box<dyn std::error::Error>> {
    tonic_build::compile_protos("../../format/Flight.proto")?;
    Ok(())
}
<|endoftext|>


<commit_before>use time::Duration;
use time::PrimitiveDateTime as DateTime;

const GIGASECOND: Duration = Duration::new(1_000_000_000, 0);

// Returns a DateTime one billion seconds after start.
pub fn after(start: DateTime) -> DateTime {
    return start + GIGASECOND;
}
<commit_msg>Prefer checked add<commit_after>use time::Duration;
use time::PrimitiveDateTime as DateTime;

const GIGASECOND: Duration = Duration::new(1_000_000_000, 0);

// Returns a DateTime one billion seconds after start.
pub fn after(start: DateTime) -> DateTime {
    return start.checked_add(GIGASECOND).expect("overflow occured");
}
<|endoftext|>


<commit_before>use std::{collections::HashMap, str::FromStr};

use sqlx::{
    any::{AnyConnectOptions, AnyPool, AnyPoolOptions, AnyRow},
    Column, ConnectOptions, Row, TypeInfo,
};

use kernel::{
    common::{
        async_trait::async_trait,
        eyre::{bail, Result},
        serde::Serialize,
        tracing::{self, log::LevelFilter},
    },
    stencila_schema::{
        ArrayValidator, BooleanValidator, CodeError, Datatable, DatatableColumn, IntegerValidator,
        Node, Null, Number, NumberValidator, StringValidator, ValidatorTypes,
    },
    Kernel, KernelSelector, KernelStatus, KernelTrait, KernelType, Task, TaskResult,
};

/// A kernel that executes SQL
#[derive(Debug, Default, Serialize)]
#[serde(crate = "kernel::common::serde")]
pub struct SqlKernel {
    url: String,

    //
    #[serde(skip)]
    connection_pool: Option<AnyPool>,

    assigned_selects: HashMap<String, Datatable>,
}

impl SqlKernel {
    pub fn new(selector: &KernelSelector) -> Self {
        let url = selector
            .config
            .clone()
            .unwrap_or_else(|| "sqlite://:memory:".to_string());

        tracing::trace!("Creating new SqlKernel with url: {}", url);

        // TODO: fallback to using $DATABASE_URL

        Self {
            url,
            connection_pool: None,
            assigned_selects: HashMap::new(),
        }
    }

    async fn connect(&mut self) -> Result<()> {
        if let Some(_pool) = &self.connection_pool {
            return Ok(());
        }

        let mut options = AnyConnectOptions::from_str(&self.url)?;
        if let Some(options) = options.as_mssql_mut() {
            options.log_statements(LevelFilter::Trace);
        }
        if let Some(options) = options.as_postgres_mut() {
            options.log_statements(LevelFilter::Trace);
        }
        if let Some(options) = options.as_sqlite_mut() {
            options.log_statements(LevelFilter::Trace);
        }

        let pool = AnyPoolOptions::new()
            .max_connections(10)
            .connect_with(options)
            .await?;
        self.connection_pool = Some(pool);

        Ok(())
    }
}

#[async_trait]
impl KernelTrait for SqlKernel {
    async fn spec(&self) -> Kernel {
        Kernel::new("sql", KernelType::Builtin, &["sql"], true, false, false)
    }

    async fn status(&self) -> Result<KernelStatus> {
        Ok(KernelStatus::Idle)
    }

    async fn get(&mut self, name: &str) -> Result<Node> {
        let _pool = self.connect().await?;
        let pool = self.connection_pool.as_ref().unwrap();
        let sql = format!("SELECT * FROM \"{}\"", name.replace('"', "-"));
        match sqlx::query(&sql).fetch_all(pool).await {
            Ok(rows) => Ok(any_rows_to_datatable(rows)),
            Err(error) => {
                bail!("No such symbol: {}", error)
            }
        }
    }

    async fn set(&mut self, _name: &str, _value: Node) -> Result<()> {
        todo!("Create a new table with name")
    }

    async fn exec_sync(&mut self, code: &str) -> Result<Task> {
        let mut task = Task::begin_sync();
        let mut outputs = Vec::new();
        let mut messages = Vec::new();

        match self.connect().await {
            Ok(..) => {
                let pool = self.connection_pool.as_ref().unwrap();
                match sqlx::query(code).fetch_all(pool).await {
                    Ok(rows) => {
                        let datatable = any_rows_to_datatable(rows);
                        outputs.push(datatable);
                    }
                    Err(error) => {
                        let message = error.to_string();
                        let message = message
                            .strip_prefix("error returned from database:")
                            .unwrap_or(&message);
                        let message = message.trim().to_string();

                        messages.push(CodeError {
                            error_message: message,
                            ..Default::default()
                        });
                    }
                }
            }
            Err(error) => {
                messages.push(CodeError {
                    error_message: error.to_string(),
                    ..Default::default()
                });
            }
        }

        task.end(TaskResult::new(outputs, messages));

        Ok(task)
    }
}

/// Transform a vector of `sqlx` [`AnyRow`]s to a Stencila [`Datatable`]
fn any_rows_to_datatable(rows: Vec<AnyRow>) -> Node {
    let rows_len = rows.len();
    let columns = if let Some(row) = rows.first() {
        row.columns()
            .iter()
            .map(|column| {
                let name = column.name().to_string();
                let type_name = column.type_info().name();
                let validator = match type_name {
                    "BOOLEAN" => {
                        Some(ValidatorTypes::BooleanValidator(BooleanValidator::default()))
                    }
                    "INTEGER" => {
                        Some(ValidatorTypes::IntegerValidator(IntegerValidator::default()))
                    }
                    "REAL" => Some(ValidatorTypes::NumberValidator(NumberValidator::default())),
                    "TEXT" => Some(ValidatorTypes::StringValidator(StringValidator::default())),
                    "BLOB" | "NULL" => None,
                    _ => {
                        tracing::debug!("Unhandled column type: {}", type_name);
                        None
                    }
                };
                (name, validator)
            })
            .collect()
    } else {
        Vec::new()
    };

    // Pre-allocate an vector of the size needed to hold all values and insert them in
    // column-first order
    let mut values: Vec<Node> = vec![Node::Null(Null {}); columns.len() * rows.len()];
    for (row_index, row) in rows.into_iter().enumerate() {
        for (col_index, (_name, validator)) in columns.iter().enumerate() {
            let position = col_index * rows_len + row_index;
            let value = match validator {
                Some(ValidatorTypes::BooleanValidator(..)) => {
                    row.try_get::<bool, usize>(col_index).map(Node::Boolean)
                }
                Some(ValidatorTypes::IntegerValidator(..)) => {
                    row.try_get::<i64, usize>(col_index).map(Node::Integer)
                }
                Some(ValidatorTypes::NumberValidator(..)) => row
                    .try_get::<f64, usize>(col_index)
                    .map(|num| Node::Number(Number(num))),
                Some(ValidatorTypes::StringValidator(..)) => {
                    row.try_get::<String, usize>(col_index).map(Node::String)
                }
                _ => row.try_get::<String, usize>(col_index).map(Node::String),
            };
            if let Ok(value) = value {
                values[position] = value;
            }
        }
    }

    let columns = columns
        .into_iter()
        .map(|(name, validator)| DatatableColumn {
            name,
            validator: validator.map(|validator| {
                Box::new(ArrayValidator {
                    items_validator: Some(Box::new(validator)),
                    ..Default::default()
                })
            }),
            values: values.drain(..rows_len).collect(),
            ..Default::default()
        })
        .collect();

    Node::Datatable(Datatable {
        columns,
        ..Default::default()
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use kernel::{common::tokio, stencila_schema::Number, KernelTrait};
    use test_utils::{assert_json_eq, common::serde_json::json};

    #[tokio::test]
    async fn get_set_exec() -> Result<()> {
        let mut kernel = SqlKernel::new(&KernelSelector::default());

        match kernel.get("a").await {
            Ok(..) => bail!("Expected an error"),
            Err(error) => assert!(error.to_string().contains("does not exist")),
        };

        match kernel.set("a", Node::String("A".to_string())).await {
            Ok(..) => bail!("Expected an error"),
            Err(error) => assert!(error
                .to_string()
                .contains("Unable to convert string `A` to a number")),
        };

        kernel.set("a", Node::Number(Number(1.23))).await?;

        let a = kernel.get("a").await?;
        assert!(matches!(a, Node::Number(..)));
        assert_json_eq!(a, json!(1.23));

        let (outputs, errors) = kernel.exec("a * 2").await?;
        assert_json_eq!(outputs, json!([2.46]));
        assert_eq!(errors.len(), 0);

        let (outputs, errors) = kernel.exec("x * 2").await?;
        assert_eq!(outputs.len(), 0);
        assert_json_eq!(
            errors,
            json!([{"type": "CodeError", "errorMessage": "Undefined variable or function: x"}])
        );

        Ok(())
    }
}
<commit_msg>fix(SQL kernel): Make SQLite db file path absolute if necessary<commit_after>use std::{
    collections::HashMap,
    path::{Path, PathBuf},
    str::FromStr,
};

use sqlx::{
    any::{AnyConnectOptions, AnyPool, AnyPoolOptions, AnyRow},
    Column, ConnectOptions, Row, TypeInfo,
};

use kernel::{
    common::{
        async_trait::async_trait,
        eyre::{bail, Result},
        serde::Serialize,
        tracing::{self, log::LevelFilter},
    },
    stencila_schema::{
        ArrayValidator, BooleanValidator, CodeError, Datatable, DatatableColumn, IntegerValidator,
        Node, Null, Number, NumberValidator, StringValidator, ValidatorTypes,
    },
    Kernel, KernelSelector, KernelStatus, KernelTrait, KernelType, Task, TaskResult,
};

/// A kernel that executes SQL
#[derive(Debug, Default, Serialize)]
#[serde(crate = "kernel::common::serde")]
pub struct SqlKernel {
    /// The kernel configuration containing the database URL
    config: Option<String>,

    /// The directory that the kernel is started in
    ///
    /// Used to be able to resolve the path to SQLite files with relative paths
    directory: Option<PathBuf>,

    /// The kernel's database connection pool
    #[serde(skip)]
    connection_pool: Option<AnyPool>,

    /// Any select queries that have been assigned to variables
    #[serde(skip)]
    assigned_selects: HashMap<String, Datatable>,
}

impl SqlKernel {
    /// Create a new `SqlKernel`
    pub fn new(selector: &KernelSelector) -> Self {
        Self {
            config: selector.config.clone(),
            ..Default::default()
        }
    }

    /// Connect to the database
    ///
    /// This is called just-in-time (called by `exec`, `get` etc) as well as in `start`
    /// so that any errors during connection can be surfaced to the user as code chunk errors.
    async fn connect(&mut self) -> Result<()> {
        if let Some(_pool) = &self.connection_pool {
            return Ok(());
        }

        // Resolve the database URL, falling back to env var and then to in-memory SQLite
        let mut url = self
            .config
            .clone()
            // TODO: fallback to using $DATABASE_URL
            .unwrap_or_else(|| "sqlite://:memory:".to_string());

        // If the URL is for a SQLite and the file path is relative then make it
        // absolute using the directory that the kernel was started in.
        if let (Some(spec), Some(directory)) = (url.strip_prefix("sqlite://"), &self.directory) {
            let path = PathBuf::from(spec);
            if spec != ":memory:" && path.is_relative() {
                url = ["sqlite://", &directory.join(path).to_string_lossy()].concat();
            }
        }

        tracing::trace!("Connecting to database: {}", url);

        let mut options = AnyConnectOptions::from_str(&url)?;
        if let Some(options) = options.as_mssql_mut() {
            options.log_statements(LevelFilter::Trace);
        }
        if let Some(options) = options.as_postgres_mut() {
            options.log_statements(LevelFilter::Trace);
        }
        if let Some(options) = options.as_sqlite_mut() {
            options.log_statements(LevelFilter::Trace);
        }

        let pool = AnyPoolOptions::new()
            .max_connections(10)
            .connect_with(options)
            .await?;
        self.connection_pool = Some(pool);

        Ok(())
    }
}

#[async_trait]
impl KernelTrait for SqlKernel {
    async fn spec(&self) -> Kernel {
        Kernel::new("sql", KernelType::Builtin, &["sql"], true, false, false)
    }

    async fn status(&self) -> Result<KernelStatus> {
        Ok(KernelStatus::Idle)
    }

    async fn start(&mut self, directory: &Path) -> Result<()> {
        self.directory = Some(directory.to_owned());

        // Log error but do not return error so that another attempt is made to
        // re-connect is made in `exec` etc
        if let Err(error) = self.connect().await {
            tracing::error!("While attempting to connect to database: {}", error)
        }

        Ok(())
    }

    async fn get(&mut self, name: &str) -> Result<Node> {
        let _pool = self.connect().await?;
        let pool = self.connection_pool.as_ref().unwrap();
        let sql = format!("SELECT * FROM \"{}\"", name.replace('"', "-"));
        match sqlx::query(&sql).fetch_all(pool).await {
            Ok(rows) => Ok(any_rows_to_datatable(rows)),
            Err(error) => {
                bail!("No such symbol: {}", error)
            }
        }
    }

    async fn set(&mut self, _name: &str, _value: Node) -> Result<()> {
        todo!("Create a new table with name")
    }

    async fn exec_sync(&mut self, code: &str) -> Result<Task> {
        let mut task = Task::begin_sync();
        let mut outputs = Vec::new();
        let mut messages = Vec::new();

        match self.connect().await {
            Ok(..) => {
                let pool = self.connection_pool.as_ref().unwrap();
                match sqlx::query(code).fetch_all(pool).await {
                    Ok(rows) => {
                        let datatable = any_rows_to_datatable(rows);
                        outputs.push(datatable);
                    }
                    Err(error) => {
                        let message = error.to_string();
                        let message = message
                            .strip_prefix("error returned from database:")
                            .unwrap_or(&message);
                        let message = message.trim().to_string();

                        messages.push(CodeError {
                            error_message: message,
                            ..Default::default()
                        });
                    }
                }
            }
            Err(error) => {
                messages.push(CodeError {
                    error_message: error.to_string(),
                    ..Default::default()
                });
            }
        }

        task.end(TaskResult::new(outputs, messages));

        Ok(task)
    }
}

/// Transform a vector of `sqlx` [`AnyRow`]s to a Stencila [`Datatable`]
fn any_rows_to_datatable(rows: Vec<AnyRow>) -> Node {
    let rows_len = rows.len();
    let columns = if let Some(row) = rows.first() {
        row.columns()
            .iter()
            .map(|column| {
                let name = column.name().to_string();
                let type_name = column.type_info().name();
                let validator = match type_name {
                    "BOOLEAN" => {
                        Some(ValidatorTypes::BooleanValidator(BooleanValidator::default()))
                    }
                    "INTEGER" => {
                        Some(ValidatorTypes::IntegerValidator(IntegerValidator::default()))
                    }
                    "REAL" => Some(ValidatorTypes::NumberValidator(NumberValidator::default())),
                    "TEXT" => Some(ValidatorTypes::StringValidator(StringValidator::default())),
                    "BLOB" | "NULL" => None,
                    _ => {
                        tracing::debug!("Unhandled column type: {}", type_name);
                        None
                    }
                };
                (name, validator)
            })
            .collect()
    } else {
        Vec::new()
    };

    // Pre-allocate an vector of the size needed to hold all values and insert them in
    // column-first order
    let mut values: Vec<Node> = vec![Node::Null(Null {}); columns.len() * rows.len()];
    for (row_index, row) in rows.into_iter().enumerate() {
        for (col_index, (_name, validator)) in columns.iter().enumerate() {
            let position = col_index * rows_len + row_index;
            let value = match validator {
                Some(ValidatorTypes::BooleanValidator(..)) => {
                    row.try_get::<bool, usize>(col_index).map(Node::Boolean)
                }
                Some(ValidatorTypes::IntegerValidator(..)) => {
                    row.try_get::<i64, usize>(col_index).map(Node::Integer)
                }
                Some(ValidatorTypes::NumberValidator(..)) => row
                    .try_get::<f64, usize>(col_index)
                    .map(|num| Node::Number(Number(num))),
                Some(ValidatorTypes::StringValidator(..)) => {
                    row.try_get::<String, usize>(col_index).map(Node::String)
                }
                _ => row.try_get::<String, usize>(col_index).map(Node::String),
            };
            if let Ok(value) = value {
                values[position] = value;
            }
        }
    }

    let columns = columns
        .into_iter()
        .map(|(name, validator)| DatatableColumn {
            name,
            validator: validator.map(|validator| {
                Box::new(ArrayValidator {
                    items_validator: Some(Box::new(validator)),
                    ..Default::default()
                })
            }),
            values: values.drain(..rows_len).collect(),
            ..Default::default()
        })
        .collect();

    Node::Datatable(Datatable {
        columns,
        ..Default::default()
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use kernel::{common::tokio, stencila_schema::Number, KernelTrait};
    use test_utils::{assert_json_eq, common::serde_json::json};

    #[tokio::test]
    async fn get_set_exec() -> Result<()> {
        let mut kernel = SqlKernel::new(&KernelSelector::default());

        match kernel.get("a").await {
            Ok(..) => bail!("Expected an error"),
            Err(error) => assert!(error.to_string().contains("does not exist")),
        };

        match kernel.set("a", Node::String("A".to_string())).await {
            Ok(..) => bail!("Expected an error"),
            Err(error) => assert!(error
                .to_string()
                .contains("Unable to convert string `A` to a number")),
        };

        kernel.set("a", Node::Number(Number(1.23))).await?;

        let a = kernel.get("a").await?;
        assert!(matches!(a, Node::Number(..)));
        assert_json_eq!(a, json!(1.23));

        let (outputs, errors) = kernel.exec("a * 2").await?;
        assert_json_eq!(outputs, json!([2.46]));
        assert_eq!(errors.len(), 0);

        let (outputs, errors) = kernel.exec("x * 2").await?;
        assert_eq!(outputs.len(), 0);
        assert_json_eq!(
            errors,
            json!([{"type": "CodeError", "errorMessage": "Undefined variable or function: x"}])
        );

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


<commit_before>#[macro_use]
extern crate objc;
#[macro_use]
extern crate objc_foundation;

use objc::{ClassDecl, Id};
use objc::runtime::Object;
use objc_foundation::{class, INSObject, NSObject};

object_struct!(MYObject);

impl MYObject {
    fn number(&self) -> uint {
        let obj = unsafe {
            &*(self as *const _ as *const Object)
        };
        unsafe {
            *obj.get_ivar::<uint>("_number")
        }
    }

    fn set_number(&mut self, number: uint) {
        let obj = unsafe {
            &mut *(self as *mut _ as *mut Object)
        };
        unsafe {
            obj.set_ivar("_number", number);
        }
    }

    fn register() {
        let superclass = class::<NSObject>();
        let mut decl = ClassDecl::new(superclass, "MYObject").unwrap();

        decl.add_ivar::<uint>("_number");
        decl.add_method(method!(
            (&mut MYObject)this, setNumber:(uint)number; {
                this.set_number(number);
            }
        ));
        decl.add_method(method!(
            (&MYObject)this, number -> uint, {
                this.number()
            }
        ));

        decl.register();
    }
}

fn main() {
    MYObject::register();
    let mut obj: Id<MYObject> = INSObject::new();
    println!("{:?}", obj);

    obj.set_number(7);
    println!("Number: {}", obj.number());

    unsafe {
        msg_send![obj, setNumber:12u];
    }
    println!("Number: {}", obj.number());
}
<commit_msg>Replaced uint with u32 in custom_class example.<commit_after>#[macro_use]
extern crate objc;
#[macro_use]
extern crate objc_foundation;

use objc::{ClassDecl, Id};
use objc::runtime::Object;
use objc_foundation::{class, INSObject, NSObject};

object_struct!(MYObject);

impl MYObject {
    fn number(&self) -> u32 {
        let obj = unsafe {
            &*(self as *const _ as *const Object)
        };
        unsafe {
            *obj.get_ivar::<u32>("_number")
        }
    }

    fn set_number(&mut self, number: u32) {
        let obj = unsafe {
            &mut *(self as *mut _ as *mut Object)
        };
        unsafe {
            obj.set_ivar("_number", number);
        }
    }

    fn register() {
        let superclass = class::<NSObject>();
        let mut decl = ClassDecl::new(superclass, "MYObject").unwrap();

        decl.add_ivar::<u32>("_number");
        decl.add_method(method!(
            (&mut MYObject)this, setNumber:(u32)number; {
                this.set_number(number);
            }
        ));
        decl.add_method(method!(
            (&MYObject)this, number -> u32, {
                this.number()
            }
        ));

        decl.register();
    }
}

fn main() {
    MYObject::register();
    let mut obj: Id<MYObject> = INSObject::new();
    println!("{:?}", obj);

    obj.set_number(7);
    println!("Number: {}", obj.number());

    unsafe {
        msg_send![obj, setNumber:12u32];
    }
    println!("Number: {}", obj.number());
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use rustc_serialize::json::Json;

use analysis::Analyzer;
use Value;

#[derive(Debug, PartialEq)]
pub enum FieldType {
    String,
    Binary,
    Number{size: u8, is_float: bool},
    Boolean,
    Date,
}

impl Default for FieldType {
    fn default() -> FieldType { FieldType::String }
}

#[derive(Debug)]
pub struct FieldMapping {
    data_type: FieldType,
    is_stored: bool,
    pub is_in_all: bool,
    boost: f64,
    analyzer: Analyzer,
}

impl Default for FieldMapping {
    fn default() -> FieldMapping {
        FieldMapping {
            data_type: FieldType::default(),
            is_stored: false,
            is_in_all: true,
            boost: 1.0f64,
            analyzer: Analyzer::Standard,
        }
    }
}

impl FieldMapping {
    pub fn process_value(&self, value: Json) -> Option<Value> {
        if value == Json::Null {
            return Some(Value::Null);
        }

        match self.data_type {
            FieldType::String => {
                match value {
                    Json::String(string) => {
                        // Analyzed strings become TSVectors. Unanalyzed strings become... strings
                        if self.analyzer == Analyzer::None {
                            Some(Value::String(string))
                        } else {
                            let tokens = self.analyzer.run(string);
                            Some(Value::TSVector(tokens))
                        }
                    }
                    Json::Array(array) => {
                        // Pack any strings into a vec, ignore nulls. Quit if we see anything else
                        let mut strings = Vec::new();

                        for item in array {
                            match item {
                                Json::String(string) => strings.push(string),
                                Json::Null => {}
                                _ => {
                                    return None;
                                }
                            }
                        }

                        self.process_value(Json::String(strings.join(" ")))
                    }
                    _ => None
                }
            }
            FieldType::Boolean => {
                Some(Value::Boolean(parse_boolean(&value)))
            }
            _ => None
        }
    }
}

#[derive(Debug)]
pub struct Mapping {
    pub fields: HashMap<String, FieldMapping>,
}

impl Mapping {
    pub fn from_json(json: &Json) -> Mapping {
        let json = json.as_object().unwrap();
        let properties_json = json.get("properties").unwrap().as_object().unwrap();

        // Parse fields
        let mut fields = HashMap::new();
        for (field_name, field_mapping_json) in properties_json.iter() {
            fields.insert(field_name.clone(), FieldMapping::from_json(field_mapping_json));
        }

        Mapping {
            fields: fields,
        }
    }
}

fn parse_boolean(json: &Json) -> bool {
    match *json {
        Json::Boolean(val) => val,
        Json::String(ref s) => {
            match s.as_ref() {
                "yes" => true,
                "no" => false,
                _ => {
                    warn!("bad boolean value {:?}", s);
                    false
                }
            }
        }
        _ => {
            // TODO: Raise error
            warn!("bad boolean value {:?}", json);
            false
        }
    }
}

impl FieldMapping {
    pub fn from_json(json: &Json) -> FieldMapping {
        let json = json.as_object().unwrap();
        let mut field_mapping = FieldMapping::default();

        for (key, value) in json.iter() {
            match key.as_ref() {
                "type" => {
                    let type_name = value.as_string().unwrap();

                    field_mapping.data_type = match type_name.as_ref() {
                        "string" => FieldType::String,
                        "integer" => FieldType::Number{size: 64, is_float: false},
                        "boolean" => FieldType::Boolean,
                        "date" => FieldType::Date,
                        _ => {
                            // TODO; make this an error
                            warn!("unimplemented type name! {}", type_name);
                            FieldType::default()
                        }
                    };
                }
                "index" => {
                    let index = value.as_string().unwrap();
                    if index == "not_analyzed" {
                        field_mapping.analyzer = Analyzer::None;
                    } else {
                        // TODO: Implement other variants and make this an error
                        warn!("unimplemented index setting! {}", index);
                    }
                }
                "index_analyzer" => {
                    if let Some(ref s) = value.as_string() {
                        if s == &"edgengram_analyzer" {
                            field_mapping.analyzer = Analyzer::EdgeNGram;
                        }
                    }
                }
                "boost" => {
                    field_mapping.boost = value.as_f64().unwrap();
                }
                "store" => {
                    field_mapping.is_stored = parse_boolean(value);
                }
                "include_in_all" => {
                    field_mapping.is_in_all = parse_boolean(value);
                }
                _ => warn!("unimplemented field mapping key! {}", key)
            }

        }

        field_mapping
    }
}
<commit_msg>Added ability to parse numbers<commit_after>use std::collections::HashMap;

use rustc_serialize::json::Json;

use analysis::Analyzer;
use Value;

#[derive(Debug, PartialEq)]
pub enum FieldType {
    String,
    Binary,
    Number{size: u8, is_float: bool},
    Boolean,
    Date,
}

impl Default for FieldType {
    fn default() -> FieldType { FieldType::String }
}

#[derive(Debug)]
pub struct FieldMapping {
    data_type: FieldType,
    is_stored: bool,
    pub is_in_all: bool,
    boost: f64,
    analyzer: Analyzer,
}

impl Default for FieldMapping {
    fn default() -> FieldMapping {
        FieldMapping {
            data_type: FieldType::default(),
            is_stored: false,
            is_in_all: true,
            boost: 1.0f64,
            analyzer: Analyzer::Standard,
        }
    }
}

impl FieldMapping {
    pub fn process_value(&self, value: Json) -> Option<Value> {
        if value == Json::Null {
            return Some(Value::Null);
        }

        match self.data_type {
            FieldType::String => {
                match value {
                    Json::String(string) => {
                        // Analyzed strings become TSVectors. Unanalyzed strings become... strings
                        if self.analyzer == Analyzer::None {
                            Some(Value::String(string))
                        } else {
                            let tokens = self.analyzer.run(string);
                            Some(Value::TSVector(tokens))
                        }
                    }
                    Json::I64(num) => self.process_value(Json::String(num.to_string())),
                    Json::U64(num) => self.process_value(Json::String(num.to_string())),
                    Json::F64(num) => self.process_value(Json::String(num.to_string())),
                    Json::Array(array) => {
                        // Pack any strings into a vec, ignore nulls. Quit if we see anything else
                        let mut strings = Vec::new();

                        for item in array {
                            match item {
                                Json::String(string) => strings.push(string),
                                Json::Null => {}
                                _ => {
                                    return None;
                                }
                            }
                        }

                        self.process_value(Json::String(strings.join(" ")))
                    }
                    _ => None
                }
            }
            FieldType::Number{size, is_float} => {
                match value {
                    // TODO check the numbers fit in "size"
                    Json::U64(num) => Some(Value::U64(num)),
                    Json::I64(num) => Some(Value::I64(num)),
                    Json::F64(num) => {
                        if !is_float {
                            return None;
                        }

                        Some(Value::F64(num))
                    }
                    _ => None
                }
            }
            FieldType::Boolean => {
                Some(Value::Boolean(parse_boolean(&value)))
            }
            _ => None
        }
    }
}

#[derive(Debug)]
pub struct Mapping {
    pub fields: HashMap<String, FieldMapping>,
}

impl Mapping {
    pub fn from_json(json: &Json) -> Mapping {
        let json = json.as_object().unwrap();
        let properties_json = json.get("properties").unwrap().as_object().unwrap();

        // Parse fields
        let mut fields = HashMap::new();
        for (field_name, field_mapping_json) in properties_json.iter() {
            fields.insert(field_name.clone(), FieldMapping::from_json(field_mapping_json));
        }

        Mapping {
            fields: fields,
        }
    }
}

fn parse_boolean(json: &Json) -> bool {
    match *json {
        Json::Boolean(val) => val,
        Json::String(ref s) => {
            match s.as_ref() {
                "yes" => true,
                "no" => false,
                _ => {
                    warn!("bad boolean value {:?}", s);
                    false
                }
            }
        }
        _ => {
            // TODO: Raise error
            warn!("bad boolean value {:?}", json);
            false
        }
    }
}

impl FieldMapping {
    pub fn from_json(json: &Json) -> FieldMapping {
        let json = json.as_object().unwrap();
        let mut field_mapping = FieldMapping::default();

        for (key, value) in json.iter() {
            match key.as_ref() {
                "type" => {
                    let type_name = value.as_string().unwrap();

                    field_mapping.data_type = match type_name.as_ref() {
                        "string" => FieldType::String,
                        "integer" => FieldType::Number{size: 64, is_float: false},
                        "boolean" => FieldType::Boolean,
                        "date" => FieldType::Date,
                        _ => {
                            // TODO; make this an error
                            warn!("unimplemented type name! {}", type_name);
                            FieldType::default()
                        }
                    };
                }
                "index" => {
                    let index = value.as_string().unwrap();
                    if index == "not_analyzed" {
                        field_mapping.analyzer = Analyzer::None;
                    } else {
                        // TODO: Implement other variants and make this an error
                        warn!("unimplemented index setting! {}", index);
                    }
                }
                "index_analyzer" => {
                    if let Some(ref s) = value.as_string() {
                        if s == &"edgengram_analyzer" {
                            field_mapping.analyzer = Analyzer::EdgeNGram;
                        }
                    }
                }
                "boost" => {
                    field_mapping.boost = value.as_f64().unwrap();
                }
                "store" => {
                    field_mapping.is_stored = parse_boolean(value);
                }
                "include_in_all" => {
                    field_mapping.is_in_all = parse_boolean(value);
                }
                _ => warn!("unimplemented field mapping key! {}", key)
            }

        }

        field_mapping
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Masaki Hara
//
// This software is released under the MIT License.
// http://opensource.org/licenses/mit-license.php

extern crate byteorder;

use self::byteorder::{LittleEndian, ByteOrder};
use std::cmp;
use std::io::{Write, Result};

const S : [u32; 64] = [
    7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22,
    5,  9, 14, 20, 5,  9, 14, 20, 5,  9, 14, 20, 5,  9, 14, 20,
    4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23,
    6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21,
];
const K : [u32; 64] = [
    0xD76AA478, 0xE8C7B756, 0x242070DB, 0xC1BDCEEE,
    0xF57C0FAF, 0x4787C62A, 0xA8304613, 0xFD469501,
    0x698098D8, 0x8B44F7AF, 0xFFFF5BB1, 0x895CD7BE,
    0x6B901122, 0xFD987193, 0xA679438E, 0x49B40821,
    0xF61E2562, 0xC040B340, 0x265E5A51, 0xE9B6C7AA,
    0xD62F105D, 0x02441453, 0xD8A1E681, 0xE7D3FBC8,
    0x21E1CDE6, 0xC33707D6, 0xF4D50D87, 0x455A14ED,
    0xA9E3E905, 0xFCEFA3F8, 0x676F02D9, 0x8D2A4C8A,
    0xFFFA3942, 0x8771F681, 0x6D9D6122, 0xFDE5380C,
    0xA4BEEA44, 0x4BDECFA9, 0xF6BB4B60, 0xBEBFBC70,
    0x289B7EC6, 0xEAA127FA, 0xD4EF3085, 0x04881D05,
    0xD9D4D039, 0xE6DB99E5, 0x1FA27CF8, 0xC4AC5665,
    0xF4292244, 0x432AFF97, 0xAB9423A7, 0xFC93A039,
    0x655B59C3, 0x8F0CCC92, 0xFFEFF47D, 0x85845DD1,
    0x6FA87E4F, 0xFE2CE6E0, 0xA3014314, 0x4E0811A1,
    0xF7537E82, 0xBD3AF235, 0x2AD7D2BB, 0xEB86D391,
];

fn update(hh: &mut [u32; 4], chunk: &[u8; 64]) {
    let mut m : [u32; 16] = [0; 16];
    for i in 0..16 {
        m[i] = LittleEndian::read_u32(&chunk[i * 4 .. i * 4 + 4]);
    }
    let mut a : u32 = hh[0];
    let mut b : u32 = hh[1];
    let mut c : u32 = hh[2];
    let mut d : u32 = hh[3];
    for i in 0..64 {
        let f : u32;
        let g : usize;
        if i < 16 {
            f = (b & c) | (!b & d);
            g = i;
        } else if i < 32 {
            f = (d & b) | (!d & c);
            g = (5 * i + 1) % 16;
        } else if i < 48 {
            f = b ^ c ^ d;
            g = (3 * i + 5) % 16;
        } else {
            f = c ^ (b | !d);
            g = 7 * i % 16;
        }
        let tmp = d;
        d = c;
        c = b;
        b = b.wrapping_add(
            (a.wrapping_add(f)
             .wrapping_add(K[i]).wrapping_add(m[g])).rotate_left(S[i]));
        a = tmp;
    }
    hh[0] = hh[0].wrapping_add(a);
    hh[1] = hh[1].wrapping_add(b);
    hh[2] = hh[2].wrapping_add(c);
    hh[3] = hh[3].wrapping_add(d);
}

pub struct MD5Writer {
    state: [u32; 4],
    position: u64,
    buf: [u8; 64],
    bufpos: usize,
}

impl Copy for MD5Writer {}
impl Clone for MD5Writer {
    fn clone(&self) -> Self {
        return MD5Writer {
            state: self.state,
            position: self.position,
            buf: self.buf,
            bufpos: self.bufpos
        };
    }
}

impl MD5Writer {
    pub fn new() -> MD5Writer {
        return MD5Writer {
            state: [0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476],
            position: 0,
            buf: [0; 64],
            bufpos: 0,
        };
    }
    pub fn sum(&self) -> [u8; 16] {
        let mut state = self.state;
        let mut buf = self.buf;
        let buflen = buf.len();
        let bufpos = self.bufpos;
        buf[self.bufpos] = 0x80;
        for i in bufpos + 1 .. buflen { buf[i] = 0; }
        if bufpos + 1 > buflen - 8 {
            update(&mut state, &buf);
            for i in 0 .. buflen { buf[i] = 0; }
        }
        LittleEndian::write_u64(&mut buf[buflen - 8 ..], self.position);
        update(&mut state, &buf);
        let mut hashbuf : [u8; 16] = [0; 16];
        for i in 0..4 {
            LittleEndian::write_u32(
                &mut hashbuf[i * 4 .. i * 4 + 4], state[i]);
        }
        return hashbuf;
    }
}

impl Write for MD5Writer {
    fn write(&mut self, mut buf: &[u8]) -> Result<usize> {
        let orig_len = buf.len();
        while buf.len() > 0 {
            let wrsize = cmp::min(self.buf.len() - self.bufpos, buf.len());
            self.buf[self.bufpos .. self.bufpos + wrsize]
                .clone_from_slice(&buf[0 .. wrsize]);
            buf = &buf[wrsize ..];
            self.position += (wrsize * 8) as u64;
            self.bufpos += wrsize;
            if self.bufpos >= self.buf.len() {
                update(&mut self.state, &self.buf);
                self.bufpos -= self.buf.len();
            }
        }
        return Ok(orig_len);
    }
    fn flush(&mut self) -> Result<()> {
        return Ok(());
    }
}
<commit_msg>Small fix of MD5.<commit_after>// Copyright (c) 2016 Masaki Hara
//
// This software is released under the MIT License.
// http://opensource.org/licenses/mit-license.php

extern crate byteorder;

use self::byteorder::{LittleEndian, ByteOrder};
use std::cmp;
use std::io::{Write, Result};

const S : [u32; 64] = [
    7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22, 7, 12, 17, 22,
    5,  9, 14, 20, 5,  9, 14, 20, 5,  9, 14, 20, 5,  9, 14, 20,
    4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23, 4, 11, 16, 23,
    6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21, 6, 10, 15, 21,
];
const K : [u32; 64] = [
    0xD76AA478, 0xE8C7B756, 0x242070DB, 0xC1BDCEEE,
    0xF57C0FAF, 0x4787C62A, 0xA8304613, 0xFD469501,
    0x698098D8, 0x8B44F7AF, 0xFFFF5BB1, 0x895CD7BE,
    0x6B901122, 0xFD987193, 0xA679438E, 0x49B40821,
    0xF61E2562, 0xC040B340, 0x265E5A51, 0xE9B6C7AA,
    0xD62F105D, 0x02441453, 0xD8A1E681, 0xE7D3FBC8,
    0x21E1CDE6, 0xC33707D6, 0xF4D50D87, 0x455A14ED,
    0xA9E3E905, 0xFCEFA3F8, 0x676F02D9, 0x8D2A4C8A,
    0xFFFA3942, 0x8771F681, 0x6D9D6122, 0xFDE5380C,
    0xA4BEEA44, 0x4BDECFA9, 0xF6BB4B60, 0xBEBFBC70,
    0x289B7EC6, 0xEAA127FA, 0xD4EF3085, 0x04881D05,
    0xD9D4D039, 0xE6DB99E5, 0x1FA27CF8, 0xC4AC5665,
    0xF4292244, 0x432AFF97, 0xAB9423A7, 0xFC93A039,
    0x655B59C3, 0x8F0CCC92, 0xFFEFF47D, 0x85845DD1,
    0x6FA87E4F, 0xFE2CE6E0, 0xA3014314, 0x4E0811A1,
    0xF7537E82, 0xBD3AF235, 0x2AD7D2BB, 0xEB86D391,
];

fn update(hh: &mut [u32; 4], chunk: &[u8; 64]) {
    let mut m : [u32; 16] = [0; 16];
    for i in 0..16 {
        m[i] = LittleEndian::read_u32(&chunk[i * 4 .. i * 4 + 4]);
    }
    let mut a : u32 = hh[0];
    let mut b : u32 = hh[1];
    let mut c : u32 = hh[2];
    let mut d : u32 = hh[3];
    for i in 0..64 {
        let f : u32;
        let g : usize;
        if i < 16 {
            f = (b & c) | (!b & d);
            g = i;
        } else if i < 32 {
            f = (d & b) | (!d & c);
            g = (5 * i + 1) % 16;
        } else if i < 48 {
            f = b ^ c ^ d;
            g = (3 * i + 5) % 16;
        } else {
            f = c ^ (b | !d);
            g = 7 * i % 16;
        }
        let tmp = d;
        d = c;
        c = b;
        b = b.wrapping_add(
            (a.wrapping_add(f)
             .wrapping_add(K[i]).wrapping_add(m[g])).rotate_left(S[i]));
        a = tmp;
    }
    hh[0] = hh[0].wrapping_add(a);
    hh[1] = hh[1].wrapping_add(b);
    hh[2] = hh[2].wrapping_add(c);
    hh[3] = hh[3].wrapping_add(d);
}

pub struct MD5Writer {
    state: [u32; 4],
    position: u64,
    buf: [u8; 64],
    bufpos: usize,
}

impl Copy for MD5Writer {}
impl Clone for MD5Writer {
    fn clone(&self) -> Self {
        return MD5Writer {
            state: self.state,
            position: self.position,
            buf: self.buf,
            bufpos: self.bufpos
        };
    }
}

impl MD5Writer {
    pub fn new() -> MD5Writer {
        return MD5Writer {
            state: [0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476],
            position: 0,
            buf: [0; 64],
            bufpos: 0,
        };
    }
    pub fn sum(&self) -> [u8; 16] {
        let mut state = self.state;
        let mut buf = self.buf;
        let buflen = buf.len();
        let bufpos = self.bufpos;
        buf[self.bufpos] = 0x80;
        for i in bufpos + 1 .. buflen { buf[i] = 0; }
        if bufpos + 1 > buflen - 8 {
            update(&mut state, &buf);
            for i in 0 .. buflen { buf[i] = 0; }
        }
        LittleEndian::write_u64(&mut buf[buflen - 8 ..], self.position);
        update(&mut state, &buf);
        let mut hashbuf : [u8; 16] = [0; 16];
        for i in 0..4 {
            LittleEndian::write_u32(
                &mut hashbuf[i * 4 .. i * 4 + 4], state[i]);
        }
        return hashbuf;
    }
}

impl Write for MD5Writer {
    fn write(&mut self, mut buf: &[u8]) -> Result<usize> {
        let orig_len = buf.len();
        while buf.len() > 0 {
            let wrsize = cmp::min(self.buf.len() - self.bufpos, buf.len());
            self.buf[self.bufpos .. self.bufpos + wrsize]
                .clone_from_slice(&buf[0 .. wrsize]);
            buf = &buf[wrsize ..];
            self.position = self.position.wrapping_add((wrsize * 8) as u64);
            self.bufpos += wrsize;
            if self.bufpos >= self.buf.len() {
                update(&mut self.state, &self.buf);
                self.bufpos -= self.buf.len();
            }
        }
        return Ok(orig_len);
    }
    fn flush(&mut self) -> Result<()> {
        return Ok(());
    }
}
<|endoftext|>