<commit_before>use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::{self, File};
use std::path::Path;
use std::error::Error;
use std::rc::Rc;
use std::cell::RefCell;

use collada::{self, COLLADA, GeometricElement, ArrayElement, PrimitiveType, VisualScene, Geometry,
              Node};

use polygon::gl_render::{GLRender, GLMeshData};
use polygon::geometry::mesh::Mesh;

use wav::Wave;
use scene::Scene;
use ecs::Entity;
use component::{MeshManager, TransformManager};

#[derive(Debug, Clone)]
pub struct ResourceManager {
    renderer: Rc<GLRender>,
    meshes: RefCell<HashMap<String, GLMeshData>>,
    audio_clips: RefCell<HashMap<String, Rc<Wave>>>,

    visual_scenes: RefCell<HashMap<String, VisualScene>>,
    geometries: RefCell<HashMap<String, Geometry>>,
}

impl ResourceManager {
    pub fn new(renderer: Rc<GLRender>) -> ResourceManager {
        ResourceManager {
            renderer: renderer,
            meshes: RefCell::new(HashMap::new()),
            audio_clips: RefCell::new(HashMap::new()),
            visual_scenes: RefCell::new(HashMap::new()),
            geometries: RefCell::new(HashMap::new()),
        }
    }

    /// TODO: Perform validity checking on data when loading (e.g. make sure all nodes have an id).
    pub fn load_model<P: AsRef<Path>>(&self, path: P) -> Result<(), String> {
        let mut visual_scenes = self.visual_scenes.borrow_mut();
        let mut geometries = self.geometries.borrow_mut();

        let path_ref = path.as_ref();
        let metadata = match fs::metadata(path_ref) {
            Err(why) => return Err(format!(
                "Unable to read metadata for {}, either it doesn't exist or the user lacks permissions, {}",
                path_ref.display(),
                &why)),
            Ok(metadata) => metadata,
        };
        if !metadata.is_file() {
            return Err(format!(
                "{} could not be loaded because it is not a file",
                path_ref.display()));
        }
        let collada_data = match COLLADA::load(path_ref) {
            Err(why) => return Err(format!(
                "couldn't open {}: {}",
                path_ref.display(),
                &why)),
            Ok(data) => data,
        };

        // Store each of the visual scenes from the collada file.
        for visual_scene in collada_data.library_visual_scenes.as_ref().unwrap().visual_scenes.iter() {
            let id = match visual_scene.id {
                None => return Err(format!(
                    "COLLADA file {} contained a <visual_scene> with no \"id\" attribute, this is unsupported.",
                    path_ref.display())),
                Some(ref id) => id.clone(),
            };
            visual_scenes.insert(id, visual_scene.clone());
        }

        // Store each of the geometries so they can be referenced later.
        for geometry in collada_data.library_geometries.as_ref().unwrap().geometries.iter() {
            let id = match geometry.id {
                None => return Err(format!(
                    "COLLADA file {} contained a <geometry> with no \"id\" attribute, this is unsupported",
                    path_ref.display())),
                Some(ref id) => id.clone(),
            };
            geometries.insert(id, geometry.clone());
        }

        Ok(())
    }

    pub fn get_mesh(&self, uri: &str) -> Result<GLMeshData, String> {
        let mut meshes = self.meshes.borrow_mut();

        // Use cached mesh data if possible.
        if let Some(mesh) = meshes.get(uri) {
            return Ok(*mesh);
        }

        // Generate mesh data since none has ben created previously.
        let visual_scenes = self.visual_scenes.borrow();

        // TODO: Handle invalid URIs (empty, invalid characters?).
        let mut uri_segments = uri.split(".");
        let root = uri_segments.next().unwrap();
        let visual_scene = match visual_scenes.get(root) {
            None => return Err(format!(
                "No source file {} found from which to load {}",
                root,
                uri)),
            Some(visual_scene) => visual_scene,
        };

        // Get the first node in the URI.
        let mut node = {
            let name = uri_segments.next().unwrap();
            let mut result: Option<&Node> = None;
            for node in &visual_scene.nodes {
                if node.id.as_ref().unwrap() == name {
                    result = Some(node);
                    break;
                }
            }

            match result {
                None => return Err(format!(
                    "No node named {} found in scene {}",
                    name,
                    root)),
                Some(node) => node,
            }
        };

        for name in uri_segments {
            let mut result: Option<&Node> = None;
            for next_node in &node.nodes {
                if next_node.id.as_ref().unwrap() == name {
                    result = Some(next_node);
                    break;
                }
            }

            match result {
                None => return Err(format!(
                    "No node named {} found when parsing {}",
                    name,
                    uri)),
                Some(next_node) =>
                    node = next_node,
            }
        }

        let mesh_data = self.gen_mesh_from_node(node, uri).unwrap();
        meshes.insert(uri.to_string(), mesh_data);
        Ok(mesh_data)
    }

    pub fn get_audio_clip(&self, path_text: &str) -> Rc<Wave> {
        let mut audio_clips = self.audio_clips.borrow_mut();

        if !audio_clips.contains_key(path_text) {
            let wave = Wave::from_file(path_text).unwrap();
            audio_clips.insert(path_text.into(), Rc::new(wave));
        }

        audio_clips.get(path_text).unwrap().clone()
    }

    pub fn instantiate_model(&self, resource: &str, scene: &Scene) -> Result<Entity, String> {
        if resource.contains(".") {
            println!("WARNING: ResourceManager::instantiate_model() doesn't yet support fully qualified URIs, only root assets may be instantiated.");
        }

        let mut uri_segments = resource.split(".");
        let root = uri_segments.next().unwrap();
        let visual_scenes = self.visual_scenes.borrow();
        let visual_scene = {
            match visual_scenes.get(root) {
                None => return Err(format!(
                    "No source file {} found from which to load {}",
                    root,
                    resource)),
                Some(visual_scene) => visual_scene,
            }
        };

        let node = {
            if visual_scene.nodes.len() == 0 {
                return Err(format!(
                    "No nodes associated with model {}",
                    resource));
            }

            if visual_scene.nodes.len() > 1 {
                println!(
                    "WARNING: Model {} has more than one node at the root level. This is not currenlty supported, only the first node will be used.",
                    resource);
            }

            &visual_scene.nodes[0]
        };

        let mut uri = String::from(resource);
        uri.push_str(".");
        uri.push_str(node.id.as_ref().unwrap());

        if let Ok(mesh) = self.get_mesh(&uri) {
            let entity = scene.create_entity();
            let mut transform_manager = scene.get_manager_mut::<TransformManager>();
            transform_manager.assign(entity);
            scene.get_manager_mut::<MeshManager>().give_mesh(entity, mesh);

            return Ok(entity);
        }

        panic!("Why wasn't {} loaded yet?", uri);
    }

    fn gen_mesh_from_node(&self, node: &collada::Node, uri: &str) -> Result<GLMeshData, String> {
        let geometry_name = {
            if node.instance_geometries.len() == 0 {
                return Err(format!("No geometry is identified by {}", uri));
            }
            if node.instance_geometries.len() > 1 {
                return Err(format!("More than one geometry is identified by {}", uri));
            }

            let url = &node.instance_geometries[0].url;
            &url[1..] // Skip the leading "#" character that starts all URLs.
        };

        let geometries = self.geometries.borrow();
        let geometry = geometries.get(geometry_name).unwrap();
        let mesh = geometry_to_mesh(geometry);

        let frag_src = load_file_text("shaders/forward_phong.frag.glsl");
        let vert_src = load_file_text("shaders/forward_phong.vert.glsl");

        let mesh_data =
            self.renderer.gen_mesh(&mesh, vert_src.as_ref(), frag_src.as_ref());

        Ok(mesh_data)
    }
}

/// Load the mesh data from a COLLADA .dae file.
///
/// The data in a COLLADA files is formatted for efficiency, but isn't necessarily
/// organized in a way that is supported by the graphics API. This method reformats the
/// data so that it can be sent straight to the GPU without further manipulation.
///
/// In order to to this, it reorganizes the normals, UVs, and other vertex attributes to
/// be in the same order as the vertex positions.
fn geometry_to_mesh(geometry: &Geometry) -> Mesh {
    let mesh = match geometry.data {
        GeometricElement::Mesh(ref mesh) => mesh,
        _ => panic!("No mesh found within geometry")
    };

    let position_data_raw = get_raw_positions(&mesh);
    let normal_data_raw = get_normals(&mesh);

    let triangles = match mesh.primitives[0] {
        PrimitiveType::Triangles(ref triangles) => triangles,
        _ => panic!("Only triangles primitives are supported currently")
    };
    let primitive_indices = &triangles.primitives;

    // Create a new array for the positions so we can add the w coordinate.
    let mut position_data: Vec<f32> = Vec::with_capacity(position_data_raw.len() / 3 * 4);

    // Create a new array for the normals and rearrange them to match the order of position attributes.
    let mut normal_data: Vec<f32> = Vec::with_capacity(position_data.len());

    // Iterate over the indices, rearranging the normal data to match the position data.
    let stride = triangles.inputs.len();
    let mut vertex_index_map: HashMap<(usize, usize), u32> = HashMap::new();
    let mut indices: Vec<u32> = Vec::new();
    let vertex_count = triangles.count * 3;
    let mut index_count = 0;
    for offset in 0..vertex_count {
        // Determine the offset of the the current vertex's attributes
        let position_index = primitive_indices[offset * stride];
        let normal_index = primitive_indices[offset * stride + 1];

        // Push the index of the vertex, either reusing an existing vertex or creating a new one.
        let vertex_key = (position_index, normal_index);
        let vertex_index = if vertex_index_map.contains_key(&vertex_key) {
            // Vertex has already been assembled, reuse the index.
            (*vertex_index_map.get(&vertex_key).unwrap()) as u32
        } else {
            // Assemble new vertex.
            let vertex_index = index_count;
            index_count += 1;
            vertex_index_map.insert(vertex_key, vertex_index as u32);

            // Append position to position data array.
            for offset in 0..3 {
                position_data.push(position_data_raw[position_index * 3 + offset]);
            }
            position_data.push(1.0);

            // Append normal to normal data array.
            for offset in 0..3 {
                normal_data.push(normal_data_raw[normal_index * 3 + offset]);
            }

            vertex_index
        };

        indices.push(vertex_index);
    }

    let mesh = Mesh::from_raw_data(position_data.as_ref(), normal_data.as_ref(), indices.as_ref());

    mesh
}

fn get_raw_positions(mesh: &collada::Mesh) -> &[f32] {
    // TODO: Consult the correct element (<triangles> for now) to determine which source has position data.
    let position_data: &[f32] = match mesh.sources[0].array_element {
        ArrayElement::Float(ref float_array) => float_array.as_ref(),
        _ => panic!("Only float arrays supported for vertex position array")
    };
    assert!(position_data.len() > 0);

    position_data
}

