<commit_before>// ToDO: add tests
<commit_msg>test(hostid) add a test for hostid (#1533)<commit_after>use crate::common::util::*;
extern crate regex;
use self::regex::Regex;

#[test]
fn test_normal() {
    let (_, mut ucmd) = at_and_ucmd!();
    let result = ucmd.run();

    assert!(result.success);
    let re = Regex::new(r"^[0-9a-f]{8}").unwrap();
    assert!(re.is_match(&result.stdout.trim()));
}
<|endoftext|>


<commit_before>use std::{iter, ptr, vec};

use serde::de::{
    self,
    value::{Error, SeqDeserializer},
    Deserializer, IntoDeserializer,
};

#[derive(Debug)]
pub enum ValOrVec<T> {
    Val(T),
    Vec(Vec<T>),
}

impl<T> ValOrVec<T> {
    pub fn push(&mut self, new_val: T) {
        match self {
            Self::Val(val) => {
                let mut vec = Vec::with_capacity(2);
                // Safety:
                //
                // since the vec is pre-allocated, push can't panic, so there
                // is no opportunity for a panic in the unsafe code.
                unsafe {
                    let existing_val = ptr::read(val);
                    vec.push(existing_val);
                    vec.push(new_val);
                    ptr::write(self, Self::Vec(vec))
                }
            }
            Self::Vec(vec) => vec.push(new_val),
        }
    }
}

impl<T> IntoIterator for ValOrVec<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter::new(self)
    }
}

pub enum IntoIter<T> {
    Val(iter::Once<T>),
    Vec(vec::IntoIter<T>),
}

impl<T> IntoIter<T> {
    fn new(vv: ValOrVec<T>) -> Self {
        match vv {
            ValOrVec::Val(val) => Self::Val(iter::once(val)),
            ValOrVec::Vec(vec) => Self::Vec(vec.into_iter()),
        }
    }
}

impl<T> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Val(iter) => iter.next(),
            Self::Vec(iter) => iter.next(),
        }
    }
}

impl<'de, T> IntoDeserializer<'de> for ValOrVec<T>
where
    T: IntoDeserializer<'de> + Deserializer<'de, Error = Error>,
{
    type Deserializer = Self;

    fn into_deserializer(self) -> Self::Deserializer {
        self
    }
}

macro_rules! forward_to_part {
    ($($method:ident,)*) => {
        $(
            fn $method<V>(self, visitor: V) -> Result<V::Value, Self::Error>
                where V: de::Visitor<'de>
            {
                match self {
                    Self::Val(val) => val.$method(visitor),
                    Self::Vec(_) => Err(de::Error::custom("TODO: Error message")),
                }
            }
        )*
    }
}

impl<'de, T> Deserializer<'de> for ValOrVec<T>
where
    T: IntoDeserializer<'de> + Deserializer<'de, Error = Error>,
{
    type Error = Error;

    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            Self::Val(val) => val.deserialize_any(visitor),
            Self::Vec(_) => self.deserialize_seq(visitor),
        }
    }

    fn deserialize_seq<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        visitor.visit_seq(SeqDeserializer::new(self.into_iter()))
    }

    fn deserialize_enum<V>(
        self,
        name: &'static str,
        variants: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            Self::Val(val) => val.deserialize_enum(name, variants, visitor),
            Self::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_tuple<V>(
        self,
        len: usize,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            Self::Val(val) => val.deserialize_tuple(len, visitor),
            Self::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_struct<V>(
        self,
        name: &'static str,
        fields: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            Self::Val(val) => val.deserialize_struct(name, fields, visitor),
            Self::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_unit_struct<V>(
        self,
        name: &'static str,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            Self::Val(val) => val.deserialize_unit_struct(name, visitor),
            Self::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_tuple_struct<V>(
        self,
        name: &'static str,
        len: usize,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            Self::Val(val) => val.deserialize_tuple_struct(name, len, visitor),
            Self::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_newtype_struct<V>(
        self,
        name: &'static str,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            Self::Val(val) => val.deserialize_newtype_struct(name, visitor),
            Self::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    forward_to_part! {
        deserialize_bool,
        deserialize_char,
        deserialize_str,
        deserialize_string,
        deserialize_bytes,
        deserialize_byte_buf,
        deserialize_unit,
        deserialize_u8,
        deserialize_u16,
        deserialize_u32,
        deserialize_u64,
        deserialize_i8,
        deserialize_i16,
        deserialize_i32,
        deserialize_i64,
        deserialize_f32,
        deserialize_f64,
        deserialize_option,
        deserialize_identifier,
        deserialize_ignored_any,
        deserialize_map,
    }
}
<commit_msg>Fix building on 1.36.0<commit_after>use std::{iter, ptr, vec};

use serde::de::{
    self,
    value::{Error, SeqDeserializer},
    Deserializer, IntoDeserializer,
};

#[derive(Debug)]
pub enum ValOrVec<T> {
    Val(T),
    Vec(Vec<T>),
}

impl<T> ValOrVec<T> {
    pub fn push(&mut self, new_val: T) {
        match self {
            ValOrVec::Val(val) => {
                let mut vec = Vec::with_capacity(2);
                // Safety:
                //
                // since the vec is pre-allocated, push can't panic, so there
                // is no opportunity for a panic in the unsafe code.
                unsafe {
                    let existing_val = ptr::read(val);
                    vec.push(existing_val);
                    vec.push(new_val);
                    ptr::write(self, ValOrVec::Vec(vec))
                }
            }
            ValOrVec::Vec(vec) => vec.push(new_val),
        }
    }
}

impl<T> IntoIterator for ValOrVec<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter::new(self)
    }
}

pub enum IntoIter<T> {
    Val(iter::Once<T>),
    Vec(vec::IntoIter<T>),
}

impl<T> IntoIter<T> {
    fn new(vv: ValOrVec<T>) -> Self {
        match vv {
            ValOrVec::Val(val) => IntoIter::Val(iter::once(val)),
            ValOrVec::Vec(vec) => IntoIter::Vec(vec.into_iter()),
        }
    }
}

impl<T> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            IntoIter::Val(iter) => iter.next(),
            IntoIter::Vec(iter) => iter.next(),
        }
    }
}

impl<'de, T> IntoDeserializer<'de> for ValOrVec<T>
where
    T: IntoDeserializer<'de> + Deserializer<'de, Error = Error>,
{
    type Deserializer = Self;

    fn into_deserializer(self) -> Self::Deserializer {
        self
    }
}

macro_rules! forward_to_part {
    ($($method:ident,)*) => {
        $(
            fn $method<V>(self, visitor: V) -> Result<V::Value, Self::Error>
                where V: de::Visitor<'de>
            {
                match self {
                    ValOrVec::Val(val) => val.$method(visitor),
                    ValOrVec::Vec(_) => Err(de::Error::custom("TODO: Error message")),
                }
            }
        )*
    }
}