fn get_normals(mesh: &collada::Mesh) -> &[f32] {
    // TODO: Consult the correct element (<triangles> for now) to determine which source has normal data.
    let normal_data: &[f32] = match mesh.sources[1].array_element {
        ArrayElement::Float(ref float_array) => float_array.as_ref(),
        _ => panic!("Only float arrays supported for vertex normal array")
    };
    assert!(normal_data.len() > 0);

    normal_data
}

pub fn load_file_text(path: &str) -> String {
    let file_path = Path::new(path);
    let mut file = match File::open(&file_path) {
        // The `desc` field of `IoError` is a string that describes the error
        Err(why) => panic!("couldn't open {}: {}", file_path.display(), Error::description(&why)),
        Ok(file) => file,
    };
    let mut contents = String::new();
    match file.read_to_string(&mut contents) {
        Err(why) => panic!("couldn't read {}: {}", file_path.display(), Error::description(&why)),
        Ok(_) => ()
    }
    contents
}
<commit_msg>Mesh Generation Cleanup<commit_after>use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::{self, File};
use std::path::Path;
use std::error::Error;
use std::rc::Rc;
use std::cell::RefCell;

use collada::{self, COLLADA, GeometricElement, ArrayElement, PrimitiveType, VisualScene, Geometry,
              Node};

use polygon::gl_render::{GLRender, GLMeshData};
use polygon::geometry::mesh::Mesh;

use wav::Wave;
use scene::Scene;
use ecs::Entity;
use component::{MeshManager, TransformManager};

#[derive(Debug, Clone)]
pub struct ResourceManager {
    renderer: Rc<GLRender>,
    meshes: RefCell<HashMap<String, GLMeshData>>,
    audio_clips: RefCell<HashMap<String, Rc<Wave>>>,

    visual_scenes: RefCell<HashMap<String, VisualScene>>,
    geometries: RefCell<HashMap<String, Geometry>>,
}

impl ResourceManager {
    pub fn new(renderer: Rc<GLRender>) -> ResourceManager {
        ResourceManager {
            renderer: renderer,
            meshes: RefCell::new(HashMap::new()),
            audio_clips: RefCell::new(HashMap::new()),
            visual_scenes: RefCell::new(HashMap::new()),
            geometries: RefCell::new(HashMap::new()),
        }
    }

    /// TODO: Perform validity checking on data when loading (e.g. make sure all nodes have an id).
    pub fn load_model<P: AsRef<Path>>(&self, path: P) -> Result<(), String> {
        let mut visual_scenes = self.visual_scenes.borrow_mut();
        let mut geometries = self.geometries.borrow_mut();

        let path_ref = path.as_ref();
        let metadata = match fs::metadata(path_ref) {
            Err(why) => return Err(format!(
                "Unable to read metadata for {}, either it doesn't exist or the user lacks permissions, {}",
                path_ref.display(),
                &why)),
            Ok(metadata) => metadata,
        };
        if !metadata.is_file() {
            return Err(format!(
                "{} could not be loaded because it is not a file",
                path_ref.display()));
        }
        let collada_data = match COLLADA::load(path_ref) {
            Err(why) => return Err(format!(
                "couldn't open {}: {}",
                path_ref.display(),
                &why)),
            Ok(data) => data,
        };

        // Store each of the visual scenes from the collada file.
        for visual_scene in collada_data.library_visual_scenes.as_ref().unwrap().visual_scenes.iter() {
            let id = match visual_scene.id {
                None => return Err(format!(
                    "COLLADA file {} contained a <visual_scene> with no \"id\" attribute, this is unsupported.",
                    path_ref.display())),
                Some(ref id) => id.clone(),
            };
            visual_scenes.insert(id, visual_scene.clone());
        }

        // Store each of the geometries so they can be referenced later.
        for geometry in collada_data.library_geometries.as_ref().unwrap().geometries.iter() {
            let id = match geometry.id {
                None => return Err(format!(
                    "COLLADA file {} contained a <geometry> with no \"id\" attribute, this is unsupported",
                    path_ref.display())),
                Some(ref id) => id.clone(),
            };
            geometries.insert(id, geometry.clone());
        }

        Ok(())
    }

    pub fn get_mesh(&self, uri: &str) -> Result<GLMeshData, String> {
        // Use cached mesh data if possible.
        if let Some(mesh) = self.get_cached_mesh(uri) {
            return Ok(mesh);
        }

        // Generate mesh data since none has ben created previously.
        let visual_scenes = self.visual_scenes.borrow();

        // TODO: Handle invalid URIs (empty, invalid characters?).
        let mut uri_segments = uri.split(".");
        let root = uri_segments.next().unwrap();
        let visual_scene = match visual_scenes.get(root) {
            None => return Err(format!(
                "No source file {} found from which to load {}",
                root,
                uri)),
            Some(visual_scene) => visual_scene,
        };

        // Get the first node in the URI.
        let mut node = {
            let name = uri_segments.next().unwrap();
            let mut result: Option<&Node> = None;
            for node in &visual_scene.nodes {
                if node.id.as_ref().unwrap() == name {
                    result = Some(node);
                    break;
                }
            }

            match result {
                None => return Err(format!(
                    "No node named {} found in scene {}",
                    name,
                    root)),
                Some(node) => node,
            }
        };

        for name in uri_segments {
            let mut result: Option<&Node> = None;
            for next_node in &node.nodes {
                if next_node.id.as_ref().unwrap() == name {
                    result = Some(next_node);
                    break;
                }
            }

            match result {
                None => return Err(format!(
                    "No node named {} found when parsing {}",
                    name,
                    uri)),
                Some(next_node) =>
                    node = next_node,
            }
        }

        let mesh_data = self.gen_mesh_from_node(node, uri).unwrap();
        Ok(mesh_data)
    }

    pub fn get_audio_clip(&self, path_text: &str) -> Rc<Wave> {
        let mut audio_clips = self.audio_clips.borrow_mut();

        if !audio_clips.contains_key(path_text) {
            let wave = Wave::from_file(path_text).unwrap();
            audio_clips.insert(path_text.into(), Rc::new(wave));
        }

        audio_clips.get(path_text).unwrap().clone()
    }

    pub fn instantiate_model(&self, resource: &str, scene: &Scene) -> Result<Entity, String> {
        if resource.contains(".") {
            println!("WARNING: ResourceManager::instantiate_model() doesn't yet support fully qualified URIs, only root assets may be instantiated.");
        }

        let mut uri_segments = resource.split(".");
        let root = uri_segments.next().unwrap();
        let visual_scenes = self.visual_scenes.borrow();
        let visual_scene = {
            match visual_scenes.get(root) {
                None => return Err(format!(
                    "No source file {} found from which to load {}",
                    root,
                    resource)),
                Some(visual_scene) => visual_scene,
            }
        };

        let node = {
            if visual_scene.nodes.len() == 0 {
                return Err(format!(
                    "No nodes associated with model {}",
                    resource));
            }

            if visual_scene.nodes.len() > 1 {
                println!(
                    "WARNING: Model {} has more than one node at the root level. This is not currenlty supported, only the first node will be used.",
                    resource);
            }

            &visual_scene.nodes[0]
        };

        let mut uri = String::from(resource);
        uri.push_str(".");
        uri.push_str(node.id.as_ref().unwrap());

        let mesh_data = if let Some(mesh_data) = self.get_cached_mesh(&uri) {
            mesh_data
        } else {
            match self.gen_mesh_from_node(node, &uri) {
                Err(message) => return Err(message),
                Ok(mesh_data) => mesh_data,
            }
        };

        let entity = scene.create_entity();
        let mut transform_manager = scene.get_manager_mut::<TransformManager>();
        transform_manager.assign(entity);
        scene.get_manager_mut::<MeshManager>().give_mesh(entity, mesh_data);

        return Ok(entity);
    }

    fn gen_mesh_from_node(&self, node: &collada::Node, uri: &str) -> Result<GLMeshData, String> {
        let geometry_name = {
            if node.instance_geometries.len() == 0 {
                return Err(format!("No geometry is identified by {}", uri));
            }
            if node.instance_geometries.len() > 1 {
                return Err(format!("More than one geometry is identified by {}", uri));
            }

            let url = &node.instance_geometries[0].url;
            &url[1..] // Skip the leading "#" character that starts all URLs.
        };

        let geometries = self.geometries.borrow();
        let geometry = geometries.get(geometry_name).unwrap();
        self.gen_mesh(geometry, uri)
    }

    fn has_cached_mesh(&self, uri: &str) -> bool {
        self.meshes.borrow().contains_key(uri)
    }

    fn get_cached_mesh(&self, uri: &str) -> Option<GLMeshData> {
        match self.meshes.borrow().get(uri) {
            None => None,
            Some(mesh_ref) => Some(*mesh_ref),
        }
    }

    fn gen_mesh(&self, geometry: &Geometry, uri: &str) -> Result<GLMeshData, String> {
        assert!(!self.has_cached_mesh(uri), "Attempting to create a new mesh for {} when the uri is already in the meshes map", uri);

        let mesh = geometry_to_mesh(geometry);

        let frag_src = load_file_text("shaders/forward_phong.frag.glsl");
        let vert_src = load_file_text("shaders/forward_phong.vert.glsl");

        let mesh_data =
            self.renderer.gen_mesh(&mesh, vert_src.as_ref(), frag_src.as_ref());
        self.meshes.borrow_mut().insert(uri.into(), mesh_data);

        Ok(mesh_data)
    }
}

/// Load the mesh data from a COLLADA .dae file.
///
/// The data in a COLLADA files is formatted for efficiency, but isn't necessarily
/// organized in a way that is supported by the graphics API. This method reformats the
/// data so that it can be sent straight to the GPU without further manipulation.
///
/// In order to to this, it reorganizes the normals, UVs, and other vertex attributes to
/// be in the same order as the vertex positions.
fn geometry_to_mesh(geometry: &Geometry) -> Mesh {
    let mesh = match geometry.data {
        GeometricElement::Mesh(ref mesh) => mesh,
        _ => panic!("No mesh found within geometry")
    };

    let position_data_raw = get_raw_positions(&mesh);
    let normal_data_raw = get_normals(&mesh);

    let triangles = match mesh.primitives[0] {
        PrimitiveType::Triangles(ref triangles) => triangles,
        _ => panic!("Only triangles primitives are supported currently")
    };
    let primitive_indices = &triangles.primitives;

    // Create a new array for the positions so we can add the w coordinate.
    let mut position_data: Vec<f32> = Vec::with_capacity(position_data_raw.len() / 3 * 4);

    // Create a new array for the normals and rearrange them to match the order of position attributes.
    let mut normal_data: Vec<f32> = Vec::with_capacity(position_data.len());

    // Iterate over the indices, rearranging the normal data to match the position data.
    let stride = triangles.inputs.len();
    let mut vertex_index_map: HashMap<(usize, usize), u32> = HashMap::new();
    let mut indices: Vec<u32> = Vec::new();
    let vertex_count = triangles.count * 3;
    let mut index_count = 0;
    for offset in 0..vertex_count {
        // Determine the offset of the the current vertex's attributes
        let position_index = primitive_indices[offset * stride];
        let normal_index = primitive_indices[offset * stride + 1];

        // Push the index of the vertex, either reusing an existing vertex or creating a new one.
        let vertex_key = (position_index, normal_index);
        let vertex_index = if vertex_index_map.contains_key(&vertex_key) {
            // Vertex has already been assembled, reuse the index.
            (*vertex_index_map.get(&vertex_key).unwrap()) as u32
        } else {
            // Assemble new vertex.
            let vertex_index = index_count;
            index_count += 1;
            vertex_index_map.insert(vertex_key, vertex_index as u32);

            // Append position to position data array.
            for offset in 0..3 {
                position_data.push(position_data_raw[position_index * 3 + offset]);
            }
            position_data.push(1.0);

            // Append normal to normal data array.
            for offset in 0..3 {
                normal_data.push(normal_data_raw[normal_index * 3 + offset]);
            }

            vertex_index
        };

        indices.push(vertex_index);
    }

    let mesh = Mesh::from_raw_data(position_data.as_ref(), normal_data.as_ref(), indices.as_ref());

    mesh
}

fn get_raw_positions(mesh: &collada::Mesh) -> &[f32] {
    // TODO: Consult the correct element (<triangles> for now) to determine which source has position data.
    let position_data: &[f32] = match mesh.sources[0].array_element {
        ArrayElement::Float(ref float_array) => float_array.as_ref(),
        _ => panic!("Only float arrays supported for vertex position array")
    };
    assert!(position_data.len() > 0);

    position_data
}

fn get_normals(mesh: &collada::Mesh) -> &[f32] {
    // TODO: Consult the correct element (<triangles> for now) to determine which source has normal data.
    let normal_data: &[f32] = match mesh.sources[1].array_element {
        ArrayElement::Float(ref float_array) => float_array.as_ref(),
        _ => panic!("Only float arrays supported for vertex normal array")
    };
    assert!(normal_data.len() > 0);

    normal_data
}

pub fn load_file_text(path: &str) -> String {
    let file_path = Path::new(path);
    let mut file = match File::open(&file_path) {
        // The `desc` field of `IoError` is a string that describes the error
        Err(why) => panic!("couldn't open {}: {}", file_path.display(), Error::description(&why)),
        Ok(file) => file,
    };
    let mut contents = String::new();
    match file.read_to_string(&mut contents) {
        Err(why) => panic!("couldn't read {}: {}", file_path.display(), Error::description(&why)),
        Ok(_) => ()
    }
    contents
}
<|endoftext|>


<commit_before>use std::mem;
use std::fmt;
use std::io::{self, Read};
use std::time::Duration;

use futures::{Async, Poll, Stream};
use serde::de::DeserializeOwned;
use serde_json;

use client::KeepCoreThreadAlive;
use header::Headers;
use {async_impl, StatusCode, Url, wait};

/// A Response to a submitted `Request`.
pub struct Response {
    inner: async_impl::Response,
    body: async_impl::ReadableChunks<WaitBody>,
    _thread_handle: KeepCoreThreadAlive,
}

impl fmt::Debug for Response {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.inner, f)
    }
}

impl Response {
    /// Get the final `Url` of this `Response`.
    ///
    /// # Example
    ///
    /// ```rust
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let resp = reqwest::get("http://httpbin.org/redirect/1")?;
    /// assert_eq!(resp.url().as_str(), "http://httpbin.org/get");
    /// # Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn url(&self) -> &Url {
        self.inner.url()
    }

    /// Get the `StatusCode` of this `Response`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let resp = reqwest::get("http://httpbin.org/get")?;
    /// if resp.status().is_success() {
    ///     println!("success!");
    /// } else if resp.status().is_server_error() {
    ///     println!("server error!");
    /// } else {
    ///     println!("Something else happened. Status: {:?}", resp.status());
    /// }
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// ```rust
    /// use reqwest::Client;
    /// use reqwest::StatusCode;
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let client = Client::new();
    /// let resp = client.post("http://httpbin.org/post")
    ///             .body("possibly too large")
    ///             .send()?;
    /// match resp.status() {
    ///     StatusCode::Ok => println!("success!"),
    ///     StatusCode::PayloadTooLarge => {
    ///         println!("Request payload is too large!");
    ///     }
    ///     s => println!("Received response status: {:?}", s),
    /// };
    /// # Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn status(&self) -> StatusCode {
        self.inner.status()
    }

    /// Get the `Headers` of this `Response`.
    ///
    /// # Example
    ///
    /// Checking the `Content-Length` header before reading the response body.
    ///
    /// ```rust
    /// # use std::io::{Read, Write};
    /// # use reqwest::Client;
    /// # use reqwest::header::ContentLength;
    /// #
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let client = Client::new();
    /// let mut resp = client.head("http://httpbin.org/bytes/3000").send()?;
    /// if resp.status().is_success() {
    ///     let len = resp.headers().get::<ContentLength>()
    ///                 .map(|ct_len| **ct_len)
    ///                 .unwrap_or(0);
    ///     // limit 1mb response
    ///     if len <= 1_000_000 {
    ///         let mut buf = Vec::with_capacity(len as usize);
    ///         let mut resp = reqwest::get("http://httpbin.org/bytes/3000")?;
    ///         if resp.status().is_success() {
    ///             ::std::io::copy(&mut resp, &mut buf)?;
    ///         }
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn headers(&self) -> &Headers {
        self.inner.headers()
    }

    /// Try and deserialize the response body as JSON using `serde`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # extern crate reqwest;
    /// # #[macro_use] extern crate serde_derive;
    /// #
    /// # use reqwest::Error;
    /// #
    /// #[derive(Deserialize)]
    /// struct Ip {
    ///     origin: String,
    /// }
    ///
    /// # fn run() -> Result<(), Error> {
    /// let json: Ip = reqwest::get("http://httpbin.org/ip")?.json()?;
    /// # Ok(())
    /// # }
    /// #
    /// # fn main() { }
    /// ```
    ///
    /// # Errors
    ///
    /// This method fails whenever the response body is not in JSON format
    /// or it cannot be properly deserialized to target type `T`. For more
    /// details please see [`serde_json::from_reader`].
    /// [`serde_json::from_reader`]: https://docs.serde.rs/serde_json/fn.from_reader.html
    #[inline]
    pub fn json<T: DeserializeOwned>(&mut self) -> ::Result<T> {
        // There's 2 ways we could implement this:
        //
        // 1. Just using from_reader(self), making use of our blocking read adapter
        // 2. Just use self.inner.json().wait()
        //
        // Doing 1 is pretty easy, but it means we have the `serde_json` code
        // in more than one place, doing basically the same thing.
        //
        // Doing 2 would mean `serde_json` is only in one place, but we'd
        // need to update the sync Response to lazily make a blocking read
        // adapter, so that our `inner` could possibly still have the original
        // body.
        //
        // Went for easier for now, just to get it working.
        serde_json::from_reader(self).map_err(::error::from)
    }

    /// Copy the response body into a writer.
    ///
    /// This function internally uses [`std::io::copy`] and hence will continuously read data from
    /// the body and then write it into writer in a streaming fashion until EOF is met.
    ///
    /// On success, the total number of bytes that were copied to `writer` is returned.
    /// [`std::io::copy`]: https://doc.rust-lang.org/std/io/fn.copy.html
    ///
    /// # Example
    ///
    /// ```rust
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let mut resp = reqwest::get("http://httpbin.org/range/5")?;
    /// let mut buf: Vec<u8> = vec![];
    /// resp.copy_to(&mut buf)?;
    /// assert_eq!(b"abcde", buf.as_slice());
    /// # Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn copy_to<W: ?Sized>(&mut self, w: &mut W) -> io::Result<u64>
        where W: io::Write
    {
        io::copy(self, w)
    }

    /// Turn a response into an error if the server returned an error.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # extern crate reqwest;
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let res = reqwest::get("http://httpbin.org/status/400")?
    ///     .error_for_status();
    /// if let Err(err) = res {
    ///     assert_eq!(err.status(), Some(reqwest::StatusCode::BadRequest));
    /// }
    /// # Ok(())
    /// # }
    /// # fn main() {}
    /// ```
    #[inline]
    pub fn error_for_status(self) -> ::Result<Self> {
        let Response { body, inner, _thread_handle } = self;
        inner.error_for_status().map(move |inner| {
            Response {
                inner: inner,
                body: body,
                _thread_handle: _thread_handle,
            }
        })
    }
}

impl Read for Response {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.body.read(buf)
    }
}

struct WaitBody {
    inner: wait::WaitStream<async_impl::Decoder>
}

impl Stream for WaitBody {
    type Item = <async_impl::Decoder as Stream>::Item;
    type Error = <async_impl::Decoder as Stream>::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        match self.inner.next() {
            Some(Ok(chunk)) => Ok(Async::Ready(Some(chunk))),
            Some(Err(e)) => {
                let req_err = match e {
                    wait::Waited::TimedOut => ::error::timedout(None),
                    wait::Waited::Err(e) => e,
                };

                Err(req_err)
            },
            None => Ok(Async::Ready(None)),
        }
    }
}

// pub(crate)

pub fn new(mut res: async_impl::Response, timeout: Option<Duration>, thread: KeepCoreThreadAlive) -> Response {
    let body = mem::replace(res.body_mut(), async_impl::Decoder::empty());
    let body = async_impl::ReadableChunks::new(WaitBody {
        inner: wait::stream(body, timeout)
    });

    Response {
        inner: res,
        body: body,
        _thread_handle: thread,
    }
}
<commit_msg>feat: add Response::text()<commit_after>use std::mem;
use std::fmt;
use std::io::{self, Read};
use std::time::Duration;

use futures::{Async, Poll, Stream};
use serde::de::DeserializeOwned;
use serde_json;

use client::KeepCoreThreadAlive;
use header::Headers;
use {async_impl, StatusCode, Url, wait};

/// A Response to a submitted `Request`.
pub struct Response {
    inner: async_impl::Response,
    body: async_impl::ReadableChunks<WaitBody>,
    _thread_handle: KeepCoreThreadAlive,
}

impl fmt::Debug for Response {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.inner, f)
    }
}