impl<'de, T> Deserializer<'de> for ValOrVec<T>
where
    T: IntoDeserializer<'de> + Deserializer<'de, Error = Error>,
{
    type Error = Error;

    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            ValOrVec::Val(val) => val.deserialize_any(visitor),
            ValOrVec::Vec(_) => self.deserialize_seq(visitor),
        }
    }

    fn deserialize_seq<V>(self, visitor: V) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        visitor.visit_seq(SeqDeserializer::new(self.into_iter()))
    }

    fn deserialize_enum<V>(
        self,
        name: &'static str,
        variants: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            ValOrVec::Val(val) => val.deserialize_enum(name, variants, visitor),
            ValOrVec::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_tuple<V>(
        self,
        len: usize,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            ValOrVec::Val(val) => val.deserialize_tuple(len, visitor),
            ValOrVec::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_struct<V>(
        self,
        name: &'static str,
        fields: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            ValOrVec::Val(val) => val.deserialize_struct(name, fields, visitor),
            ValOrVec::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_unit_struct<V>(
        self,
        name: &'static str,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            ValOrVec::Val(val) => val.deserialize_unit_struct(name, visitor),
            ValOrVec::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_tuple_struct<V>(
        self,
        name: &'static str,
        len: usize,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            ValOrVec::Val(val) => {
                val.deserialize_tuple_struct(name, len, visitor)
            }
            ValOrVec::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    fn deserialize_newtype_struct<V>(
        self,
        name: &'static str,
        visitor: V,
    ) -> Result<V::Value, Self::Error>
    where
        V: de::Visitor<'de>,
    {
        match self {
            ValOrVec::Val(val) => val.deserialize_newtype_struct(name, visitor),
            ValOrVec::Vec(_) => Err(de::Error::custom("TODO: Error message")),
        }
    }

    forward_to_part! {
        deserialize_bool,
        deserialize_char,
        deserialize_str,
        deserialize_string,
        deserialize_bytes,
        deserialize_byte_buf,
        deserialize_unit,
        deserialize_u8,
        deserialize_u16,
        deserialize_u32,
        deserialize_u64,
        deserialize_i8,
        deserialize_i16,
        deserialize_i32,
        deserialize_i64,
        deserialize_f32,
        deserialize_f64,
        deserialize_option,
        deserialize_identifier,
        deserialize_ignored_any,
        deserialize_map,
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;
extern crate docopt;
extern crate image;
extern crate gpx;
extern crate geo;
extern crate chrono;
extern crate rayon;

use std::fs;
use std::fs::File;
use std::io::BufReader;
use std::path;

use docopt::Docopt;
use gpx::read;
use gpx::{Gpx, Track};
use geo::Point;
use image::ImageBuffer;
use rayon::prelude::*;

const USAGE: &'static str = "
Generate video from GPX files.

Usage:
  derivers <top-lat> <left-lng> <bottom-lat> <right-lng> <width> <height> <directory>
";

#[derive(Debug, Deserialize)]
struct CommandArgs {
    arg_top_lat: f64,
    arg_left_lng: f64,
    arg_bottom_lat: f64,
    arg_right_lng: f64,
    arg_width: u32,
    arg_height: u32,
    arg_directory: String,
}

#[derive(Debug)]
struct ImageFrame {
    top_left: Point<f64>,
    bottom_right: Point<f64>,
    width: u32,
    height: u32,
}

impl ImageFrame {
    pub fn from(args: &CommandArgs) -> ImageFrame {
        // h == w * (top - bottom) / (right - left)
        let height = (args.arg_width as f64) *
            ((args.arg_top_lat - args.arg_bottom_lat) /
             (args.arg_right_lng - args.arg_left_lng));

        println!("Computed height: {:?}", height);

        ImageFrame {
            top_left: Point::new(args.arg_left_lng, args.arg_top_lat),
            bottom_right: Point::new(args.arg_right_lng, args.arg_bottom_lat),
            width: args.arg_width,
            height: args.arg_height,
        }
    }

    pub fn get_image(&self) -> image::DynamicImage {
        let buf = ImageBuffer::from_pixel(self.width, self.height, image::Rgb([0, 0, 0]));

        image::ImageRgb8(buf)
    }

    // Using simple equirectangular projection for now
    pub fn project_to_screen(&self, coord: &Point<f64>) -> Option<(u32, u32)> {
        // lng is x pos
        let x_pos = self.top_left.lng() - coord.lng();
        let y_pos = self.top_left.lat() - coord.lat();

        let x_offset = x_pos / (self.top_left.lng() - self.bottom_right.lng());
        let y_offset = y_pos / (self.top_left.lat() - self.bottom_right.lat());

        let (x, y) = ((x_offset * self.width as f64),
                      (y_offset * self.height as f64));

        if (x < 0.0 || x as u32 >= self.width) ||
            (y < 0.0 || y as u32 >= self.height) {
            None
        } else {
            Some((x as u32, y as u32))
        }
    }
}

#[derive(Debug)]
struct Activity {
    name: String,
    date: chrono::DateTime<chrono::Utc>,
    track_points: Vec<Point<f64>>,
}

fn parse_gpx(path: &path::PathBuf) -> Option<Activity> {
    let file = File::open(path).unwrap();
    let reader = BufReader::new(file);

    let gpx: Gpx = read(reader).unwrap();

    // Nothing to do if there are no tracks
    if gpx.tracks.len() == 0 {
        return None;
    } else if gpx.tracks.len() > 1 {
        println!("Warning! more than 1 track, just taking first");
    }

    let track: &Track = &gpx.tracks[0];

    let mut activity = Activity {
        name: track.name.clone().unwrap_or(String::from("Untitled")),
        date: chrono::Utc::now(),
        track_points: vec![],
    };

    if let Some(metadata) = gpx.metadata {
        if let Some(time) = metadata.time {
            activity.date = time;
        }
    }

    // Append all the waypoints.
    for seg in track.segments.iter() {
        let points = seg.points.iter().map(|ref wpt| wpt.point());
        activity.track_points.extend(points);
    }

    if activity.track_points.len() == 0 {
        None
    } else {
        Some(activity)
    }
}

fn main() {
    let args: CommandArgs = Docopt::new(USAGE)
        .and_then(|d| d.deserialize())
        .unwrap_or_else(|e| e.exit());

    println!("{:?}", args);

    let frame = ImageFrame::from(&args);
    let mut image = frame.get_image();

    let paths: Vec<path::PathBuf> = fs::read_dir(args.arg_directory)
        .unwrap()
        .into_iter()
        .map(|p| p.unwrap().path())
        .collect();

    let activities: Vec<Activity> = paths
        .into_par_iter()
        .filter_map(|ref p| parse_gpx(p))
        .collect();

    let fout = &mut File::create("heatmap.ppm").unwrap();

    for act in activities {
        println!("Activity: {:?}", act.name);

        let pixels: Vec<(u32, u32)> = act.track_points
            .par_iter()
            .filter_map(|ref pt| frame.project_to_screen(pt))
            .collect();

        for (x, y) in pixels.into_iter() {
            let pixel = image.as_mut_rgb8().unwrap().get_pixel_mut(x, y);

            let c = if pixel[0] == 255 {
                pixel[0]
            } else if pixel[0] == 0 {
                25
            } else {
                pixel[0] + 5
            };

            *pixel = image::Rgb([c, c, c]);
        }

        image.save(fout, image::PPM).unwrap();
    }
}
<commit_msg>Write more frames<commit_after>#[macro_use]
extern crate serde_derive;
extern crate docopt;
extern crate image;
extern crate gpx;
extern crate geo;
extern crate chrono;
extern crate rayon;

use std::fs;
use std::fs::File;
use std::io::BufReader;
use std::path;

use docopt::Docopt;
use gpx::read;
use gpx::{Gpx, Track};
use geo::Point;
use image::ImageBuffer;
use rayon::prelude::*;

const USAGE: &'static str = "
Generate video from GPX files.

Usage:
  derivers <top-lat> <left-lng> <bottom-lat> <right-lng> <width> <height> <directory>
";

#[derive(Debug, Deserialize)]
struct CommandArgs {
    arg_top_lat: f64,
    arg_left_lng: f64,
    arg_bottom_lat: f64,
    arg_right_lng: f64,
    arg_width: u32,
    arg_height: u32,
    arg_directory: String,
}

#[derive(Debug)]
struct ImageFrame {
    top_left: Point<f64>,
    bottom_right: Point<f64>,
    width: u32,
    height: u32,
}

impl ImageFrame {
    pub fn from(args: &CommandArgs) -> ImageFrame {
        // h == w * (top - bottom) / (right - left)
        let height = (args.arg_width as f64) *
            ((args.arg_top_lat - args.arg_bottom_lat) / (args.arg_right_lng - args.arg_left_lng));

        println!("Computed height: {:?}", height);

        ImageFrame {
            top_left: Point::new(args.arg_left_lng, args.arg_top_lat),
            bottom_right: Point::new(args.arg_right_lng, args.arg_bottom_lat),
            width: args.arg_width,
            height: height as u32,
        }
    }

    pub fn get_image(&self) -> image::DynamicImage {
        let buf = ImageBuffer::from_pixel(self.width, self.height, image::Rgb([255, 255, 255]));

        image::ImageRgb8(buf)
    }

    // Using simple equirectangular projection for now
    pub fn project_to_screen(&self, coord: &Point<f64>) -> Option<(u32, u32)> {
        // lng is x pos
        let x_pos = self.top_left.lng() - coord.lng();
        let y_pos = self.top_left.lat() - coord.lat();

        let x_offset = x_pos / (self.top_left.lng() - self.bottom_right.lng());
        let y_offset = y_pos / (self.top_left.lat() - self.bottom_right.lat());

        let (x, y) = (
            (x_offset * self.width as f64),
            (y_offset * self.height as f64),
        );

        if (x < 0.0 || x as u32 >= self.width) || (y < 0.0 || y as u32 >= self.height) {
            None
        } else {
            Some((x as u32, y as u32))
        }
    }
}

#[derive(Debug)]
struct Activity {
    name: String,
    date: chrono::DateTime<chrono::Utc>,
    track_points: Vec<Point<f64>>,
}

fn parse_gpx(path: &path::PathBuf) -> Option<Activity> {
    let file = File::open(path).unwrap();
    let reader = BufReader::new(file);

    let gpx: Gpx = read(reader).unwrap();

    // Nothing to do if there are no tracks
    if gpx.tracks.len() == 0 {
        return None;
    } else if gpx.tracks.len() > 1 {
        println!("Warning! more than 1 track, just taking first");
    }

    let track: &Track = &gpx.tracks[0];

    let mut activity = Activity {
        name: track.name.clone().unwrap_or(String::from("Untitled")),
        date: chrono::Utc::now(),
        track_points: vec![],
    };

    if let Some(metadata) = gpx.metadata {
        if let Some(time) = metadata.time {
            activity.date = time;
        }
    }

    // Append all the waypoints.
    for seg in track.segments.iter() {
        let points = seg.points.iter().map(|ref wpt| wpt.point());
        activity.track_points.extend(points);
    }

    if activity.track_points.len() == 0 {
        None
    } else {
        Some(activity)
    }
}

fn main() {
    let args: CommandArgs = Docopt::new(USAGE)
        .and_then(|d| d.deserialize())
        .unwrap_or_else(|e| e.exit());

    println!("{:?}", args);

    let frame = ImageFrame::from(&args);
    let mut image = frame.get_image();

    let paths: Vec<path::PathBuf> = fs::read_dir(args.arg_directory)
        .unwrap()
        .into_iter()
        .map(|p| p.unwrap().path())
        .collect();

    let mut activities: Vec<Activity> = paths
        .into_par_iter()
        .filter_map(|ref p| parse_gpx(p))
        .collect();

    activities.sort_by_key_unstable(|a| a.date);

    let fout = &mut File::create("heatmap.ppm").unwrap();
    let mut counter = 0;

    for act in activities {
        println!("Activity: {:?}", act.name);

        let pixels: Vec<(u32, u32)> = act.track_points
            .par_iter()
            .filter_map(|ref pt| frame.project_to_screen(pt))
            .collect();

        for (x, y) in pixels.into_iter() {
            {
                let pixel = image.as_mut_rgb8().unwrap().get_pixel_mut(x, y);

                let c = if pixel[0] == 0 {
                    pixel[0]
                } else if pixel[0] == 255 {
                    pixel[0] - 25
                } else {
                    pixel[0] - 5
                };

                *pixel = image::Rgb([c, c, c]);
                counter += 1;
            }

            if counter == 50 {
                image.save(fout, image::PPM).unwrap();
                counter = 0;
            }
        }

    }
}
<|endoftext|>


<commit_before>use clap::Parser;
use dicom_core::{dicom_value, smallvec};
use dicom_core::{DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_dump::DumpOptions;
use dicom_encoding::transfer_syntax;
use dicom_object::{mem::InMemDicomObject, open_file, StandardDataDictionary};
use dicom_transfer_syntax_registry::{entries, TransferSyntaxRegistry};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use query::parse_queries;
use smallvec::smallvec;
use snafu::{prelude::*, ErrorCompat};
use std::io::{stderr, Read};
use std::path::PathBuf;
use tracing::{debug, error, info, warn, Level};
use transfer_syntax::TransferSyntaxIndex;

mod query;

/// DICOM C-FIND SCU
#[derive(Debug, Parser)]
struct App {
    /// socket address to FIND SCP (example: "127.0.0.1:1045")
    addr: String,
    /// a DICOM file representing the query object
    file: Option<PathBuf>,
    /// sequence of queries
    #[clap(short('q'))]
    query: Vec<String>,

    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling AE title
    #[clap(long = "calling-ae-title", default_value = "FIND-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[clap(long = "called-ae-title")]
    called_ae_title: Option<String>,
    /// the maximum PDU length
    #[clap(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,

    /// use patient root information model
    #[clap(long, conflicts_with = "study")]
    patient: bool,
    /// use study root information model (default)
    #[clap(long, conflicts_with = "patient")]
    study: bool,
}

fn report<E: 'static>(err: &E)
where
    E: std::error::Error,
{
    error!("{}", err);
    if let Some(source) = err.source() {
        eprintln!();
        eprintln!("Caused by:");
        for (i, e) in std::iter::successors(Some(source), |e| e.source()).enumerate() {
            eprintln!("   {}: {}", i, e);
        }
    }
}

fn report_backtrace<E: 'static>(err: &E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    let env_backtrace = std::env::var("RUST_BACKTRACE").unwrap_or_default();
    let env_lib_backtrace = std::env::var("RUST_LIB_BACKTRACE").unwrap_or_default();
    if env_lib_backtrace == "1" || (env_backtrace == "1" && env_lib_backtrace != "0") {
        if let Some(backtrace) = ErrorCompat::backtrace(&err) {
            eprintln!();
            eprintln!("Backtrace:");
            eprintln!("{}", backtrace);
        }
    }
}

fn report_with_backtrace<E: 'static>(err: E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    report(&err);
    report_backtrace(&err);
}

fn main() {
    run().unwrap_or_else(|err| {
        report_with_backtrace(err);
        std::process::exit(-2);
    });
}

#[derive(Debug, Snafu)]
enum Error {
    /// Could not initialize SCU
    InitScu {
        source: dicom_ul::association::client::Error,
    },

    /// Could not construct DICOM command
    CreateCommand { source: dicom_object::Error },

    /// Could not read DICOM command
    ReadCommand { source: dicom_object::Error },

    /// Could not dump DICOM output
    DumpOutput { source: std::io::Error },

    #[snafu(whatever, display("{}", message))]
    Other {
        message: String,
        #[snafu(source(from(Box<dyn std::error::Error + 'static>, Some)))]
        source: Option<Box<dyn std::error::Error + 'static>>,
    },
}

fn build_query(
    file: Option<PathBuf>,
    q: Vec<String>,
    patient: bool,
    study: bool,
    verbose: bool,
) -> Result<InMemDicomObject, Error> {
    match (file, q) {
        (Some(file), q) => {
            if !q.is_empty() {
                whatever!("Conflicted file with query terms");
            }

            if verbose {
                info!("Opening file '{}'...", file.display());
            }

            open_file(file)
                .context(CreateCommandSnafu)
                .map(|file| file.into_inner())
        }
        (None, q) => {
            let mut obj =
                parse_queries(&q).whatever_context("Could not build query object from terms")?;

            // (0008,0052) CS QueryRetrieveLevel
            let level = match (patient, study) {
                (true, false) => "PATIENT",
                (false, true) |
                (false, false) => "STUDY",
                _ => unreachable!(),
            };
            obj.put(DataElement::new(
                tags::QUERY_RETRIEVE_LEVEL,
                VR::CS,
                PrimitiveValue::from(level),
            ));

            Ok(obj)
        }
    }
}

fn run() -> Result<(), Error> {
    let App {
        addr,
        file,
        verbose,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
        patient,
        study,
        query,
    } = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if verbose { Level::DEBUG } else { Level::INFO })
            .finish(),
    )
    .unwrap_or_else(|e| {
        report(&e);
    });

    let dcm_query = build_query(file, query, patient, study, verbose)?;

    let abstract_syntax = match (patient, study) {
        // Patient Root Query/Retrieve Information Model - FIND
        (true, false) => "1.2.840.10008.5.1.4.1.2.1.1",
        // Study Root Query/Retrieve Information Model – FIND (default)
        (false, false) | (false, true) => "1.2.840.10008.5.1.4.1.2.2.1",
        // Series
        _ => unreachable!("Unexpected flag combination"),
    };

    if verbose {
        info!("Establishing association with '{}'...", &addr);
    }

    let mut scu_opt = ClientAssociationOptions::new()
        .with_abstract_syntax(abstract_syntax)
        .calling_ae_title(calling_ae_title)
        .max_pdu_length(max_pdu_length);

    if let Some(called_ae_title) = called_ae_title {
        scu_opt = scu_opt.called_ae_title(called_ae_title);
    }

    let mut scu = scu_opt.establish_with(&addr).context(InitScuSnafu)?;

    if verbose {
        info!("Association established");
    }

    let pc_selected = if let Some(pc_selected) = scu.presentation_contexts().first() {
        pc_selected
    } else {
        error!("Could not choose a presentation context");
        let _ = scu.abort();
        std::process::exit(-2);
    };
    let pc_selected_id = pc_selected.id;

    let ts = if let Some(ts) = TransferSyntaxRegistry.get(&pc_selected.transfer_syntax) {
        ts
    } else {
        error!("Poorly negotiated transfer syntax");
        let _ = scu.abort();
        std::process::exit(-2);
    };

    if verbose {
        debug!("Transfer Syntax: {}", ts.name());
    }

    let cmd = find_req_command(abstract_syntax, 1);

    let mut cmd_data = Vec::with_capacity(128);
    cmd.write_dataset_with_ts(&mut cmd_data, &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased())
        .whatever_context("Failed to write command")?;

    let implicit_vr_le = entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

    let mut iod_data = Vec::with_capacity(128);
    dcm_query
        .write_dataset_with_ts(&mut iod_data, &implicit_vr_le)
        .whatever_context("failed to write identifier dataset")?;

    let nbytes = cmd_data.len() + iod_data.len();

    if verbose {
        debug!("Sending query ({} B)...", nbytes);
    }

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Command,
            is_last: true,
            data: cmd_data,
        }],
    };
    scu.send(&pdu).whatever_context("Could not send command")?;

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Data,
            is_last: true,
            data: iod_data,
        }],
    };
    scu.send(&pdu)
        .whatever_context("Could not send C-Find request")?;

    if verbose {
        debug!("Awaiting response...");
    }

    let mut i = 0;
    loop {
        let rsp_pdu = scu
            .receive()
            .whatever_context("Failed to receive response from remote node")?;

        match rsp_pdu {
            Pdu::PData { data } => {
                let data_value = &data[0];

                let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                    &data_value.data[..],
                    &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
                )
                .context(ReadCommandSnafu)?;
                if verbose {
                    eprintln!("Match #{} Response command:", i);
                    DumpOptions::new()
                        .dump_object_to(stderr(), &cmd_obj)
                        .context(DumpOutputSnafu)?;
                }
                let status = cmd_obj
                    .element(tags::STATUS)
                    .whatever_context("status code from response is missing")?
                    .to_int::<u16>()
                    .whatever_context("failed to read status code")?;
                if status == 0 {
                    if verbose {
                        debug!("Matching is complete");
                    }
                    break;
                } else if status == 0xFF00 || status == 0xFF01 {
                    if verbose {
                        debug!("Operation pending: {:x}", status);
                    }

                    // fetch DICOM data

                    let dcm = {
                        let mut rsp = scu.receive_pdata();
                        let mut response_data = Vec::new();
                        rsp.read_to_end(&mut response_data)
                            .whatever_context("Failed to read response data")?;

                        InMemDicomObject::read_dataset_with_ts(&response_data[..], &implicit_vr_le)
                            .whatever_context("Could not read response data set")?
                    };

                    println!("------------ Match #{} ------------", i);
                    DumpOptions::new()
                        .dump_object(&dcm)
                        .context(DumpOutputSnafu)?;
                    i += 1;
                } else {
                    warn!("Operation failed (status code {})", status);
                    break;
                }
            }

            pdu @ Pdu::Unknown { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::AssociationAC { .. }
            | pdu @ Pdu::AssociationRJ { .. }
            | pdu @ Pdu::ReleaseRQ
            | pdu @ Pdu::ReleaseRP
            | pdu @ Pdu::AbortRQ { .. } => {
                error!("Unexpected SCP response: {:?}", pdu);
                let _ = scu.abort();
                std::process::exit(-2);
            }
        }
    }
    let _ = scu.release();

    Ok(())
}