impl Response {
    /// Get the final `Url` of this `Response`.
    ///
    /// # Example
    ///
    /// ```rust
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let resp = reqwest::get("http://httpbin.org/redirect/1")?;
    /// assert_eq!(resp.url().as_str(), "http://httpbin.org/get");
    /// # Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn url(&self) -> &Url {
        self.inner.url()
    }

    /// Get the `StatusCode` of this `Response`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let resp = reqwest::get("http://httpbin.org/get")?;
    /// if resp.status().is_success() {
    ///     println!("success!");
    /// } else if resp.status().is_server_error() {
    ///     println!("server error!");
    /// } else {
    ///     println!("Something else happened. Status: {:?}", resp.status());
    /// }
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// ```rust
    /// use reqwest::Client;
    /// use reqwest::StatusCode;
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let client = Client::new();
    /// let resp = client.post("http://httpbin.org/post")
    ///             .body("possibly too large")
    ///             .send()?;
    /// match resp.status() {
    ///     StatusCode::Ok => println!("success!"),
    ///     StatusCode::PayloadTooLarge => {
    ///         println!("Request payload is too large!");
    ///     }
    ///     s => println!("Received response status: {:?}", s),
    /// };
    /// # Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn status(&self) -> StatusCode {
        self.inner.status()
    }

    /// Get the `Headers` of this `Response`.
    ///
    /// # Example
    ///
    /// Checking the `Content-Length` header before reading the response body.
    ///
    /// ```rust
    /// # use std::io::{Read, Write};
    /// # use reqwest::Client;
    /// # use reqwest::header::ContentLength;
    /// #
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let client = Client::new();
    /// let mut resp = client.head("http://httpbin.org/bytes/3000").send()?;
    /// if resp.status().is_success() {
    ///     let len = resp.headers().get::<ContentLength>()
    ///                 .map(|ct_len| **ct_len)
    ///                 .unwrap_or(0);
    ///     // limit 1mb response
    ///     if len <= 1_000_000 {
    ///         let mut buf = Vec::with_capacity(len as usize);
    ///         let mut resp = reqwest::get("http://httpbin.org/bytes/3000")?;
    ///         if resp.status().is_success() {
    ///             ::std::io::copy(&mut resp, &mut buf)?;
    ///         }
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn headers(&self) -> &Headers {
        self.inner.headers()
    }

    /// Try and deserialize the response body as JSON using `serde`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # extern crate reqwest;
    /// # #[macro_use] extern crate serde_derive;
    /// #
    /// # use reqwest::Error;
    /// #
    /// #[derive(Deserialize)]
    /// struct Ip {
    ///     origin: String,
    /// }
    ///
    /// # fn run() -> Result<(), Error> {
    /// let json: Ip = reqwest::get("http://httpbin.org/ip")?.json()?;
    /// # Ok(())
    /// # }
    /// #
    /// # fn main() { }
    /// ```
    ///
    /// # Errors
    ///
    /// This method fails whenever the response body is not in JSON format
    /// or it cannot be properly deserialized to target type `T`. For more
    /// details please see [`serde_json::from_reader`].
    /// [`serde_json::from_reader`]: https://docs.serde.rs/serde_json/fn.from_reader.html
    #[inline]
    pub fn json<T: DeserializeOwned>(&mut self) -> ::Result<T> {
        // There's 2 ways we could implement this:
        //
        // 1. Just using from_reader(self), making use of our blocking read adapter
        // 2. Just use self.inner.json().wait()
        //
        // Doing 1 is pretty easy, but it means we have the `serde_json` code
        // in more than one place, doing basically the same thing.
        //
        // Doing 2 would mean `serde_json` is only in one place, but we'd
        // need to update the sync Response to lazily make a blocking read
        // adapter, so that our `inner` could possibly still have the original
        // body.
        //
        // Went for easier for now, just to get it working.
        serde_json::from_reader(self).map_err(::error::from)
    }

    /// Get the response text.
    ///
    /// # Example
    ///
    /// ```rust
    /// # extern crate reqwest;
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let content = reqwest::get("http://httpbin.org/range/26")?.text()?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn text(&mut self) -> io::Result<String> {
        let len = self.headers().get::<::header::ContentLength>()
            .map(|ct_len| **ct_len)
            .unwrap_or(0);
        let mut content = String::with_capacity(len as usize);
        self.read_to_string(&mut content)?;
        Ok(content)
    }

    /// Copy the response body into a writer.
    ///
    /// This function internally uses [`std::io::copy`] and hence will continuously read data from
    /// the body and then write it into writer in a streaming fashion until EOF is met.
    ///
    /// On success, the total number of bytes that were copied to `writer` is returned.
    /// [`std::io::copy`]: https://doc.rust-lang.org/std/io/fn.copy.html
    ///
    /// # Example
    ///
    /// ```rust
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let mut resp = reqwest::get("http://httpbin.org/range/5")?;
    /// let mut buf: Vec<u8> = vec![];
    /// resp.copy_to(&mut buf)?;
    /// assert_eq!(b"abcde", buf.as_slice());
    /// # Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn copy_to<W: ?Sized>(&mut self, w: &mut W) -> io::Result<u64>
        where W: io::Write
    {
        io::copy(self, w)
    }

    /// Turn a response into an error if the server returned an error.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # extern crate reqwest;
    /// # fn run() -> Result<(), Box<::std::error::Error>> {
    /// let res = reqwest::get("http://httpbin.org/status/400")?
    ///     .error_for_status();
    /// if let Err(err) = res {
    ///     assert_eq!(err.status(), Some(reqwest::StatusCode::BadRequest));
    /// }
    /// # Ok(())
    /// # }
    /// # fn main() {}
    /// ```
    #[inline]
    pub fn error_for_status(self) -> ::Result<Self> {
        let Response { body, inner, _thread_handle } = self;
        inner.error_for_status().map(move |inner| {
            Response {
                inner: inner,
                body: body,
                _thread_handle: _thread_handle,
            }
        })
    }
}

impl Read for Response {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.body.read(buf)
    }
}

struct WaitBody {
    inner: wait::WaitStream<async_impl::Decoder>
}

impl Stream for WaitBody {
    type Item = <async_impl::Decoder as Stream>::Item;
    type Error = <async_impl::Decoder as Stream>::Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        match self.inner.next() {
            Some(Ok(chunk)) => Ok(Async::Ready(Some(chunk))),
            Some(Err(e)) => {
                let req_err = match e {
                    wait::Waited::TimedOut => ::error::timedout(None),
                    wait::Waited::Err(e) => e,
                };

                Err(req_err)
            },
            None => Ok(Async::Ready(None)),
        }
    }
}

// pub(crate)

pub fn new(mut res: async_impl::Response, timeout: Option<Duration>, thread: KeepCoreThreadAlive) -> Response {
    let body = mem::replace(res.body_mut(), async_impl::Decoder::empty());
    let body = async_impl::ReadableChunks::new(WaitBody {
        inner: wait::stream(body, timeout)
    });

    Response {
        inner: res,
        body: body,
        _thread_handle: thread,
    }
}
<|endoftext|>


<commit_before>use common::{Header, HTTPVersion, StatusCode};
use std::io::{IoResult, MemReader};
use std::io::fs::File;
use std::io::util;
use std::io::util::NullReader;
use chunks::ChunksEncoder;

/// Object representing an HTTP response.
#[experimental]
pub struct Response<R> {
    reader: R,
    status_code: StatusCode,
    headers: Vec<Header>,
    http_version: HTTPVersion,
}

impl<R: Reader> Response<R> {
    #[experimental]
    pub fn new(status_code: StatusCode, mut headers: Vec<Header>,
               data: R, data_length: Option<uint>) -> Response<R>
    {
        // add Content-Length if not in the headers
        if data_length.is_some() {
            if headers.iter().find(|h| h.field.equiv(&"Content-Length")).is_none() {
                headers.unshift(
                    Header{
                        field: from_str("Content-Length").unwrap(),
                        value: format!("{}", data_length.unwrap())
                    }
                );
            }
        }

        Response {
            reader: data,
            status_code: status_code,
            headers: headers,
            http_version: HTTPVersion(1, 1),
        }
    }

    /// Returns the same request, but with an additional header.
    #[experimental]
    pub fn with_header(mut self, header: Header) -> Response<R> {
        self.headers.push(header);
        self
    }

    /// Returns the same request, but with a different status code.
    #[experimental]
    pub fn with_status_code(mut self, code: StatusCode) -> Response<R> {
        self.status_code = code;
        self
    }

    /// Forces an HTTP version.
    #[experimental]
    pub fn with_http_version(mut self, version: HTTPVersion) -> Response<R> {
        self.http_version = version;
        self
    }

    /// Cleans-up the headers so that they can be returned.
    #[experimental]
    fn purify_headers(&mut self) {
        // removing some unwanted headers, like Connection
        self.headers.retain(|h| !h.field.equiv(&"Connection"));

        // add Server if not in the headers
        if self.headers.iter().find(|h| h.field.equiv(&"Server")).is_none() {
            self.headers.unshift(
                Header{field: from_str("Server").unwrap(), value: "tiny-http (Rust)".to_string()}
            );
        }

        // add Connection: close
        /*self.headers.push(
            Header{field: "Connection".to_string(), value: "close".to_string()}
        );*/
    }

    /// Prints the HTTP response to a writer.
    #[experimental]
    pub fn raw_print<W: Writer>(mut self, mut writer: W) -> IoResult<()> {
        self.purify_headers();

        // if we don't have a Content-Length, or if the Content-Length is too big, using chunks writer
        let chunks_threshold = 32768;
        let use_chunks =
            self.http_version >= HTTPVersion(1, 1) &&
            self.headers.iter().find(|h| h.field.equiv(&"Content-Length"))
                .and_then(|h| from_str::<uint>(h.value.as_slice()))
                .filtered(|val| *val < chunks_threshold)
                .is_none();

        // add transfer-encoding header
        if use_chunks {
            self.headers.push(
                Header{field: from_str("Transfer-Encoding").unwrap(), value: "chunked".to_string()}
            )
        }

        // writing status line
        try!(write!(writer, "HTTP/{} {} {}\r\n",
            self.http_version,
            self.status_code.as_uint(),
            self.status_code.get_default_reason_phrase()
        ));

        // writing headers
        for header in self.headers.iter() {
            try!(write!(writer, "{}: {}\r\n", header.field, header.value));
        }

        // separator between header and data
        try!(write!(writer, "\r\n"));

        // writing data
        if use_chunks {
            let mut writer = ChunksEncoder::new(writer);
            try!(util::copy(&mut self.reader, &mut writer));
        } else {
            try!(util::copy(&mut self.reader, &mut writer));
        }

        Ok(())
    }
}

impl Response<File> {
    #[experimental]
    pub fn from_file(file: &Path) -> IoResult<Response<File>> {
        let mut file = try!(File::open(file));
        let stats = try!(file.stat());

        Ok(Response::new(
            StatusCode(200),
            Vec::new(),
            file,
            Some(stats.size as uint)
        ))
    }
}

impl Response<MemReader> {
    #[experimental]
    pub fn from_data(data: Vec<u8>) -> Response<MemReader> {
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            Vec::new(),
            MemReader::new(data),
            Some(data_len)
        )
    }

    #[experimental]
    pub fn from_string(data: String) -> Response<MemReader> {
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            vec!(
                from_str("Content-Type: text/plain; charset=UTF-8").unwrap()
            ),
            MemReader::new(data.into_bytes()),
            Some(data_len)
        )
    }
}