fn find_req_command(
    sop_class_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + even_len(sop_class_uid.len())   // SOP Class UID
            + 8 + 2 // command field
            + 8 + 2 // message ID
            + 8 + 2 // priority
            + 8 + 2, // data set type
        ),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        PrimitiveValue::from(sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        // 0020H: C-FIND-RQ message
        dicom_value!(U16, [0x0020]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        // medium
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    obj
}

fn even_len(l: usize) -> u32 {
    ((l + 1) & !1) as u32
}

#[cfg(test)]
mod tests {
    use super::even_len;
    #[test]
    fn test_even_len() {
        assert_eq!(even_len(0), 0);
        assert_eq!(even_len(1), 2);
        assert_eq!(even_len(2), 2);
        assert_eq!(even_len(3), 4);
        assert_eq!(even_len(4), 4);
        assert_eq!(even_len(5), 6);
    }
}
<commit_msg>[storescp] Report 0 results<commit_after>use clap::Parser;
use dicom_core::{dicom_value, smallvec};
use dicom_core::{DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_dump::DumpOptions;
use dicom_encoding::transfer_syntax;
use dicom_object::{mem::InMemDicomObject, open_file, StandardDataDictionary};
use dicom_transfer_syntax_registry::{entries, TransferSyntaxRegistry};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use query::parse_queries;
use smallvec::smallvec;
use snafu::{prelude::*, ErrorCompat};
use std::io::{stderr, Read};
use std::path::PathBuf;
use tracing::{debug, error, info, warn, Level};
use transfer_syntax::TransferSyntaxIndex;

mod query;

/// DICOM C-FIND SCU
#[derive(Debug, Parser)]
struct App {
    /// socket address to FIND SCP (example: "127.0.0.1:1045")
    addr: String,
    /// a DICOM file representing the query object
    file: Option<PathBuf>,
    /// sequence of queries
    #[clap(short('q'))]
    query: Vec<String>,

    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling AE title
    #[clap(long = "calling-ae-title", default_value = "FIND-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[clap(long = "called-ae-title")]
    called_ae_title: Option<String>,
    /// the maximum PDU length
    #[clap(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,

    /// use patient root information model
    #[clap(long, conflicts_with = "study")]
    patient: bool,
    /// use study root information model (default)
    #[clap(long, conflicts_with = "patient")]
    study: bool,
}

fn report<E: 'static>(err: &E)
where
    E: std::error::Error,
{
    error!("{}", err);
    if let Some(source) = err.source() {
        eprintln!();
        eprintln!("Caused by:");
        for (i, e) in std::iter::successors(Some(source), |e| e.source()).enumerate() {
            eprintln!("   {}: {}", i, e);
        }
    }
}

fn report_backtrace<E: 'static>(err: &E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    let env_backtrace = std::env::var("RUST_BACKTRACE").unwrap_or_default();
    let env_lib_backtrace = std::env::var("RUST_LIB_BACKTRACE").unwrap_or_default();
    if env_lib_backtrace == "1" || (env_backtrace == "1" && env_lib_backtrace != "0") {
        if let Some(backtrace) = ErrorCompat::backtrace(&err) {
            eprintln!();
            eprintln!("Backtrace:");
            eprintln!("{}", backtrace);
        }
    }
}

fn report_with_backtrace<E: 'static>(err: E)
where
    E: std::error::Error,
    E: ErrorCompat,
{
    report(&err);
    report_backtrace(&err);
}

fn main() {
    run().unwrap_or_else(|err| {
        report_with_backtrace(err);
        std::process::exit(-2);
    });
}

#[derive(Debug, Snafu)]
enum Error {
    /// Could not initialize SCU
    InitScu {
        source: dicom_ul::association::client::Error,
    },

    /// Could not construct DICOM command
    CreateCommand { source: dicom_object::Error },

    /// Could not read DICOM command
    ReadCommand { source: dicom_object::Error },

    /// Could not dump DICOM output
    DumpOutput { source: std::io::Error },

    #[snafu(whatever, display("{}", message))]
    Other {
        message: String,
        #[snafu(source(from(Box<dyn std::error::Error + 'static>, Some)))]
        source: Option<Box<dyn std::error::Error + 'static>>,
    },
}

fn build_query(
    file: Option<PathBuf>,
    q: Vec<String>,
    patient: bool,
    study: bool,
    verbose: bool,
) -> Result<InMemDicomObject, Error> {
    match (file, q) {
        (Some(file), q) => {
            if !q.is_empty() {
                whatever!("Conflicted file with query terms");
            }

            if verbose {
                info!("Opening file '{}'...", file.display());
            }

            open_file(file)
                .context(CreateCommandSnafu)
                .map(|file| file.into_inner())
        }
        (None, q) => {
            let mut obj =
                parse_queries(&q).whatever_context("Could not build query object from terms")?;

            // (0008,0052) CS QueryRetrieveLevel
            let level = match (patient, study) {
                (true, false) => "PATIENT",
                (false, true) |
                (false, false) => "STUDY",
                _ => unreachable!(),
            };
            obj.put(DataElement::new(
                tags::QUERY_RETRIEVE_LEVEL,
                VR::CS,
                PrimitiveValue::from(level),
            ));

            Ok(obj)
        }
    }
}

fn run() -> Result<(), Error> {
    let App {
        addr,
        file,
        verbose,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
        patient,
        study,
        query,
    } = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if verbose { Level::DEBUG } else { Level::INFO })
            .finish(),
    )
    .unwrap_or_else(|e| {
        report(&e);
    });

    let dcm_query = build_query(file, query, patient, study, verbose)?;

    let abstract_syntax = match (patient, study) {
        // Patient Root Query/Retrieve Information Model - FIND
        (true, false) => "1.2.840.10008.5.1.4.1.2.1.1",
        // Study Root Query/Retrieve Information Model – FIND (default)
        (false, false) | (false, true) => "1.2.840.10008.5.1.4.1.2.2.1",
        // Series
        _ => unreachable!("Unexpected flag combination"),
    };

    if verbose {
        info!("Establishing association with '{}'...", &addr);
    }

    let mut scu_opt = ClientAssociationOptions::new()
        .with_abstract_syntax(abstract_syntax)
        .calling_ae_title(calling_ae_title)
        .max_pdu_length(max_pdu_length);

    if let Some(called_ae_title) = called_ae_title {
        scu_opt = scu_opt.called_ae_title(called_ae_title);
    }

    let mut scu = scu_opt.establish_with(&addr).context(InitScuSnafu)?;

    if verbose {
        info!("Association established");
    }

    let pc_selected = if let Some(pc_selected) = scu.presentation_contexts().first() {
        pc_selected
    } else {
        error!("Could not choose a presentation context");
        let _ = scu.abort();
        std::process::exit(-2);
    };
    let pc_selected_id = pc_selected.id;

    let ts = if let Some(ts) = TransferSyntaxRegistry.get(&pc_selected.transfer_syntax) {
        ts
    } else {
        error!("Poorly negotiated transfer syntax");
        let _ = scu.abort();
        std::process::exit(-2);
    };

    if verbose {
        debug!("Transfer Syntax: {}", ts.name());
    }

    let cmd = find_req_command(abstract_syntax, 1);

    let mut cmd_data = Vec::with_capacity(128);
    cmd.write_dataset_with_ts(&mut cmd_data, &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased())
        .whatever_context("Failed to write command")?;

    let implicit_vr_le = entries::IMPLICIT_VR_LITTLE_ENDIAN.erased();

    let mut iod_data = Vec::with_capacity(128);
    dcm_query
        .write_dataset_with_ts(&mut iod_data, &implicit_vr_le)
        .whatever_context("failed to write identifier dataset")?;

    let nbytes = cmd_data.len() + iod_data.len();

    if verbose {
        debug!("Sending query ({} B)...", nbytes);
    }

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Command,
            is_last: true,
            data: cmd_data,
        }],
    };
    scu.send(&pdu).whatever_context("Could not send command")?;

    let pdu = Pdu::PData {
        data: vec![PDataValue {
            presentation_context_id: pc_selected_id,
            value_type: PDataValueType::Data,
            is_last: true,
            data: iod_data,
        }],
    };
    scu.send(&pdu)
        .whatever_context("Could not send C-Find request")?;

    if verbose {
        debug!("Awaiting response...");
    }

    let mut i = 0;
    loop {
        let rsp_pdu = scu
            .receive()
            .whatever_context("Failed to receive response from remote node")?;

        match rsp_pdu {
            Pdu::PData { data } => {
                let data_value = &data[0];

                let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                    &data_value.data[..],
                    &entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
                )
                .context(ReadCommandSnafu)?;
                if verbose {
                    eprintln!("Match #{} Response command:", i);
                    DumpOptions::new()
                        .dump_object_to(stderr(), &cmd_obj)
                        .context(DumpOutputSnafu)?;
                }
                let status = cmd_obj
                    .element(tags::STATUS)
                    .whatever_context("status code from response is missing")?
                    .to_int::<u16>()
                    .whatever_context("failed to read status code")?;
                if status == 0 {
                    if verbose {
                        debug!("Matching is complete");
                    }
                    if i == 0 {
                        info!("No results matching query");
                    }
                    break;
                } else if status == 0xFF00 || status == 0xFF01 {
                    if verbose {
                        debug!("Operation pending: {:x}", status);
                    }

                    // fetch DICOM data

                    let dcm = {
                        let mut rsp = scu.receive_pdata();
                        let mut response_data = Vec::new();
                        rsp.read_to_end(&mut response_data)
                            .whatever_context("Failed to read response data")?;

                        InMemDicomObject::read_dataset_with_ts(&response_data[..], &implicit_vr_le)
                            .whatever_context("Could not read response data set")?
                    };

                    println!("------------ Match #{} ------------", i);
                    DumpOptions::new()
                        .dump_object(&dcm)
                        .context(DumpOutputSnafu)?;
                    i += 1;
                } else {
                    warn!("Operation failed (status code {})", status);
                    break;
                }
            }

            pdu @ Pdu::Unknown { .. }
            | pdu @ Pdu::AssociationRQ { .. }
            | pdu @ Pdu::AssociationAC { .. }
            | pdu @ Pdu::AssociationRJ { .. }
            | pdu @ Pdu::ReleaseRQ
            | pdu @ Pdu::ReleaseRP
            | pdu @ Pdu::AbortRQ { .. } => {
                error!("Unexpected SCP response: {:?}", pdu);
                let _ = scu.abort();
                std::process::exit(-2);
            }
        }
    }
    let _ = scu.release();

    Ok(())
}

fn find_req_command(
    sop_class_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + even_len(sop_class_uid.len())   // SOP Class UID
            + 8 + 2 // command field
            + 8 + 2 // message ID
            + 8 + 2 // priority
            + 8 + 2, // data set type
        ),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        PrimitiveValue::from(sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        // 0020H: C-FIND-RQ message
        dicom_value!(U16, [0x0020]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        // medium
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    obj
}

fn even_len(l: usize) -> u32 {
    ((l + 1) & !1) as u32
}

#[cfg(test)]
mod tests {
    use super::even_len;
    #[test]
    fn test_even_len() {
        assert_eq!(even_len(0), 0);
        assert_eq!(even_len(1), 2);
        assert_eq!(even_len(2), 2);
        assert_eq!(even_len(3), 4);
        assert_eq!(even_len(4), 4);
        assert_eq!(even_len(5), 6);
    }
}
<|endoftext|>


<commit_before>use std::io::IoResult;
use std::io::net::tcp::TcpStream;
use std::io::net::ip::SocketAddr;

pub struct ClosableTcpStream {
    stream: TcpStream,
    close: Receiver<()>,
    close_read: bool,
    close_write: bool,
}

impl ClosableTcpStream {
    pub fn new(stream: TcpStream, close_read: bool, close_write: bool) -> (ClosableTcpStream, Sender<()>) {
        let (tx, rx) = channel();

        let acc = ClosableTcpStream {
            stream: stream,
            close: rx,
            close_read: close_read,
            close_write: close_write,
        };

        (acc, tx)
    }

    pub fn peer_name(&mut self) -> IoResult<SocketAddr> {
        self.stream.peer_name()
    }
}

impl Drop for ClosableTcpStream {
    fn drop(&mut self) {
        if self.close_read {
            self.stream.close_read().ok();      // ignoring outcome
        }
        if self.close_write {
            self.stream.close_write().ok();     // ignoring outcome
        }
    }
}

impl Reader for ClosableTcpStream {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        use std::io;

        loop {
            if self.close.try_recv().is_ok() {
                return Err(io::standard_error(io::Closed));
            }

            self.stream.set_read_timeout(Some(100));

            match self.stream.read(buf) {
                Err(ref err) if err.kind == io::TimedOut
                    => continue,
                a => return a
            };
        }
    }
}