impl Response<NullReader> {
    #[experimental]
    pub fn empty() -> Response<NullReader> {
        Response::new(
            StatusCode(204),
            Vec::new(),
            NullReader,
            Some(0)
        )
    }

    #[experimental]
    pub fn not_modified() -> Response<NullReader> {
        Response::new(
            StatusCode(304),
            Vec::new(),
            NullReader,
            Some(0)
        )
    }
}
<commit_msg>Add doc for Response<commit_after>use common::{Header, HTTPVersion, StatusCode};
use std::io::{IoResult, MemReader};
use std::io::fs::File;
use std::io::util;
use std::io::util::NullReader;
use chunks::ChunksEncoder;

/// Object representing an HTTP response.
///
/// Some headers cannot be changed. Trying to define the value
///  of one of these will have no effect:
///
///  - `Accept-Ranges`
///  - `Connection`
///  - `Content-Range`
///  - `Trailer`
///  - `Transfer-Encoding`
///  - `Upgrade`
///
/// Some headers have special behaviors:
///
///  - `Content-Encoding`: If you define this header, the library
///      will assume that the data from the `Reader` has the specified encoding
///      and will just pass-through.
///  - `Content-Length`: If you define this header to `N`, only the first `N` bytes
///      of the `Reader` will be read. If the `Reader` reaches `EOF` before `N` bytes have
///      been read, `0`s will be sent.
#[experimental]
pub struct Response<R> {
    reader: R,
    status_code: StatusCode,
    headers: Vec<Header>,
    http_version: HTTPVersion,
}

impl<R: Reader> Response<R> {
    #[experimental]
    pub fn new(status_code: StatusCode, mut headers: Vec<Header>,
               data: R, data_length: Option<uint>) -> Response<R>
    {
        // add Content-Length if not in the headers
        if data_length.is_some() {
            if headers.iter().find(|h| h.field.equiv(&"Content-Length")).is_none() {
                headers.unshift(
                    Header{
                        field: from_str("Content-Length").unwrap(),
                        value: format!("{}", data_length.unwrap())
                    }
                );
            }
        }

        Response {
            reader: data,
            status_code: status_code,
            headers: headers,
            http_version: HTTPVersion(1, 1),
        }
    }

    /// Returns the same request, but with an additional header.
    #[experimental]
    pub fn with_header(mut self, header: Header) -> Response<R> {
        self.headers.push(header);
        self
    }

    /// Returns the same request, but with a different status code.
    #[experimental]
    pub fn with_status_code(mut self, code: StatusCode) -> Response<R> {
        self.status_code = code;
        self
    }

    /// Forces an HTTP version.
    #[experimental]
    pub fn with_http_version(mut self, version: HTTPVersion) -> Response<R> {
        self.http_version = version;
        self
    }

    /// Cleans-up the headers so that they can be returned.
    #[experimental]
    fn purify_headers(&mut self) {
        // removing some unwanted headers, like Connection
        self.headers.retain(|h| !h.field.equiv(&"Connection"));

        // add Server if not in the headers
        if self.headers.iter().find(|h| h.field.equiv(&"Server")).is_none() {
            self.headers.unshift(
                Header{field: from_str("Server").unwrap(), value: "tiny-http (Rust)".to_string()}
            );
        }

        // add Connection: close
        /*self.headers.push(
            Header{field: "Connection".to_string(), value: "close".to_string()}
        );*/
    }

    /// Prints the HTTP response to a writer.
    #[experimental]
    pub fn raw_print<W: Writer>(mut self, mut writer: W) -> IoResult<()> {
        self.purify_headers();

        // if we don't have a Content-Length, or if the Content-Length is too big, using chunks writer
        let chunks_threshold = 32768;
        let use_chunks =
            self.http_version >= HTTPVersion(1, 1) &&
            self.headers.iter().find(|h| h.field.equiv(&"Content-Length"))
                .and_then(|h| from_str::<uint>(h.value.as_slice()))
                .filtered(|val| *val < chunks_threshold)
                .is_none();

        // add transfer-encoding header
        if use_chunks {
            self.headers.push(
                Header{field: from_str("Transfer-Encoding").unwrap(), value: "chunked".to_string()}
            )
        }

        // writing status line
        try!(write!(writer, "HTTP/{} {} {}\r\n",
            self.http_version,
            self.status_code.as_uint(),
            self.status_code.get_default_reason_phrase()
        ));

        // writing headers
        for header in self.headers.iter() {
            try!(write!(writer, "{}: {}\r\n", header.field, header.value));
        }

        // separator between header and data
        try!(write!(writer, "\r\n"));

        // writing data
        if use_chunks {
            let mut writer = ChunksEncoder::new(writer);
            try!(util::copy(&mut self.reader, &mut writer));
        } else {
            try!(util::copy(&mut self.reader, &mut writer));
        }

        Ok(())
    }
}

impl Response<File> {
    #[experimental]
    pub fn from_file(file: &Path) -> IoResult<Response<File>> {
        let mut file = try!(File::open(file));
        let stats = try!(file.stat());

        Ok(Response::new(
            StatusCode(200),
            Vec::new(),
            file,
            Some(stats.size as uint)
        ))
    }
}

impl Response<MemReader> {
    #[experimental]
    pub fn from_data(data: Vec<u8>) -> Response<MemReader> {
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            Vec::new(),
            MemReader::new(data),
            Some(data_len)
        )
    }

    #[experimental]
    pub fn from_string(data: String) -> Response<MemReader> {
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            vec!(
                from_str("Content-Type: text/plain; charset=UTF-8").unwrap()
            ),
            MemReader::new(data.into_bytes()),
            Some(data_len)
        )
    }
}

impl Response<NullReader> {
    #[experimental]
    pub fn empty() -> Response<NullReader> {
        Response::new(
            StatusCode(204),
            Vec::new(),
            NullReader,
            Some(0)
        )
    }

    #[experimental]
    pub fn not_modified() -> Response<NullReader> {
        Response::new(
            StatusCode(304),
            Vec::new(),
            NullReader,
            Some(0)
        )
    }
}
<|endoftext|>


<commit_before>extern crate mustache;

use std::str;
use std::sync::{Arc, RWLock};
use std::collections::HashMap;
use std::io::{IoResult, File};
use std::io::util::copy;
use std::path::BytesContainer;
use http;
use time;
use mimes::get_media_type;

///A container for the response
pub struct Response<'a, 'b: 'a> {
    ///the original `http::server::ResponseWriter`
    pub origin: &'a mut http::server::ResponseWriter<'b>,
    templates: Arc<RWLock<HashMap<&'static str, mustache::Template>>>
}

impl<'a, 'b> Response<'a, 'b> {

    pub fn from_internal<'c, 'd>(response: &'c mut http::server::ResponseWriter<'d>,
                                 templates: Arc<RWLock<HashMap<&'static str, mustache::Template>>>
                                ) -> Response<'c, 'd>
    {
        Response {
            origin: response,
            templates: templates
        }
    }

    /// Writes a response
    ///
    /// # Example
    /// ```{rust,ignore}
    /// response.send("hello world");
    /// ```
    pub fn send<T: BytesContainer> (&mut self, text: T) {
        // TODO: This needs to be more sophisticated to return the correct headers
        // not just "some headers" :)
        Response::set_headers(self.origin);
        let _ = self.origin.write(text.container_as_bytes());
    }

    /// sets the content type by it's short form.
    ///
    /// # Example
    /// ```{rust,ignore}
    /// response.set_content_type("html");
    /// ```
    pub fn set_content_type(&mut self, text: &str) {
        // TODO: make this a chaining API. (Fight the lifetime hell!)
        self.origin.headers.content_type = get_media_type(text);
    }

    fn set_headers(response_writer: &mut http::server::ResponseWriter) {
        response_writer.headers.date = Some(time::now_utc());

        // we don't need to set this https://github.com/Ogeon/rustful/issues/3#issuecomment-44787613
        response_writer.headers.content_length = None;
        response_writer.headers.content_type = response_writer.headers.content_type
                                                                      .clone()
                                                                      .or(get_media_type("txt"));

        response_writer.headers.server = Some(String::from_str("Nickel"));
    }

    pub fn send_file(&mut self, path: &Path) -> IoResult<()> {
        let mut file = try!(File::open(path));
        self.origin.headers.content_length = None;

        self.origin.headers.content_type = path.extension_str().and_then(get_media_type);
        self.origin.headers.server = Some(String::from_str("Nickel"));
        copy(&mut file, self.origin)
    }

    /// Renders the given template bound with the given data.
    ///
    /// # Example
    /// ```{rust,ignore}
    /// let mut data = HashMap::<&'static str, &'static str>::new();
    /// data.insert("name", "user");
    /// response.render("examples/assets/template.tpl", &data);
    /// ```
    pub fn render(&mut self, path: &'static str, data: &HashMap<&'static str, &'static str>)
    {
        let mut templates = self.templates.write();
        let template = templates.find_or_insert_with(path, |_|
                     {
                         mustache::compile_str(str::from_utf8(
                            File::open(&Path::new(path))
                                 .read_to_end().ok().expect(format!("Couldn't open the template file: {}", path).as_slice()).as_slice()
                        ).expect("Coulnt't read template file as utf8")
                    )
                }
            );
        let _ = template.render(self.origin, data);
    }
}

#[test]
fn matches_content_type () {
    let path = &Path::new("test.txt");
    let content_type = path.extension_str().and_then(get_media_type).unwrap();

    assert_eq!(content_type.type_.as_slice(), "text");
    assert_eq!(content_type.subtype.as_slice(), "plain");
}
<commit_msg>Let response.render() take any encodable type.<commit_after>extern crate mustache;

use std::str;
use std::sync::{Arc, RWLock};
use std::collections::HashMap;
use std::io::{IoResult, File};
use std::io::util::copy;
use std::path::BytesContainer;
use serialize::Encodable;
use http;
use time;
use mimes::get_media_type;

///A container for the response
pub struct Response<'a, 'b: 'a> {
    ///the original `http::server::ResponseWriter`
    pub origin: &'a mut http::server::ResponseWriter<'b>,
    templates: Arc<RWLock<HashMap<&'static str, mustache::Template>>>
}

impl<'a, 'b> Response<'a, 'b> {

    pub fn from_internal<'c, 'd>(response: &'c mut http::server::ResponseWriter<'d>,
                                 templates: Arc<RWLock<HashMap<&'static str, mustache::Template>>>
                                ) -> Response<'c, 'd>
    {
        Response {
            origin: response,
            templates: templates
        }
    }

    /// Writes a response
    ///
    /// # Example
    /// ```{rust,ignore}
    /// response.send("hello world");
    /// ```
    pub fn send<T: BytesContainer> (&mut self, text: T) {
        // TODO: This needs to be more sophisticated to return the correct headers
        // not just "some headers" :)
        Response::set_headers(self.origin);
        let _ = self.origin.write(text.container_as_bytes());
    }