impl Writer for ClosableTcpStream {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        use std::io;

        loop {
            if self.close.try_recv().is_ok() {
                return Err(io::standard_error(io::Closed));
            }

            self.stream.set_write_timeout(Some(100));

            match self.stream.write(buf) {
                Err(ref err) if err.kind == io::TimedOut
                    => continue,
                Err(err) => {
                    match err.kind {
                        io::ShortWrite(nb) => return self.write(buf.slice_from(nb)),
                        _ => return Err(err)
                    };
                }
                a => return a
            };
        }
    }
}
<commit_msg>Fix ClosableTcpStream not having flush()<commit_after>use std::io::IoResult;
use std::io::net::tcp::TcpStream;
use std::io::net::ip::SocketAddr;

pub struct ClosableTcpStream {
    stream: TcpStream,
    close: Receiver<()>,
    close_read: bool,
    close_write: bool,
}

impl ClosableTcpStream {
    pub fn new(stream: TcpStream, close_read: bool, close_write: bool) -> (ClosableTcpStream, Sender<()>) {
        let (tx, rx) = channel();

        let acc = ClosableTcpStream {
            stream: stream,
            close: rx,
            close_read: close_read,
            close_write: close_write,
        };

        (acc, tx)
    }

    pub fn peer_name(&mut self) -> IoResult<SocketAddr> {
        self.stream.peer_name()
    }
}

impl Drop for ClosableTcpStream {
    fn drop(&mut self) {
        if self.close_read {
            self.stream.close_read().ok();      // ignoring outcome
        }
        if self.close_write {
            self.stream.close_write().ok();     // ignoring outcome
        }
    }
}

impl Reader for ClosableTcpStream {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        use std::io;

        loop {
            if self.close.try_recv().is_ok() {
                return Err(io::standard_error(io::Closed));
            }

            self.stream.set_read_timeout(Some(100));

            match self.stream.read(buf) {
                Err(ref err) if err.kind == io::TimedOut
                    => continue,
                a => return a
            };
        }
    }
}

impl Writer for ClosableTcpStream {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        use std::io;

        loop {
            if self.close.try_recv().is_ok() {
                return Err(io::standard_error(io::Closed));
            }

            self.stream.set_write_timeout(Some(100));

            match self.stream.write(buf) {
                Err(ref err) if err.kind == io::TimedOut
                    => continue,
                Err(err) => {
                    match err.kind {
                        io::ShortWrite(nb) => return self.write(buf.slice_from(nb)),
                        _ => return Err(err)
                    };
                }
                a => return a
            };
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        self.stream.flush()
    }
}
<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) kwantam <kwantam@gmail.com>
// *     * 20150507 ~ added big_ routines to prevent overflow when num > 2^63
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use std::mem::swap;