    /// sets the content type by it's short form.
    ///
    /// # Example
    /// ```{rust,ignore}
    /// response.set_content_type("html");
    /// ```
    pub fn set_content_type(&mut self, text: &str) {
        // TODO: make this a chaining API. (Fight the lifetime hell!)
        self.origin.headers.content_type = get_media_type(text);
    }

    fn set_headers(response_writer: &mut http::server::ResponseWriter) {
        response_writer.headers.date = Some(time::now_utc());

        // we don't need to set this https://github.com/Ogeon/rustful/issues/3#issuecomment-44787613
        response_writer.headers.content_length = None;
        response_writer.headers.content_type = response_writer.headers.content_type
                                                                      .clone()
                                                                      .or(get_media_type("txt"));

        response_writer.headers.server = Some(String::from_str("Nickel"));
    }

    pub fn send_file(&mut self, path: &Path) -> IoResult<()> {
        let mut file = try!(File::open(path));
        self.origin.headers.content_length = None;

        self.origin.headers.content_type = path.extension_str().and_then(get_media_type);
        self.origin.headers.server = Some(String::from_str("Nickel"));
        copy(&mut file, self.origin)
    }

    /// Renders the given template bound with the given data.
    ///
    /// # Example
    /// ```{rust,ignore}
    /// let mut data = HashMap::<&'static str, &'static str>::new();
    /// data.insert("name", "user");
    /// response.render("examples/assets/template.tpl", &data);
    /// ```
    pub fn render<'a, T: Encodable<mustache::Encoder<'a>, mustache::Error>>
        (&mut self, path: &'static str, data: &T)
    {
        let mut templates = self.templates.write();
        let template = templates.find_or_insert_with(path, |_|
                     {
                         mustache::compile_str(str::from_utf8(
                            File::open(&Path::new(path))
                                 .read_to_end().ok().expect(format!("Couldn't open the template file: {}", path).as_slice()).as_slice()
                        ).expect("Coulnt't read template file as utf8")
                    )
                }
            );
        let _ = template.render(self.origin, data);
    }
}

#[test]
fn matches_content_type () {
    let path = &Path::new("test.txt");
    let content_type = path.extension_str().and_then(get_media_type).unwrap();

    assert_eq!(content_type.type_.as_slice(), "text");
    assert_eq!(content_type.subtype.as_slice(), "plain");
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::AptosPackageHooks;
use aptos::move_tool::MemberId;
use aptos_crypto::ed25519::Ed25519PrivateKey;
use aptos_crypto::{PrivateKey, Uniform};
use aptos_types::{
    access_path::AccessPath,
    account_address::AccountAddress,
    state_store::state_key::StateKey,
    transaction::{EntryFunction, SignedTransaction, TransactionPayload, TransactionStatus},
};
use cached_packages::aptos_stdlib;
use framework::{BuildOptions, BuiltPackage};
use language_e2e_tests::{
    account::{Account, AccountData},
    executor::FakeExecutor,
};
use move_deps::move_core_types::language_storage::{ResourceKey, StructTag, TypeTag};
use move_deps::move_package::package_hooks::register_package_hooks;
use project_root::get_project_root;
use rand::{
    rngs::{OsRng, StdRng},
    Rng, SeedableRng,
};
use serde::de::DeserializeOwned;
use std::collections::BTreeMap;
use std::path::Path;

/// A simple test harness for defining Move e2e tests.
///
/// Tests defined via this harness typically live in the `<crate>/tests` directory, the standard
/// Rust place for defining integration tests.
///
/// For defining a set of new tests around a specific area, you add a new Rust source
/// `tested_area.rs` to the `tests` directory of your crate. You also will create a directory
/// `tested_area.data` which lives side-by-side with the Rust source. In this directory, you
/// place any number of Move packages you need for running the tests. In addition, the test
/// infrastructure will place baseline (golden) files in the `tested_area.data` using the `.exp`
/// (expected) ending.  For examples, see e.g. the `tests/code_publishing.rs` test in this crate.
///
/// NOTE: This harness currently is a wrapper around existing legacy e2e testing infra. We
/// eventually plan to retire the legacy code, and are rather keen to know what of the legacy
/// test infra we want to maintain and also which existing tests to preserve.
pub struct MoveHarness {
    /// The executor being used.
    pub executor: FakeExecutor,
    /// The current transaction sequence number, by account address.
    txn_seq_no: BTreeMap<AccountAddress, u64>,
}

impl MoveHarness {
    /// Creates a new harness.
    pub fn new() -> Self {
        register_package_hooks(Box::new(AptosPackageHooks {}));
        Self {
            executor: FakeExecutor::from_fresh_genesis(),
            txn_seq_no: BTreeMap::default(),
        }
    }

    pub fn new_mainnet() -> Self {
        Self {
            executor: FakeExecutor::from_mainnet_genesis(),
            txn_seq_no: BTreeMap::default(),
        }
    }

    pub fn new_no_parallel() -> Self {
        Self {
            executor: FakeExecutor::from_fresh_genesis().set_not_parallel(),
            txn_seq_no: BTreeMap::default(),
        }
    }

    /// Creates an account for the given static address. This address needs to be static so
    /// we can load regular Move code to there without need to rewrite code addresses.
    pub fn new_account_at(&mut self, addr: AccountAddress) -> Account {
        // The below will use the genesis keypair but that should be fine.
        let acc = Account::new_genesis_account(addr);
        // Mint the account 10M Aptos coins (with 8 decimals).
        let data = AccountData::with_account(acc, 1_000_000_000_000_000, 10);
        self.executor.add_account_data(&data);
        self.txn_seq_no.insert(addr, 10);
        data.account().clone()
    }

    // Creates an account with a randomly generated address and key pair
    pub fn new_account_with_key_pair(&mut self) -> Account {
        let mut rng = StdRng::from_seed(OsRng.gen());

        let privkey = Ed25519PrivateKey::generate(&mut rng);
        let pubkey = privkey.public_key();
        let acc = Account::with_keypair(privkey, pubkey);
        let data = AccountData::with_account(acc.clone(), 1_000_000_000_000_000, 10);
        self.executor.add_account_data(&data);
        self.txn_seq_no.insert(*acc.address(), 10);
        data.account().clone()
    }

    /// Gets the account where the Aptos framework is installed (0x1).
    pub fn aptos_framework_account(&mut self) -> Account {
        self.new_account_at(AccountAddress::ONE)
    }

    /// Runs a signed transaction. On success, applies the write set.
    pub fn run(&mut self, txn: SignedTransaction) -> TransactionStatus {
        let output = self.executor.execute_transaction(txn);
        if matches!(output.status(), TransactionStatus::Keep(_)) {
            self.executor.apply_write_set(output.write_set());
        }
        output.status().to_owned()
    }

    /// Runs a block of signed transactions. On success, applies the write set.
    pub fn run_block(&mut self, txn_block: Vec<SignedTransaction>) -> Vec<TransactionStatus> {
        let mut result = vec![];
        for output in self.executor.execute_block(txn_block).unwrap() {
            if matches!(output.status(), TransactionStatus::Keep(_)) {
                self.executor.apply_write_set(output.write_set());
            }
            result.push(output.status().to_owned())
        }
        result
    }

    /// Creates a transaction, based on provided payload.
    pub fn create_transaction_payload(
        &mut self,
        account: &Account,
        payload: TransactionPayload,
    ) -> SignedTransaction {
        let seq_no_ref = self.txn_seq_no.get_mut(account.address()).unwrap();
        let seq_no = *seq_no_ref;
        *seq_no_ref += 1;
        account
            .transaction()
            .sequence_number(seq_no)
            .max_gas_amount(1_000_000)
            .gas_unit_price(1)
            .payload(payload)
            .sign()
    }

    /// Runs a transaction, based on provided payload. If the transaction succeeds, any generated
    /// writeset will be applied to storage.
    pub fn run_transaction_payload(
        &mut self,
        account: &Account,
        payload: TransactionPayload,
    ) -> TransactionStatus {
        let txn = self.create_transaction_payload(account, payload);
        self.run(txn)
    }

    /// Creates a transaction which runs the specified entry point `fun`. Arguments need to be
    /// provided in bcs-serialized form.
    pub fn create_entry_function(
        &mut self,
        account: &Account,
        fun: MemberId,
        ty_args: Vec<TypeTag>,
        args: Vec<Vec<u8>>,
    ) -> SignedTransaction {
        let MemberId {
            module_id,
            member_id: function_id,
        } = fun;
        self.create_transaction_payload(
            account,
            TransactionPayload::EntryFunction(EntryFunction::new(
                module_id,
                function_id,
                ty_args,
                args,
            )),
        )
    }

    /// Run the specified entry point `fun`. Arguments need to be provided in bcs-serialized form.
    pub fn run_entry_function(
        &mut self,
        account: &Account,
        fun: MemberId,
        ty_args: Vec<TypeTag>,
        args: Vec<Vec<u8>>,
    ) -> TransactionStatus {
        let txn = self.create_entry_function(account, fun, ty_args, args);
        self.run(txn)
    }

    /// Creates a transaction which publishes the Move Package found at the given path on behalf
    /// of the given account.
    pub fn create_publish_package(
        &mut self,
        account: &Account,
        path: &Path,
        options: Option<BuildOptions>,
    ) -> SignedTransaction {
        let package = BuiltPackage::build(path.to_owned(), options.unwrap_or_default())
            .expect("building package must succeed");
        let code = package.extract_code();
        let metadata = package
            .extract_metadata()
            .expect("extracting package metadata must succeed");
        self.create_transaction_payload(
            account,
            aptos_stdlib::code_publish_package_txn(
                bcs::to_bytes(&metadata).expect("PackageMetadata has BCS"),
                code,
            ),
        )
    }

    /// Runs transaction which publishes the Move Package.
    pub fn publish_package(&mut self, account: &Account, path: &Path) -> TransactionStatus {
        let txn = self.create_publish_package(account, path, None);
        self.run(txn)
    }

    /// Runs transaction which publishes the Move Package.
    pub fn publish_package_with_options(
        &mut self,
        account: &Account,
        path: &Path,
        options: BuildOptions,
    ) -> TransactionStatus {
        let txn = self.create_publish_package(account, path, Some(options));
        self.run(txn)
    }

    pub fn fast_forward(&mut self, seconds: u64) {
        let current_time = self.executor.get_block_time();
        self.executor
            .set_block_time(current_time + seconds * 1_000_000)
    }

    pub fn new_epoch(&mut self) {
        self.fast_forward(7200);
        self.executor.new_block()
    }

    pub fn new_block_with_metadata(
        &mut self,
        proposer: AccountAddress,
        failed_proposer_indices: Vec<u32>,
    ) {
        self.fast_forward(1);
        self.executor
            .new_block_with_metadata(proposer, failed_proposer_indices);
    }

    pub fn read_state_value(&self, state_key: &StateKey) -> Option<Vec<u8>> {
        self.executor.read_state_value(state_key).and_then(|bytes| {
            if bytes.is_empty() {
                None
            } else {
                Some(bytes)
            }
        })
    }

    /// Reads the raw, serialized data of a resource.
    pub fn read_resource_raw(
        &self,
        addr: &AccountAddress,
        struct_tag: StructTag,
    ) -> Option<Vec<u8>> {
        let path = AccessPath::resource_access_path(ResourceKey::new(*addr, struct_tag));
        self.read_state_value(&StateKey::AccessPath(path))
    }

    /// Reads the resource data `T`.
    pub fn read_resource<T: DeserializeOwned>(
        &self,
        addr: &AccountAddress,
        struct_tag: StructTag,
    ) -> Option<T> {
        Some(
            bcs::from_bytes::<T>(&self.read_resource_raw(addr, struct_tag)?).expect(
                "serialization expected to succeed (Rust type incompatible with Move type?)",
            ),
        )
    }

    /// Checks whether resource exists.
    pub fn exists_resource(&self, addr: &AccountAddress, struct_tag: StructTag) -> bool {
        self.read_resource_raw(addr, struct_tag).is_some()
    }
}

/// Enables golden files for the given harness. The golden file will be stored side-by-side
/// with the data directory of a Rust source, named after the test function.
#[macro_export]
macro_rules! enable_golden {
    ($h:expr) => {
        $h.internal_set_golden(std::file!(), language_e2e_tests::current_function_name!())
    };
}

impl MoveHarness {
    /// Internal function to support the `enable_golden` macro.
    pub fn internal_set_golden(&mut self, file_macro_value: &str, function_macro_value: &str) {
        // The result of `std::file!` gives us a name relative to the project root,
        // so we need to add that to it. We also want to replace the extension `.rs` with `.data`.
        let mut path = get_project_root().unwrap().join(file_macro_value);
        path.set_extension("data");
        // The result of the `current_function` macro gives us the fully qualified
        // We only want the trailing simple name.
        let fun = function_macro_value.split("::").last().unwrap();
        self.executor
            .set_golden_file_at(&path.display().to_string(), fun)
    }
}

/// Helper to assert transaction is successful
#[macro_export]
macro_rules! assert_success {
    ($s:expr) => {{
        use aptos_types::transaction::*;
        assert_eq!($s, TransactionStatus::Keep(ExecutionStatus::Success))
    }};
}

/// Helper to assert transaction aborts.
#[macro_export]
macro_rules! assert_abort {
    ($s:expr, $c:pat) => {{
        assert!(matches!(
            $s,
            aptos_types::transaction::TransactionStatus::Keep(
                aptos_types::transaction::ExecutionStatus::MoveAbort { code: $c, .. }
            ),
        ));
    }};
}

/// Helper to assert vm status code.
#[macro_export]
macro_rules! assert_vm_status {
    ($s:expr, $c:pat) => {{
        use aptos_types::transaction::*;
        assert!(matches!(
            $s,
            TransactionStatus::Keep(ExecutionStatus::MiscellaneousError(Some($c)))
        ));
    }};
}
<commit_msg>[move-tests] remove the no parallel as unused<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::AptosPackageHooks;
use aptos::move_tool::MemberId;
use aptos_crypto::ed25519::Ed25519PrivateKey;
use aptos_crypto::{PrivateKey, Uniform};
use aptos_types::{
    access_path::AccessPath,
    account_address::AccountAddress,
    state_store::state_key::StateKey,
    transaction::{EntryFunction, SignedTransaction, TransactionPayload, TransactionStatus},
};
use cached_packages::aptos_stdlib;
use framework::{BuildOptions, BuiltPackage};
use language_e2e_tests::{
    account::{Account, AccountData},
    executor::FakeExecutor,
};
use move_deps::move_core_types::language_storage::{ResourceKey, StructTag, TypeTag};
use move_deps::move_package::package_hooks::register_package_hooks;
use project_root::get_project_root;
use rand::{
    rngs::{OsRng, StdRng},
    Rng, SeedableRng,
};
use serde::de::DeserializeOwned;
use std::collections::BTreeMap;
use std::path::Path;

/// A simple test harness for defining Move e2e tests.
///
/// Tests defined via this harness typically live in the `<crate>/tests` directory, the standard
/// Rust place for defining integration tests.
///
/// For defining a set of new tests around a specific area, you add a new Rust source
/// `tested_area.rs` to the `tests` directory of your crate. You also will create a directory
/// `tested_area.data` which lives side-by-side with the Rust source. In this directory, you
/// place any number of Move packages you need for running the tests. In addition, the test
/// infrastructure will place baseline (golden) files in the `tested_area.data` using the `.exp`
/// (expected) ending.  For examples, see e.g. the `tests/code_publishing.rs` test in this crate.
///
/// NOTE: This harness currently is a wrapper around existing legacy e2e testing infra. We
/// eventually plan to retire the legacy code, and are rather keen to know what of the legacy
/// test infra we want to maintain and also which existing tests to preserve.
pub struct MoveHarness {
    /// The executor being used.
    pub executor: FakeExecutor,
    /// The current transaction sequence number, by account address.
    txn_seq_no: BTreeMap<AccountAddress, u64>,
}

impl MoveHarness {
    /// Creates a new harness.
    pub fn new() -> Self {
        register_package_hooks(Box::new(AptosPackageHooks {}));
        Self {
            executor: FakeExecutor::from_fresh_genesis(),
            txn_seq_no: BTreeMap::default(),
        }
    }

    pub fn new_mainnet() -> Self {
        Self {
            executor: FakeExecutor::from_mainnet_genesis(),
            txn_seq_no: BTreeMap::default(),
        }
    }

    /// Creates an account for the given static address. This address needs to be static so
    /// we can load regular Move code to there without need to rewrite code addresses.
    pub fn new_account_at(&mut self, addr: AccountAddress) -> Account {
        // The below will use the genesis keypair but that should be fine.
        let acc = Account::new_genesis_account(addr);
        // Mint the account 10M Aptos coins (with 8 decimals).
        let data = AccountData::with_account(acc, 1_000_000_000_000_000, 10);
        self.executor.add_account_data(&data);
        self.txn_seq_no.insert(addr, 10);
        data.account().clone()
    }

    // Creates an account with a randomly generated address and key pair
    pub fn new_account_with_key_pair(&mut self) -> Account {
        let mut rng = StdRng::from_seed(OsRng.gen());

        let privkey = Ed25519PrivateKey::generate(&mut rng);
        let pubkey = privkey.public_key();
        let acc = Account::with_keypair(privkey, pubkey);
        let data = AccountData::with_account(acc.clone(), 1_000_000_000_000_000, 10);
        self.executor.add_account_data(&data);
        self.txn_seq_no.insert(*acc.address(), 10);
        data.account().clone()
    }

    /// Gets the account where the Aptos framework is installed (0x1).
    pub fn aptos_framework_account(&mut self) -> Account {
        self.new_account_at(AccountAddress::ONE)
    }

    /// Runs a signed transaction. On success, applies the write set.
    pub fn run(&mut self, txn: SignedTransaction) -> TransactionStatus {
        let output = self.executor.execute_transaction(txn);
        if matches!(output.status(), TransactionStatus::Keep(_)) {
            self.executor.apply_write_set(output.write_set());
        }
        output.status().to_owned()
    }

    /// Runs a block of signed transactions. On success, applies the write set.
    pub fn run_block(&mut self, txn_block: Vec<SignedTransaction>) -> Vec<TransactionStatus> {
        let mut result = vec![];
        for output in self.executor.execute_block(txn_block).unwrap() {
            if matches!(output.status(), TransactionStatus::Keep(_)) {
                self.executor.apply_write_set(output.write_set());
            }
            result.push(output.status().to_owned())
        }
        result
    }

    /// Creates a transaction, based on provided payload.
    pub fn create_transaction_payload(
        &mut self,
        account: &Account,
        payload: TransactionPayload,
    ) -> SignedTransaction {
        let seq_no_ref = self.txn_seq_no.get_mut(account.address()).unwrap();
        let seq_no = *seq_no_ref;
        *seq_no_ref += 1;
        account
            .transaction()
            .sequence_number(seq_no)
            .max_gas_amount(1_000_000)
            .gas_unit_price(1)
            .payload(payload)
            .sign()
    }

    /// Runs a transaction, based on provided payload. If the transaction succeeds, any generated
    /// writeset will be applied to storage.
    pub fn run_transaction_payload(
        &mut self,
        account: &Account,
        payload: TransactionPayload,
    ) -> TransactionStatus {
        let txn = self.create_transaction_payload(account, payload);
        self.run(txn)
    }

    /// Creates a transaction which runs the specified entry point `fun`. Arguments need to be
    /// provided in bcs-serialized form.
    pub fn create_entry_function(
        &mut self,
        account: &Account,
        fun: MemberId,
        ty_args: Vec<TypeTag>,
        args: Vec<Vec<u8>>,
    ) -> SignedTransaction {
        let MemberId {
            module_id,
            member_id: function_id,
        } = fun;
        self.create_transaction_payload(
            account,
            TransactionPayload::EntryFunction(EntryFunction::new(
                module_id,
                function_id,
                ty_args,
                args,
            )),
        )
    }

    /// Run the specified entry point `fun`. Arguments need to be provided in bcs-serialized form.
    pub fn run_entry_function(
        &mut self,
        account: &Account,
        fun: MemberId,
        ty_args: Vec<TypeTag>,
        args: Vec<Vec<u8>>,
    ) -> TransactionStatus {
        let txn = self.create_entry_function(account, fun, ty_args, args);
        self.run(txn)
    }

    /// Creates a transaction which publishes the Move Package found at the given path on behalf
    /// of the given account.
    pub fn create_publish_package(
        &mut self,
        account: &Account,
        path: &Path,
        options: Option<BuildOptions>,
    ) -> SignedTransaction {
        let package = BuiltPackage::build(path.to_owned(), options.unwrap_or_default())
            .expect("building package must succeed");
        let code = package.extract_code();
        let metadata = package
            .extract_metadata()
            .expect("extracting package metadata must succeed");
        self.create_transaction_payload(
            account,
            aptos_stdlib::code_publish_package_txn(
                bcs::to_bytes(&metadata).expect("PackageMetadata has BCS"),
                code,
            ),
        )
    }

    /// Runs transaction which publishes the Move Package.
    pub fn publish_package(&mut self, account: &Account, path: &Path) -> TransactionStatus {
        let txn = self.create_publish_package(account, path, None);
        self.run(txn)
    }

    /// Runs transaction which publishes the Move Package.
    pub fn publish_package_with_options(
        &mut self,
        account: &Account,
        path: &Path,
        options: BuildOptions,
    ) -> TransactionStatus {
        let txn = self.create_publish_package(account, path, Some(options));
        self.run(txn)
    }

    pub fn fast_forward(&mut self, seconds: u64) {
        let current_time = self.executor.get_block_time();
        self.executor
            .set_block_time(current_time + seconds * 1_000_000)
    }

    pub fn new_epoch(&mut self) {
        self.fast_forward(7200);
        self.executor.new_block()
    }

    pub fn new_block_with_metadata(
        &mut self,
        proposer: AccountAddress,
        failed_proposer_indices: Vec<u32>,
    ) {
        self.fast_forward(1);
        self.executor
            .new_block_with_metadata(proposer, failed_proposer_indices);
    }

    pub fn read_state_value(&self, state_key: &StateKey) -> Option<Vec<u8>> {
        self.executor.read_state_value(state_key).and_then(|bytes| {
            if bytes.is_empty() {
                None
            } else {
                Some(bytes)
            }
        })
    }

    /// Reads the raw, serialized data of a resource.
    pub fn read_resource_raw(
        &self,
        addr: &AccountAddress,
        struct_tag: StructTag,
    ) -> Option<Vec<u8>> {
        let path = AccessPath::resource_access_path(ResourceKey::new(*addr, struct_tag));
        self.read_state_value(&StateKey::AccessPath(path))
    }

    /// Reads the resource data `T`.
    pub fn read_resource<T: DeserializeOwned>(
        &self,
        addr: &AccountAddress,
        struct_tag: StructTag,
    ) -> Option<T> {
        Some(
            bcs::from_bytes::<T>(&self.read_resource_raw(addr, struct_tag)?).expect(
                "serialization expected to succeed (Rust type incompatible with Move type?)",
            ),
        )
    }

    /// Checks whether resource exists.
    pub fn exists_resource(&self, addr: &AccountAddress, struct_tag: StructTag) -> bool {
        self.read_resource_raw(addr, struct_tag).is_some()
    }
}

/// Enables golden files for the given harness. The golden file will be stored side-by-side
/// with the data directory of a Rust source, named after the test function.
#[macro_export]
macro_rules! enable_golden {
    ($h:expr) => {
        $h.internal_set_golden(std::file!(), language_e2e_tests::current_function_name!())
    };
}

impl MoveHarness {
    /// Internal function to support the `enable_golden` macro.
    pub fn internal_set_golden(&mut self, file_macro_value: &str, function_macro_value: &str) {
        // The result of `std::file!` gives us a name relative to the project root,
        // so we need to add that to it. We also want to replace the extension `.rs` with `.data`.
        let mut path = get_project_root().unwrap().join(file_macro_value);
        path.set_extension("data");
        // The result of the `current_function` macro gives us the fully qualified
        // We only want the trailing simple name.
        let fun = function_macro_value.split("::").last().unwrap();
        self.executor
            .set_golden_file_at(&path.display().to_string(), fun)
    }
}

/// Helper to assert transaction is successful
#[macro_export]
macro_rules! assert_success {
    ($s:expr) => {{
        use aptos_types::transaction::*;
        assert_eq!($s, TransactionStatus::Keep(ExecutionStatus::Success))
    }};
}

/// Helper to assert transaction aborts.
#[macro_export]
macro_rules! assert_abort {
    ($s:expr, $c:pat) => {{
        assert!(matches!(
            $s,
            aptos_types::transaction::TransactionStatus::Keep(
                aptos_types::transaction::ExecutionStatus::MoveAbort { code: $c, .. }
            ),
        ));
    }};
}

/// Helper to assert vm status code.
#[macro_export]
macro_rules! assert_vm_status {
    ($s:expr, $c:pat) => {{
        use aptos_types::transaction::*;
        assert!(matches!(
            $s,
            TransactionStatus::Keep(ExecutionStatus::MiscellaneousError(Some($c)))
        ));
    }};
}
<|endoftext|>


<commit_before>use crate::{aws::Aws, util::unix_timestamp_now};
use rusoto_core::Region;
use rusoto_logs::{
    CloudWatchLogs, CloudWatchLogsClient, DeleteLogStreamRequest, DescribeLogStreamsRequest,
};
use std::{thread, time::Duration};

pub struct LogPruner {
    aws: Aws,
    aws_logs: CloudWatchLogsClient,
}

impl LogPruner {
    pub fn new(aws: Aws) -> Self {
        let aws_logs = CloudWatchLogsClient::new(Region::UsWest2);
        Self { aws, aws_logs }
    }

    pub fn prune_logs(&self) {
        let log_group = self.aws.workplace();
        let mut next_token = None;
        let now = unix_timestamp_now();
        let mut errors = 0usize;
        let mut deleted = 0usize;
        let max_age = Duration::from_secs(3 * 24 * 3600); //3 days
        loop {
            let mut request = DescribeLogStreamsRequest::default();
            request.limit = Some(50);
            request.log_group_name = log_group.to_string();
            request.order_by = Some("LastEventTime".to_string());
            request.next_token = next_token.clone();
            let response = self.aws_logs.describe_log_streams(request).sync();
            let response = match response {
                Err(err) => {
                    errors += 1;
                    if errors > 5 {
                        panic!("Too many aws errors, aborting");
                    }
                    println!("Retrying aws error: {:?}", err);
                    thread::sleep(Duration::from_millis(600));
                    continue;
                }
                Ok(r) => r,
            };
            errors = 0;
            let streams = response.log_streams.expect("No log_streams");
            for stream in streams {
                let last_event = stream.last_event_timestamp.unwrap_or(0);
                let last_event = Duration::from_millis(last_event as u64);
                if now > last_event {
                    let age = now - last_event;
                    if age < max_age {
                        continue;
                    }
                    let log_stream_name = stream.log_stream_name.expect("No log_stream_name");
                    loop {
                        let request = DeleteLogStreamRequest {
                            log_group_name: log_group.clone(),
                            log_stream_name: log_stream_name.clone(),
                        };
                        let response = self.aws_logs.delete_log_stream(request).sync();
                        if let Err(e) = response {
                            println!("Retrying aws error during log stream removal: {:?}", e);
                            thread::sleep(Duration::from_millis(600));
                            continue;
                        }
                        break;
                    }
                    thread::sleep(Duration::from_millis(200));
                    deleted += 1;
                    if deleted % 100 == 0 {
                        println!("Deleted {} streams", deleted);
                    }
                }
            }
            next_token = response.next_token;
            if next_token.is_none() {
                break;
            }
            // AWS limit is 5 calls to describe_log_streams per second
            thread::sleep(Duration::from_millis(200));
        }
    }
}
<commit_msg>[cluster_test] Minor improvements to --prune-logs<commit_after>use crate::{aws::Aws, util::unix_timestamp_now};
use rusoto_core::Region;
use rusoto_logs::{
    CloudWatchLogs, CloudWatchLogsClient, DeleteLogStreamRequest, DescribeLogStreamsRequest,
};
use std::{thread, time::Duration};

pub struct LogPruner {
    aws: Aws,
    aws_logs: CloudWatchLogsClient,
}

impl LogPruner {
    pub fn new(aws: Aws) -> Self {
        let aws_logs = CloudWatchLogsClient::new(Region::UsWest2);
        Self { aws, aws_logs }
    }

    pub fn prune_logs(&self) {
        let log_group = self.aws.workplace();
        let mut next_token = None;
        let now = unix_timestamp_now();
        let mut errors = 0usize;
        let mut deleted = 0usize;
        let max_age = Duration::from_secs(3 * 24 * 3600); //3 days
        loop {
            let mut request = DescribeLogStreamsRequest::default();
            request.limit = Some(50);
            request.log_group_name = log_group.to_string();
            request.order_by = Some("LastEventTime".to_string());
            request.next_token = next_token.clone();
            let response = self.aws_logs.describe_log_streams(request).sync();
            let response = match response {
                Err(err) => {
                    errors += 1;
                    if errors > 5 {
                        panic!("Too many aws errors, aborting");
                    }
                    println!("Retrying aws error: {:?}", err);
                    thread::sleep(Duration::from_millis(600));
                    continue;
                }
                Ok(r) => r,
            };
            errors = 0;
            let streams = response.log_streams.expect("No log_streams");
            for stream in streams {
                let last_event = stream.last_event_timestamp.unwrap_or(0);
                let last_event = Duration::from_millis(last_event as u64);
                if now > last_event {
                    let age = now - last_event;
                    if age < max_age {
                        continue;
                    }
                    let log_stream_name = stream.log_stream_name.expect("No log_stream_name");
                    let mut errors_delete = 0usize;
                    loop {
                        let request = DeleteLogStreamRequest {
                            log_group_name: log_group.clone(),
                            log_stream_name: log_stream_name.clone(),
                        };
                        let response = self.aws_logs.delete_log_stream(request).sync();
                        if let Err(e) = response {
                            errors_delete += 1;
                            if errors_delete > 5 {
                                panic!("Too many aws errors, aborting");
                            }
                            println!("Retrying aws error during log stream removal: {:?}", e);
                            thread::sleep(Duration::from_millis(600));
                            continue;
                        }
                        break;
                    }
                    thread::sleep(Duration::from_millis(200));
                    deleted += 1;
                    if deleted % 100 == 0 {
                        let age_days = age.as_secs() / 3600 / 24;
                        println!("Deleted {} streams, last age: {} days", deleted, age_days);
                    }
                }
            }
            next_token = response.next_token;
            if next_token.is_none() {
                break;
            }
            // AWS limit is 5 calls to describe_log_streams per second
            thread::sleep(Duration::from_millis(200));
        }
    }
}
<|endoftext|>


<commit_before>#![feature(test)]
extern crate parsetree;
extern crate test;
extern crate nom;

use test::Bencher;
use parsetree::Tree;
use parsetree::examples::prelude::*;

#[bench]
fn tree_traversal_parse(bencher: &mut Bencher) {
    // Create a tree
    let mut tree = Tree::new();

    // Create some parsers
    let example_parser_1 = tree.new_parser(ExampleParser1);
    let example_parser_2 = tree.new_parser(ExampleParser2);

    // Combine the parsers
    tree.link(example_parser_1, example_parser_2);

    // Traverse the tree and find the "best" parsing result
    let input = [0xff; 100];

    bencher.iter(|| {
        assert_eq!(tree.traverse(example_parser_1, &input, vec![]).len(), 2);
    });
    bencher.bytes = input.len() as u64;
}

#[bench]
fn tree_traversal_only(bencher: &mut Bencher) {
    // Create a tree
    let mut tree = Tree::new();

    // Create some parsers
    let example_parser_1 = tree.new_parser(ExampleParser1);
    let example_parser_2 = tree.new_parser(ExampleParser2);

    // Combine the parsers
    tree.link(example_parser_1, example_parser_2);

    bencher.iter(|| { for _ in example_parser_1.descendants(&tree.arena) {} });
}

#[bench]
fn tree_generation(bencher: &mut Bencher) {
    bencher.iter(|| {
        // Create a tree
        let mut tree = Tree::new();

        // Create some parsers
        let example_parser_1 = tree.new_parser(ExampleParser1);
        let example_parser_2 = tree.new_parser(ExampleParser2);

        // Combine the parsers
        tree.link(example_parser_1, example_parser_2);
    });
}
<commit_msg>Removed benchmarks for now<commit_after><|endoftext|>