// This is incorrectly reported as dead code,
//  presumably when included in build.rs.
#[allow(dead_code)]
pub(crate) fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    type I: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::I;
    fn to_u64(&self, n: Self::I) -> u64;
    fn add(&self, a: Self::I, b: Self::I) -> Self::I;
    fn mul(&self, a: Self::I, b: Self::I) -> Self::I;

    fn pow(&self, mut a: Self::I, mut b: u64) -> Self::I {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::I {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::I {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::I {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery {
    a: u64,
    n: u64,
}

impl Montgomery {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: u128) -> u64 {
        debug_assert!(x < (self.n as u128) << 64);
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let m = (x as u64).wrapping_mul(*a);
        let nm = (*n as u128) * (m as u128);
        let (xnm, overflow) = (x as u128).overflowing_add(nm); // x + n*m
        debug_assert_eq!(xnm % (1 << 64), 0);

        // (x + n*m) / R
        // in case of overflow, this is (2¹²⁸ + xnm)/2⁶⁴ - n = xnm/2⁶⁴ + (2⁶⁴ - n)
        let y = (xnm >> 64) as u64 + if !overflow { 0 } else { n.wrapping_neg() };

        if y >= *n {
            y - n
        } else {
            y
        }
    }
}

impl Arithmetic for Montgomery {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type I = u64;

    fn new(n: u64) -> Self {
        let a = inv_mod_u64(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(a), 1_u64.wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n
    }

    fn from_u64(&self, x: u64) -> Self::I {
        // TODO: optimise!
        assert!(x < self.n);
        let r = (((x as u128) << 64) % self.n as u128) as u64;
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::I) -> u64 {
        self.reduce(n as u128)
    }

    fn add(&self, a: Self::I, b: Self::I) -> Self::I {
        let r = a + b;

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 = (((a_r as u128) + (b_r as u128)) % (self.n as u128)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::I, b: Self::I) -> Self::I {
        let r = self.reduce((a as u128) * (b as u128));

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 = (((a_r as u128) * (b_r as u128)) % (self.n as u128)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn inv_mod_u64(a: u64) -> u64 {
    assert!(a % 2 == 1);
    let mut t = 0u64;
    let mut newt = 1u64;
    let mut r = 0u64;
    let mut newr = a;

    while newr != 0 {
        let quot = if r == 0 {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            std::u64::MAX
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(quot.wrapping_mul(newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(quot.wrapping_mul(newr));
        r = newr;
        newr = newrp;
    }

    assert_eq!(r, 1);
    t
}

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

    #[test]
    fn test_inverter() {
        // All odd integers from 1 to 20 000
        let mut test_values = (0..10_000u64).map(|i| 2 * i + 1);

        assert!(test_values.all(|x| x.wrapping_mul(inv_mod_u64(x)) == 1));
    }

    #[test]
    fn test_montgomery_add() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_montgomery_mult() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_montgomery_roundtrip() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }
}
<commit_msg>factor::numeric::Montgomery::add: Deal with rare overflow case<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) Wiktor Kuropatwa <wiktor.kuropatwa@gmail.com>
// * (c) kwantam <kwantam@gmail.com>
// *     * 20150507 ~ added big_ routines to prevent overflow when num > 2^63
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

use std::mem::swap;

// This is incorrectly reported as dead code,
//  presumably when included in build.rs.
#[allow(dead_code)]
pub(crate) fn gcd(mut a: u64, mut b: u64) -> u64 {
    while b > 0 {
        a %= b;
        swap(&mut a, &mut b);
    }
    a
}

pub(crate) trait Arithmetic: Copy + Sized {
    type I: Copy + Sized + Eq;

    fn new(m: u64) -> Self;
    fn modulus(&self) -> u64;
    fn from_u64(&self, n: u64) -> Self::I;
    fn to_u64(&self, n: Self::I) -> u64;
    fn add(&self, a: Self::I, b: Self::I) -> Self::I;
    fn mul(&self, a: Self::I, b: Self::I) -> Self::I;

    fn pow(&self, mut a: Self::I, mut b: u64) -> Self::I {
        let (_a, _b) = (a, b);
        let mut result = self.one();
        while b > 0 {
            if b & 1 != 0 {
                result = self.mul(result, a);
            }
            a = self.mul(a, a);
            b >>= 1;
        }

        // Check that r (reduced back to the usual representation) equals
        //  a^b % n, unless the latter computation overflows
        // Temporarily commented-out, as there u64::checked_pow is not available
        //  on the minimum supported Rust version, nor is an appropriate method
        //  for compiling the check conditionally.
        //debug_assert!(self
        //    .to_u64(_a)
        //    .checked_pow(_b as u32)
        //    .map(|r| r % self.modulus() == self.to_u64(result))
        //    .unwrap_or(true));

        result
    }

    fn one(&self) -> Self::I {
        self.from_u64(1)
    }
    fn minus_one(&self) -> Self::I {
        self.from_u64(self.modulus() - 1)
    }
    fn zero(&self) -> Self::I {
        self.from_u64(0)
    }
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct Montgomery {
    a: u64,
    n: u64,
}

impl Montgomery {
    /// computes x/R mod n efficiently
    fn reduce(&self, x: u128) -> u64 {
        debug_assert!(x < (self.n as u128) << 64);
        // TODO: optimiiiiiiise
        let Montgomery { a, n } = self;
        let m = (x as u64).wrapping_mul(*a);
        let nm = (*n as u128) * (m as u128);
        let (xnm, overflow) = (x as u128).overflowing_add(nm); // x + n*m
        debug_assert_eq!(xnm % (1 << 64), 0);

        // (x + n*m) / R
        // in case of overflow, this is (2¹²⁸ + xnm)/2⁶⁴ - n = xnm/2⁶⁴ + (2⁶⁴ - n)
        let y = (xnm >> 64) as u64 + if !overflow { 0 } else { n.wrapping_neg() };

        if y >= *n {
            y - n
        } else {
            y
        }
    }
}

impl Arithmetic for Montgomery {
    // Montgomery transform, R=2⁶⁴
    // Provides fast arithmetic mod n (n odd, u64)
    type I = u64;

    fn new(n: u64) -> Self {
        let a = inv_mod_u64(n).wrapping_neg();
        debug_assert_eq!(n.wrapping_mul(a), 1_u64.wrapping_neg());
        Montgomery { a, n }
    }

    fn modulus(&self) -> u64 {
        self.n
    }

    fn from_u64(&self, x: u64) -> Self::I {
        // TODO: optimise!
        assert!(x < self.n);
        let r = (((x as u128) << 64) % self.n as u128) as u64;
        debug_assert_eq!(x, self.to_u64(r));
        r
    }

    fn to_u64(&self, n: Self::I) -> u64 {
        self.reduce(n as u128)
    }

    fn add(&self, a: Self::I, b: Self::I) -> Self::I {
        let (r, overflow) = a.overflowing_add(b);

        // In case of overflow, a+b = 2⁶⁴ + r = (2⁶⁴ - n) + r (working mod n)
        let r = if !overflow {
            r
        } else {
            r + self.n.wrapping_neg()
        };

        // Normalise to [0; n[
        let r = if r < self.n { r } else { r - self.n };

        // Check that r (reduced back to the usual representation) equals
        // a+b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 = (((a_r as u128) + (b_r as u128)) % (self.n as u128)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} + {} = [{}] + [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }

    fn mul(&self, a: Self::I, b: Self::I) -> Self::I {
        let r = self.reduce((a as u128) * (b as u128));

        // Check that r (reduced back to the usual representation) equals
        // a*b % n
        #[cfg(debug_assertions)]
        {
            let a_r = self.to_u64(a);
            let b_r = self.to_u64(b);
            let r_r = self.to_u64(r);
            let r_2 = (((a_r as u128) * (b_r as u128)) % (self.n as u128)) as u64;
            debug_assert_eq!(
                r_r, r_2,
                "[{}] = {} ≠ {} = {} * {} = [{}] * [{}] mod {}; a = {}",
                r, r_r, r_2, a_r, b_r, a, b, self.n, self.a
            );
        }
        r
    }
}

// extended Euclid algorithm
// precondition: a is odd
pub(crate) fn inv_mod_u64(a: u64) -> u64 {
    assert!(a % 2 == 1);
    let mut t = 0u64;
    let mut newt = 1u64;
    let mut r = 0u64;
    let mut newr = a;

    while newr != 0 {
        let quot = if r == 0 {
            // special case when we're just starting out
            // This works because we know that
            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
            std::u64::MAX
        } else {
            r
        } / newr;

        let newtp = t.wrapping_sub(quot.wrapping_mul(newt));
        t = newt;
        newt = newtp;

        let newrp = r.wrapping_sub(quot.wrapping_mul(newr));
        r = newr;
        newr = newrp;
    }

    assert_eq!(r, 1);
    t
}

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

    #[test]
    fn test_inverter() {
        // All odd integers from 1 to 20 000
        let mut test_values = (0..10_000u64).map(|i| 2 * i + 1);

        assert!(test_values.all(|x| x.wrapping_mul(inv_mod_u64(x)) == 1));
    }

    #[test]
    fn test_montgomery_add() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    println!("{n:?}, {x:?}, {y:?}", n = n, x = x, y = y);
                    assert_eq!((x + y) % n, m.to_u64(m.add(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_montgomery_mult() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let m_x = m.from_u64(x);
                for y in 0..=x {
                    let m_y = m.from_u64(y);
                    assert_eq!((x * y) % n, m.to_u64(m.mul(m_x, m_y)));
                }
            }
        }
    }

    #[test]
    fn test_montgomery_roundtrip() {
        for n in 0..100 {
            let n = 2 * n + 1;
            let m = Montgomery::new(n);
            for x in 0..n {
                let x_ = m.from_u64(x);
                assert_eq!(x, m.to_u64(x_));
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate sled;
extern crate libc;

use std::fs;
use std::thread;
use std::time::Duration;

const CYCLE: usize = 16; // 65536;

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        None => return (0, 0),
        Some((_k, v)) => slice_to_u32(&*v),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for (mut k, v) in iter {
        if v == highest_vec {
            contiguous += 1;
        } else {
            k.reverse();
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    let lowest_vec = u32_to_vec(lowest);

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for (mut k, v) in tree.scan(&*low_beginning) {
        if v != lowest_vec {
            k.reverse();
        }
        assert_eq!(
            v,
            lowest_vec,
            "expected key {} to have value {}, instead it had value {}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v)
        );
    }

    (contiguous, highest)
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; 4] = unsafe { std::mem::transmute(u) };
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; 4];
    buf.copy_from_slice(b);

    unsafe { std::mem::transmute(buf) }
}

fn run(config: sled::FinalConfig) {
    let tree = config.tree();

    // flush to ensure the initial root is stable.
    tree.flush();

    let (key, highest) = verify(&tree);

    thread::spawn(|| {
        thread::sleep(Duration::from_millis(100));
        unsafe {
            libc::raise(9);
        }
    });

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let mut key = u32_to_vec((hu % CYCLE) as u32);
        key.reverse();
        let value = u32_to_vec((hu / CYCLE) as u32);
        tree.set(key, value);
    }
}

fn run_without_snapshot() {
    let config = sled::Config::default()
        .io_bufs(2)
        .blink_fanout(15)
        .page_consolidation_threshold(10)
        .cache_fixup_threshold(1)
        .cache_bits(6)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        // drop io_buf_size to 1<<16, then 1<<17 to tease out
        // low hanging fruit more quickly
        .io_buf_size(100_000) // 1<<16 is 65k but might cause stalling
        .path("test_crashes".to_string())
        .snapshot_after_ops(1 << 56)
        .build();

    run(config);
}

fn run_with_snapshot() {
    let config = sled::Config::default()
        .io_bufs(2)
        .blink_fanout(15)
        .page_consolidation_threshold(10)
        .cache_fixup_threshold(1)
        .cache_bits(6)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        // drop io_buf_size to 1<<16, then 1<<17 to tease out
        // low hanging fruit more quickly
        .io_buf_size(100_000) // 1<<16 is 65k but might cause stalling
        .path("test_crashes_with_snapshot".to_string())
        .snapshot_after_ops(1 << 10)
        .build();

    run(config);
}

#[test]
fn test_crash_recovery_with_runtime_snapshot() {
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_with_snapshot()
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup();
                panic!("child exited abnormally");
            }
        }
    }
    cleanup();
}

#[test]
fn test_crash_recovery_no_runtime_snapshot() {
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_without_snapshot()
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup();
                panic!("child exited abnormally");
            }
        }
    }
    cleanup();
}

fn cleanup() {
    let _res = fs::remove_file("test_crashes.db");
    let _res = fs::remove_file("test_crashes.conf");
    for dir_entry in fs::read_dir(".").unwrap() {
        if let Ok(de) = dir_entry {
            let path_buf = de.path();
            let path = path_buf.as_path();
            let path_str = path.to_str().unwrap();
            if path_str.starts_with("./test_crashes.snap.") {
                let _res = fs::remove_file(path);
            }
        }
    }
}
<commit_msg>properly clean up crash test files<commit_after>extern crate sled;
extern crate libc;

use std::fs;
use std::thread;
use std::time::Duration;

const CYCLE: usize = 16; // 65536;

/// Verifies that the keys in the tree are correctly recovered.
/// Panics if they are incorrect.
/// Returns the key that should be resumed at, and the current cycle value.
fn verify(tree: &sled::Tree) -> (u32, u32) {
    // key 0 should always be the highest value, as that's where we increment
    // at some point, it might go down by one
    // it should never return, or go down again after that
    let mut iter = tree.iter();
    let highest = match iter.next() {
        None => return (0, 0),
        Some((_k, v)) => slice_to_u32(&*v),
    };

    let highest_vec = u32_to_vec(highest);

    // find how far we got
    let mut contiguous: u32 = 0;
    let mut lowest = 0;
    for (mut k, v) in iter {
        if v == highest_vec {
            contiguous += 1;
        } else {
            k.reverse();
            let expected = if highest == 0 {
                CYCLE as u32 - 1
            } else {
                (highest - 1) % CYCLE as u32
            };
            let actual = slice_to_u32(&*v);
            assert_eq!(expected, actual);
            lowest = actual;
            break;
        }
    }

    let lowest_vec = u32_to_vec(lowest);

    // ensure nothing changes after this point
    let low_beginning = u32_to_vec(contiguous + 1);

    for (mut k, v) in tree.scan(&*low_beginning) {
        if v != lowest_vec {
            k.reverse();
        }
        assert_eq!(
            v,
            lowest_vec,
            "expected key {} to have value {}, instead it had value {}",
            slice_to_u32(&*k),
            lowest,
            slice_to_u32(&*v)
        );
    }

    (contiguous, highest)
}

fn u32_to_vec(u: u32) -> Vec<u8> {
    let buf: [u8; 4] = unsafe { std::mem::transmute(u) };
    buf.to_vec()
}

fn slice_to_u32(b: &[u8]) -> u32 {
    let mut buf = [0u8; 4];
    buf.copy_from_slice(b);

    unsafe { std::mem::transmute(buf) }
}

fn run(config: sled::FinalConfig) {
    let tree = config.tree();

    // flush to ensure the initial root is stable.
    tree.flush();

    let (key, highest) = verify(&tree);

    thread::spawn(|| {
        thread::sleep(Duration::from_millis(100));
        unsafe {
            libc::raise(9);
        }
    });

    let mut hu = ((highest as usize) * CYCLE) + key as usize;
    assert_eq!(hu % CYCLE, key as usize);
    assert_eq!(hu / CYCLE, highest as usize);

    loop {
        hu += 1;

        if hu / CYCLE >= CYCLE {
            hu = 0;
        }

        let mut key = u32_to_vec((hu % CYCLE) as u32);
        key.reverse();
        let value = u32_to_vec((hu / CYCLE) as u32);
        tree.set(key, value);
    }
}

fn run_without_snapshot() {
    let config = sled::Config::default()
        .io_bufs(2)
        .blink_fanout(15)
        .page_consolidation_threshold(10)
        .cache_fixup_threshold(1)
        .cache_bits(6)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        // drop io_buf_size to 1<<16, then 1<<17 to tease out
        // low hanging fruit more quickly
        .io_buf_size(100_000) // 1<<16 is 65k but might cause stalling
        .path("test_crashes".to_string())
        .snapshot_after_ops(1 << 56)
        .build();

    run(config);
}

fn run_with_snapshot() {
    let config = sled::Config::default()
        .io_bufs(2)
        .blink_fanout(15)
        .page_consolidation_threshold(10)
        .cache_fixup_threshold(1)
        .cache_bits(6)
        .cache_capacity(128 * 1024 * 1024)
        .flush_every_ms(Some(100))
        // drop io_buf_size to 1<<16, then 1<<17 to tease out
        // low hanging fruit more quickly
        .io_buf_size(100_000) // 1<<16 is 65k but might cause stalling
        .path("test_crashes_with_snapshot".to_string())
        .snapshot_after_ops(1 << 10)
        .build();

    run(config);
}

#[test]
fn test_crash_recovery_with_runtime_snapshot() {
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_with_snapshot()
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup();
                panic!("child exited abnormally");
            }
        }
    }
    cleanup_with_snapshots();
}

#[test]
fn test_crash_recovery_no_runtime_snapshot() {
    for _ in 0..100 {
        let child = unsafe { libc::fork() };
        if child == 0 {
            run_without_snapshot()
        } else {
            let mut status = 0;
            unsafe {
                libc::waitpid(child, &mut status as *mut libc::c_int, 0);
            }
            if status != 9 {
                cleanup();
                panic!("child exited abnormally");
            }
        }
    }
    cleanup();
}

fn cleanup_with_snapshots() {
    let _res = fs::remove_file("test_crashes_with_snapshot.db");
    let _res = fs::remove_file("test_crashes_with_snapshot.conf");
    for dir_entry in fs::read_dir(".").unwrap() {
        if let Ok(de) = dir_entry {
            let path_buf = de.path();
            let path = path_buf.as_path();
            let path_str = path.to_str().unwrap();
            if path_str.starts_with("./test_crashes_with_snapshot.snap.") {
                let _res = fs::remove_file(path);
            }
        }
    }
}

fn cleanup() {
    let _res = fs::remove_file("test_crashes.db");
    let _res = fs::remove_file("test_crashes.conf");
    for dir_entry in fs::read_dir(".").unwrap() {
        if let Ok(de) = dir_entry {
            let path_buf = de.path();
            let path = path_buf.as_path();
            let path_str = path.to_str().unwrap();
            if path_str.starts_with("./test_crashes.snap.") {
                let _res = fs::remove_file(path);
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Test data production using high level producers.

use std::collections::HashMap;
use std::time::{Duration, Instant};

use futures::stream::{FuturesUnordered, StreamExt};

use rdkafka::client::DefaultClientContext;
use rdkafka::config::ClientConfig;
use rdkafka::error::{KafkaError, RDKafkaError};
use rdkafka::message::{Headers, Message, OwnedHeaders};
use rdkafka::producer::future_producer::FutureRecord;
use rdkafka::producer::FutureProducer;

use crate::utils::*;

mod utils;

fn future_producer(config_overrides: HashMap<&str, &str>) -> FutureProducer<DefaultClientContext> {
    let mut config = ClientConfig::new();
    config
        .set("bootstrap.servers", "localhost")
        .set("message.timeout.ms", "5000");
    for (key, value) in config_overrides {
        config.set(key, value);
    }
    config.create().expect("Failed to create producer")
}

#[tokio::test]
async fn test_future_producer_send() {
    let producer = future_producer(HashMap::new());
    let topic_name = rand_test_topic();

    let results: FuturesUnordered<_> = (0..10)
        .map(|_| {
            producer.send(
                FutureRecord::to(&topic_name).payload("A").key("B"),
                Duration::from_secs(0),
            )
        })
        .collect();

    let results: Vec<_> = results.collect().await;
    assert!(results.len() == 10);
    for (i, result) in results.into_iter().enumerate() {
        let (partition, offset) = result.unwrap();
        assert_eq!(partition, 1);
        assert_eq!(offset, i as i64);
    }
}

#[tokio::test]
async fn test_future_producer_send_full() {
    // Connect to a nonexistent Kafka broker with no message timeout and a tiny
    // producer queue, so we can permanently fill up the queue by sending a
    // single message.
    let mut config = HashMap::new();
    config.insert("bootstrap.servers", "");
    config.insert("message.timeout.ms", "0");
    config.insert("queue.buffering.max.messages", "1");
    let producer = &future_producer(config);
    let topic_name = &rand_test_topic();

    // Fill up the queue.
    producer
        .send_result(FutureRecord::to(&topic_name).payload("A").key("B"))
        .unwrap();

    let send_message = |timeout| async move {
        let start = Instant::now();
        let res = producer
            .send(FutureRecord::to(&topic_name).payload("A").key("B"), timeout)
            .await;
        match res {
            Ok(_) => panic!("send unexpectedly succeeded"),
            Err((KafkaError::MessageProduction(RDKafkaError::QueueFull), _)) => start.elapsed(),
            Err((e, _)) => panic!("got incorrect error: {}", e),
        }
    };

    // Sending a message with no timeout should return a `QueueFull` error
    // approximately immediately.
    let elapsed = send_message(Duration::from_secs(0)).await;
    assert!(elapsed < Duration::from_millis(10));

    // Sending a message with a 1s timeout should return a `QueueFull` error
    // in about 1s.
    let elapsed = send_message(Duration::from_secs(1)).await;
    assert!(elapsed > Duration::from_millis(800));
    assert!(elapsed < Duration::from_millis(1200));
}

#[tokio::test]
async fn test_future_producer_send_fail() {
    let producer = future_producer(HashMap::new());

    let future = producer.send(
        FutureRecord::to("topic")
            .payload("payload")
            .key("key")
            .partition(100) // Fail
            .headers(
                OwnedHeaders::new()
                    .add("0", "A")
                    .add("1", "B")
                    .add("2", "C"),
            ),
        Duration::from_secs(10),
    );

    match future.await {
        Err((kafka_error, owned_message)) => {
            assert_eq!(
                kafka_error.to_string(),
                "Message production error: UnknownPartition (Local: Unknown partition)"
            );
            assert_eq!(owned_message.topic(), "topic");
            let headers = owned_message.headers().unwrap();
            assert_eq!(headers.count(), 3);
            assert_eq!(headers.get_as::<str>(0).unwrap(), ("0", Ok("A")));
            assert_eq!(headers.get_as::<str>(1).unwrap(), ("1", Ok("B")));
            assert_eq!(headers.get_as::<str>(2).unwrap(), ("2", Ok("C")));
        }
        e => {
            panic!("Unexpected return value: {:?}", e);
        }
    }
}
<commit_msg>Fix test_future_producer_send_full under Valgrind<commit_after>//! Test data production using high level producers.

use std::collections::HashMap;
use std::time::{Duration, Instant};

use futures::stream::{FuturesUnordered, StreamExt};

use rdkafka::client::DefaultClientContext;
use rdkafka::config::ClientConfig;
use rdkafka::error::{KafkaError, RDKafkaError};
use rdkafka::message::{Headers, Message, OwnedHeaders};
use rdkafka::producer::future_producer::FutureRecord;
use rdkafka::producer::FutureProducer;
use rdkafka::util::Timeout;

use crate::utils::*;

mod utils;

fn future_producer(config_overrides: HashMap<&str, &str>) -> FutureProducer<DefaultClientContext> {
    let mut config = ClientConfig::new();
    config
        .set("bootstrap.servers", "localhost")
        .set("message.timeout.ms", "5000");
    for (key, value) in config_overrides {
        config.set(key, value);
    }
    config.create().expect("Failed to create producer")
}

#[tokio::test]
async fn test_future_producer_send() {
    let producer = future_producer(HashMap::new());
    let topic_name = rand_test_topic();

    let results: FuturesUnordered<_> = (0..10)
        .map(|_| {
            producer.send(
                FutureRecord::to(&topic_name).payload("A").key("B"),
                Duration::from_secs(0),
            )
        })
        .collect();

    let results: Vec<_> = results.collect().await;
    assert!(results.len() == 10);
    for (i, result) in results.into_iter().enumerate() {
        let (partition, offset) = result.unwrap();
        assert_eq!(partition, 1);
        assert_eq!(offset, i as i64);
    }
}

#[tokio::test]
async fn test_future_producer_send_full() {
    // Connect to a nonexistent Kafka broker with a long message timeout and a
    // tiny producer queue, so we can fill up the queue for a while by sending a
    // single message.
    let mut config = HashMap::new();
    config.insert("bootstrap.servers", "");
    config.insert("message.timeout.ms", "5000");
    config.insert("queue.buffering.max.messages", "1");
    let producer = &future_producer(config);
    let topic_name = &rand_test_topic();

    // Fill up the queue.
    producer
        .send_result(FutureRecord::to(&topic_name).payload("A").key("B"))
        .unwrap();

    let send_message = |timeout| async move {
        let start = Instant::now();
        let res = producer
            .send(FutureRecord::to(&topic_name).payload("A").key("B"), timeout)
            .await;
        match res {
            Ok(_) => panic!("send unexpectedly succeeded"),
            Err((KafkaError::MessageProduction(RDKafkaError::QueueFull), _)) => start.elapsed(),
            Err((e, _)) => panic!("got incorrect error: {}", e),
        }
    };

    // Sending a message with no timeout should return a `QueueFull` error
    // approximately immediately.
    let elapsed = send_message(Duration::from_secs(0)).await;
    assert!(elapsed < Duration::from_millis(20));

    // Sending a message with a 1s timeout should return a `QueueFull` error
    // in about 1s.
    let elapsed = send_message(Duration::from_secs(1)).await;
    assert!(elapsed > Duration::from_millis(800));
    assert!(elapsed < Duration::from_millis(1200));

    producer.flush(Timeout::Never);
}

#[tokio::test]
async fn test_future_producer_send_fail() {
    let producer = future_producer(HashMap::new());

    let future = producer.send(
        FutureRecord::to("topic")
            .payload("payload")
            .key("key")
            .partition(100) // Fail
            .headers(
                OwnedHeaders::new()
                    .add("0", "A")
                    .add("1", "B")
                    .add("2", "C"),
            ),
        Duration::from_secs(10),
    );

    match future.await {
        Err((kafka_error, owned_message)) => {
            assert_eq!(
                kafka_error.to_string(),
                "Message production error: UnknownPartition (Local: Unknown partition)"
            );
            assert_eq!(owned_message.topic(), "topic");
            let headers = owned_message.headers().unwrap();
            assert_eq!(headers.count(), 3);
            assert_eq!(headers.get_as::<str>(0).unwrap(), ("0", Ok("A")));
            assert_eq!(headers.get_as::<str>(1).unwrap(), ("1", Ok("B")));
            assert_eq!(headers.get_as::<str>(2).unwrap(), ("2", Ok("C")));
        }
        e => {
            panic!("Unexpected return value: {:?}", e);
        }
    }
}
<|endoftext|>


<commit_before>use graph::{BasicBlockIndex, FuncGraph};
use graph_algorithms::Graph;
use graph_algorithms::dominators::{self, Dominators, DominatorTree};
use graph_algorithms::iterate::reverse_post_order;
use graph_algorithms::loop_tree::{self, LoopTree};
use graph_algorithms::reachable::{self, Reachability};
use graph_algorithms::transpose::TransposedGraph;
use std::fmt;

pub struct Environment<'func, 'arena: 'func> {
    pub graph: &'func FuncGraph<'arena>,
    pub dominators: Dominators<FuncGraph<'arena>>,
    pub dominator_tree: DominatorTree<FuncGraph<'arena>>,
    pub postdominators: Dominators<TransposedGraph<&'func FuncGraph<'arena>>>,
    pub reachable: Reachability<FuncGraph<'arena>>,
    pub loop_tree: LoopTree<FuncGraph<'arena>>,
    pub reverse_post_order: Vec<BasicBlockIndex>,
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Point {
    pub block: BasicBlockIndex,
    pub action: usize,
}

impl<'func, 'arena> Environment<'func, 'arena> {
    pub fn new(graph: &'func FuncGraph<'arena>) -> Self {
        let rpo = reverse_post_order(graph, graph.start_node());
        let dominators = dominators::dominators_given_rpo(graph, &rpo);
        let dominator_tree = dominators.dominator_tree();
        let reachable = reachable::reachable_given_rpo(graph, &rpo);
        let loop_tree = loop_tree::loop_tree_given(graph, &dominators);

        let postdominators = {
            let exit = graph.block_index_str("EXIT");
            let transpose = &TransposedGraph::with_start(graph, exit);
            dominators::dominators(transpose)
        };

        Environment {
            graph: graph,
            dominators: dominators,
            dominator_tree: dominator_tree,
            postdominators: postdominators,
            reachable: reachable,
            loop_tree: loop_tree,
            reverse_post_order: rpo,
        }
    }

    pub fn dump_dominators(&self) {
        let tree = self.dominators.dominator_tree();
        self.dump_dominator_tree(&tree, tree.root(), 0)
    }

    pub fn dump_postdominators(&self) {
        let tree = self.postdominators.dominator_tree();
        self.dump_dominator_tree(&tree, tree.root(), 0)
    }

    fn dump_dominator_tree<G1>(&self,
                               tree: &DominatorTree<G1>,
                               node: BasicBlockIndex,
                               indent: usize)
        where G1: Graph<Node=BasicBlockIndex>
    {
        println!("{0:1$}- {2:?}",
                 "",
                 indent,
                 node);

        for &child in tree.children(node) {
            self.dump_dominator_tree(tree, child, indent + 2)
        }
    }

    pub fn last_action(&self, block: BasicBlockIndex) -> usize {
        self.graph.block_data(block).actions.len()
    }

    pub fn end_action(&self, block: BasicBlockIndex) -> usize {
        self.last_action(block) + 1
    }

    pub fn end_point(&self, block: BasicBlockIndex) -> Point {
        Point {
            block: block,
            action: self.end_action(block)
        }
    }

    pub fn can_reach(&self, pred: Point, succ: Point) -> bool {
        if pred.block == succ.block {
            pred.action <= succ.action
        } else {
            self.reachable.can_reach(pred.block, succ.block)
        }
    }
}

impl fmt::Debug for Point {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "({:?} @ {})", self.block, self.action)
    }
}
<commit_msg>match debug to input format<commit_after>use graph::{BasicBlockIndex, FuncGraph};
use graph_algorithms::Graph;
use graph_algorithms::dominators::{self, Dominators, DominatorTree};
use graph_algorithms::iterate::reverse_post_order;
use graph_algorithms::loop_tree::{self, LoopTree};
use graph_algorithms::reachable::{self, Reachability};
use graph_algorithms::transpose::TransposedGraph;
use std::fmt;

pub struct Environment<'func, 'arena: 'func> {
    pub graph: &'func FuncGraph<'arena>,
    pub dominators: Dominators<FuncGraph<'arena>>,
    pub dominator_tree: DominatorTree<FuncGraph<'arena>>,
    pub postdominators: Dominators<TransposedGraph<&'func FuncGraph<'arena>>>,
    pub reachable: Reachability<FuncGraph<'arena>>,
    pub loop_tree: LoopTree<FuncGraph<'arena>>,
    pub reverse_post_order: Vec<BasicBlockIndex>,
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Point {
    pub block: BasicBlockIndex,
    pub action: usize,
}

impl<'func, 'arena> Environment<'func, 'arena> {
    pub fn new(graph: &'func FuncGraph<'arena>) -> Self {
        let rpo = reverse_post_order(graph, graph.start_node());
        let dominators = dominators::dominators_given_rpo(graph, &rpo);
        let dominator_tree = dominators.dominator_tree();
        let reachable = reachable::reachable_given_rpo(graph, &rpo);
        let loop_tree = loop_tree::loop_tree_given(graph, &dominators);

        let postdominators = {
            let exit = graph.block_index_str("EXIT");
            let transpose = &TransposedGraph::with_start(graph, exit);
            dominators::dominators(transpose)
        };

        Environment {
            graph: graph,
            dominators: dominators,
            dominator_tree: dominator_tree,
            postdominators: postdominators,
            reachable: reachable,
            loop_tree: loop_tree,
            reverse_post_order: rpo,
        }
    }

    pub fn dump_dominators(&self) {
        let tree = self.dominators.dominator_tree();
        self.dump_dominator_tree(&tree, tree.root(), 0)
    }

    pub fn dump_postdominators(&self) {
        let tree = self.postdominators.dominator_tree();
        self.dump_dominator_tree(&tree, tree.root(), 0)
    }

    fn dump_dominator_tree<G1>(&self,
                               tree: &DominatorTree<G1>,
                               node: BasicBlockIndex,
                               indent: usize)
        where G1: Graph<Node=BasicBlockIndex>
    {
        println!("{0:1$}- {2:?}",
                 "",
                 indent,
                 node);

        for &child in tree.children(node) {
            self.dump_dominator_tree(tree, child, indent + 2)
        }
    }

    pub fn last_action(&self, block: BasicBlockIndex) -> usize {
        self.graph.block_data(block).actions.len()
    }

    pub fn end_action(&self, block: BasicBlockIndex) -> usize {
        self.last_action(block) + 1
    }

    pub fn end_point(&self, block: BasicBlockIndex) -> Point {
        Point {
            block: block,
            action: self.end_action(block)
        }
    }

    pub fn can_reach(&self, pred: Point, succ: Point) -> bool {
        if pred.block == succ.block {
            pred.action <= succ.action
        } else {
            self.reachable.can_reach(pred.block, succ.block)
        }
    }
}

impl fmt::Debug for Point {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "{:?}/{}", self.block, self.action)
    }
}
<|endoftext|>