<commit_before>#![allow(dead_code, unused_variables)]
use alloc::heap::{Heap, Alloc, Layout};
use bytes::BytesMut;

use std::rc::Rc;

use audiosample::*;
use pixel::*;
use timeinfo::*;

pub mod error {
    error_chain! {
        errors {
            InvalidIndex
            InvalidConversion
        }
    }
}

use self::error::*;

// TODO: Document
// TODO: Change it to provide Droppable/Seekable information or use a separate enum?
#[derive(Clone, Debug, PartialEq)]
pub enum PictureType {
    UNKNOWN,
    I,
    P,
    B,
    S,
    SI,
    SP,
    SB,
    BI,
}

#[derive(Clone, Debug, PartialEq)]
pub struct VideoInfo {
    pub pic_type: PictureType,
    pub width: usize,
    pub height: usize,
    pub format: Rc<Formaton>,
}

impl VideoInfo {
    fn size(&self, align: usize) -> usize {
        let mut size = 0;
        for &component in self.format.into_iter() {
            if let Some(c) = component {
                size += c.get_data_size(self.width, self.height, align);
            }
        }
        size
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct AudioInfo {
    pub samples: usize,
    pub rate: usize,
    pub map: ChannelMap,
    pub format: Rc<Soniton>,
}

impl AudioInfo {
    fn size(&self, align: usize) -> usize {
        self.format.get_audio_size(self.samples, align) * self.map.len()
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum MediaKind {
    Video(VideoInfo),
    Audio(AudioInfo),
}

use std::convert::From;

impl From<VideoInfo> for MediaKind {
    fn from(v: VideoInfo) -> Self {
        MediaKind::Video(v)
    }
}

impl From<AudioInfo> for MediaKind {
    fn from(a: AudioInfo) -> Self {
        MediaKind::Audio(a)
    }
}

use self::MediaKind::*;

pub trait FrameBuffer {
    fn as_slice<'a>(&'a self, idx: usize) -> Result<&'a [u8]>;
    fn as_mut_slice<'a>(&'a mut self, idx: usize) -> Result<&'a mut [u8]>;
    fn linesize(&self, idx: usize) -> Result<usize>;
    fn count(&self) -> usize;
}

use std::fmt;

impl fmt::Debug for FrameBuffer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "FrameBuffer")
    }
}

#[derive(Debug)]
pub struct Frame {
    pub kind: MediaKind,
    pub buf: Box<FrameBuffer>,
    pub t: Option<TimeInfo>,
}

const ALIGNMENT: usize = 32;

struct Plane {
    buf: BytesMut,
    linesize: usize,
}

struct DefaultFrameBuffer {
    buf: BytesMut,
    planes: Vec<Plane>,
}

impl FrameBuffer for DefaultFrameBuffer {
    fn as_slice<'a>(&'a self, idx: usize) -> Result<&'a [u8]> {
        match self.planes.get(idx) {
            None => Err(Error::from_kind(ErrorKind::InvalidIndex)),
            Some(plane) => Ok(&plane.buf),
        }
    }
    fn as_mut_slice<'a>(&'a mut self, idx: usize) -> Result<&'a mut [u8]> {
        match self.planes.get_mut(idx) {
            None => Err(Error::from_kind(ErrorKind::InvalidIndex)),
            Some(plane) => Ok(&mut plane.buf),
        }
    }
    fn linesize(&self, idx: usize) -> Result<usize> {
        match self.planes.get(idx) {
            None => Err(Error::from_kind(ErrorKind::InvalidIndex)),
            Some(plane) => Ok(plane.linesize),
        }
    }
    fn count(&self) -> usize {
        self.planes.len()
    }
}

impl DefaultFrameBuffer {
    pub fn new<'a>(kind: &'a MediaKind) -> DefaultFrameBuffer {
        match kind {
            &Video(ref video) => {
                let size = video.size(ALIGNMENT);
                let data = unsafe {
                    Heap.alloc(Layout::from_size_align(size, ALIGNMENT).unwrap())
                        .unwrap()
                };
                //let data = unsafe { Heap.alloc_zeroed(Layout::from_size_align(size, ALIGNMENT)) };
                let buf = BytesMut::from(unsafe { Vec::from_raw_parts(data, size, size) });
                let mut buffer = DefaultFrameBuffer {
                    buf: buf,
                    planes: Vec::new(),
                };
                for &component in video.format.iter() {
                    if let Some(c) = component {
                        let planesize = c.get_data_size(video.width, video.height, ALIGNMENT);
                        let linesize = c.get_linesize(video.width, ALIGNMENT);
                        buffer.planes.push(Plane {
                            buf: buffer.buf.split_to(planesize),
                            linesize: linesize,
                        });
                    }
                }
                buffer
            },
            &Audio(ref audio) => {
                let size = audio.size(ALIGNMENT);
                let data = unsafe {
                    Heap.alloc(Layout::from_size_align(size, ALIGNMENT).unwrap())
                        .unwrap()
                };
                let buf = BytesMut::from(unsafe { Vec::from_raw_parts(data, size, size) });
                let mut buffer = DefaultFrameBuffer {
                    buf: buf,
                    planes: Vec::new(),
                };
                for _ in 0..audio.map.len() {
                    let size = audio.format.get_audio_size(audio.samples, ALIGNMENT);
                    buffer.planes.push(Plane {
                        buf: buffer.buf.split_to(size),
                        linesize: size,
                    });
                }
                buffer
            },
        }
    }
}

pub fn new_default_frame<T>(kind: T, t: Option<TimeInfo>) -> Frame
where T: Into<MediaKind> + Clone
{
    let k = kind.into();
    let buf = DefaultFrameBuffer::new(&k);

    Frame {
        kind: k,
        buf: box buf,
        t: t,
    }
}

use std::ptr::copy_nonoverlapping;

fn copy_plane(
    dst: &mut [u8],
    dst_linesize: usize,
    src: &[u8],
    src_linesize: usize,
    w: usize,
    h: usize,
) {
    let dst_chunks = dst.chunks_mut(dst_linesize);
    let src_chunks = src.chunks(src_linesize);

    for (d, s) in dst_chunks.zip(src_chunks).take(h) {
        unsafe {
            copy_nonoverlapping(s.as_ptr(), d.as_mut_ptr(), w);
        }
    }
}

fn copy_image<'a, IM, IU, I>(
    dst: IM,
    dst_linesizes: IU,
    src: I,
    src_linesizes: IU,
    w: usize,
    h: usize,
    fmt: &Formaton,
) where
    IM: Iterator<Item = &'a mut [u8]>,
    I: Iterator<Item = &'a [u8]>,
    IU: Iterator<Item = usize>,
{
    let dst_iter = dst.zip(dst_linesizes);
    let src_iter = src.zip(src_linesizes);
    let iter = dst_iter.zip(src_iter).zip(fmt.iter());

    for (((d, d_linesize), (s, s_linesize)), c) in iter {
        copy_plane(
            d,
            d_linesize,
            s,
            s_linesize,
            c.unwrap().get_width(w),
            c.unwrap().get_height(h),
        );
    }
}

// TODO: Add proper tests
fn copy_to_frame<'a, I, IU>(dst: &mut Frame, mut src: I, mut src_linesize: IU, w: usize, h: usize)
where
    I: Iterator<Item = &'a [u8]>,
    IU: Iterator<Item = usize>,
{
    if let MediaKind::Video(ref fmt) = dst.kind {
        let mut f_iter = fmt.format.iter();
        let w = fmt.width;
        let h = fmt.height;
        for i in 0..dst.buf.count() {
            let d_linesize = dst.buf.linesize(i).unwrap();
            let s_linesize = src_linesize.next().unwrap();
            let d = dst.buf.as_mut_slice(i).unwrap();
            let s = src.next().unwrap();
            let c = f_iter.next().unwrap();
            copy_plane(
                d,
                d_linesize,
                s,
                s_linesize,
                c.unwrap().get_width(w),
                c.unwrap().get_height(h),
            );
        }
    } else {
        unimplemented!();
    }
}

use std::mem;
use std::slice;

// TODO make it a separate trait
impl Frame {
    pub fn copy_from_slice<'a, I, IU>(&mut self, mut src: I, mut src_linesize: IU)
    where
        I: Iterator<Item = &'a [u8]>,
        IU: Iterator<Item = usize>,
    {
        if let MediaKind::Video(ref fmt) = self.kind {
            let mut f_iter = fmt.format.iter();
            let w = fmt.width;
            let h = fmt.height;
            for i in 0..self.buf.count() {
                let d_linesize = self.buf.linesize(i).unwrap();
                let s_linesize = src_linesize.next().unwrap();
                let d = self.buf.as_mut_slice(i).unwrap();
                let s = src.next().unwrap();
                let c = f_iter.next().unwrap();
                copy_plane(
                    d,
                    d_linesize,
                    s,
                    s_linesize,
                    c.unwrap().get_width(w),
                    c.unwrap().get_height(h),
                );
            }
        } else {
            unimplemented!();
        }
    }

    pub fn copy_from_raw_parts<'a, I, IU>(&mut self, mut src: I, mut src_linesize: IU)
    where
        I: Iterator<Item = *const u8>,
        IU: Iterator<Item = usize>,
    {
        if let MediaKind::Video(ref fmt) = self.kind {
            let mut f_iter = fmt.format.iter();
            let w = fmt.width;
            let h = fmt.height;
            for i in 0..self.buf.count() {
                let d_linesize = self.buf.linesize(i).unwrap();
                let s_linesize = src_linesize.next().unwrap();
                let d = self.buf.as_mut_slice(i).unwrap();
                let c = f_iter.next().unwrap();
                let r = src.next().unwrap();
                let hb = c.unwrap().get_height(h);
                let s = unsafe { slice::from_raw_parts(r, hb * s_linesize) };
                copy_plane(d, d_linesize, s, s_linesize, c.unwrap().get_width(w), hb);
                mem::forget(s);
            }
        } else {
            unimplemented!();
        }
    }
}
<commit_msg>frame: Add helpers to get slices of i16 and f32<commit_after>#![allow(dead_code, unused_variables)]
use alloc::heap::{Heap, Alloc, Layout};
use bytes::BytesMut;

use std::rc::Rc;

use audiosample::*;
use pixel::*;
use timeinfo::*;

pub mod error {
    error_chain! {
        errors {
            InvalidIndex
            InvalidConversion
        }
    }
}

use self::error::*;

// TODO: Document
// TODO: Change it to provide Droppable/Seekable information or use a separate enum?
#[derive(Clone, Debug, PartialEq)]
pub enum PictureType {
    UNKNOWN,
    I,
    P,
    B,
    S,
    SI,
    SP,
    SB,
    BI,
}

#[derive(Clone, Debug, PartialEq)]
pub struct VideoInfo {
    pub pic_type: PictureType,
    pub width: usize,
    pub height: usize,
    pub format: Rc<Formaton>,
}

impl VideoInfo {
    fn size(&self, align: usize) -> usize {
        let mut size = 0;
        for &component in self.format.into_iter() {
            if let Some(c) = component {
                size += c.get_data_size(self.width, self.height, align);
            }
        }
        size
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct AudioInfo {
    pub samples: usize,
    pub rate: usize,
    pub map: ChannelMap,
    pub format: Rc<Soniton>,
}

impl AudioInfo {
    fn size(&self, align: usize) -> usize {
        self.format.get_audio_size(self.samples, align) * self.map.len()
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum MediaKind {
    Video(VideoInfo),
    Audio(AudioInfo),
}

use std::convert::From;

impl From<VideoInfo> for MediaKind {
    fn from(v: VideoInfo) -> Self {
        MediaKind::Video(v)
    }
}

impl From<AudioInfo> for MediaKind {
    fn from(a: AudioInfo) -> Self {
        MediaKind::Audio(a)
    }
}

use self::MediaKind::*;

pub trait FrameBuffer {
    fn as_slice<'a>(&'a self, idx: usize) -> Result<&'a [u8]>;
    fn as_mut_slice<'a>(&'a mut self, idx: usize) -> Result<&'a mut [u8]>;
    fn linesize(&self, idx: usize) -> Result<usize>;
    fn count(&self) -> usize;
}

pub trait FrameBufferConv<T> : FrameBuffer {
    fn as_slice<'a>(&'a self, idx: usize) -> Result<&'a [T]> {
        let size = mem::size_of::<T>();
        if (self.linesize(idx)? % size) != 0 {
            Err(ErrorKind::InvalidConversion.into())
        } else {
            let s = FrameBuffer::as_slice(self, idx)?;
            let r = unsafe {
                slice::from_raw_parts::<T>(mem::transmute(s.as_ptr()), s.len() / size)
            };

            Ok(r)
        }
    }
    fn as_mut_slice<'a>(&'a mut self, idx: usize) -> Result<&'a mut [T]> {
        let size = mem::size_of::<T>();
        if (self.linesize(idx)? % size) != 0 {
            Err(ErrorKind::InvalidConversion.into())
        } else {
            let s = FrameBuffer::as_mut_slice(self, idx)?;
            let r = unsafe {
                slice::from_raw_parts_mut::<T>(mem::transmute(s.as_ptr()), s.len() / size)
            };

            Ok(r)
        }
    }
}

impl FrameBufferConv<i16> for FrameBuffer {}
impl FrameBufferConv<f32> for FrameBuffer {}

use std::fmt;

impl fmt::Debug for FrameBuffer {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "FrameBuffer")
    }
}

#[derive(Debug)]
pub struct Frame {
    pub kind: MediaKind,
    pub buf: Box<FrameBuffer>,
    pub t: Option<TimeInfo>,
}

const ALIGNMENT: usize = 32;

struct Plane {
    buf: BytesMut,
    linesize: usize,
}

struct DefaultFrameBuffer {
    buf: BytesMut,
    planes: Vec<Plane>,
}

impl FrameBuffer for DefaultFrameBuffer {
    fn as_slice<'a>(&'a self, idx: usize) -> Result<&'a [u8]> {
        match self.planes.get(idx) {
            None => Err(Error::from_kind(ErrorKind::InvalidIndex)),
            Some(plane) => Ok(&plane.buf),
        }
    }
    fn as_mut_slice<'a>(&'a mut self, idx: usize) -> Result<&'a mut [u8]> {
        match self.planes.get_mut(idx) {
            None => Err(Error::from_kind(ErrorKind::InvalidIndex)),
            Some(plane) => Ok(&mut plane.buf),
        }
    }
    fn linesize(&self, idx: usize) -> Result<usize> {
        match self.planes.get(idx) {
            None => Err(Error::from_kind(ErrorKind::InvalidIndex)),
            Some(plane) => Ok(plane.linesize),
        }
    }
    fn count(&self) -> usize {
        self.planes.len()
    }
}

impl DefaultFrameBuffer {
    pub fn new<'a>(kind: &'a MediaKind) -> DefaultFrameBuffer {
        match kind {
            &Video(ref video) => {
                let size = video.size(ALIGNMENT);
                let data = unsafe {
                    Heap.alloc(Layout::from_size_align(size, ALIGNMENT).unwrap())
                        .unwrap()
                };
                //let data = unsafe { Heap.alloc_zeroed(Layout::from_size_align(size, ALIGNMENT)) };
                let buf = BytesMut::from(unsafe { Vec::from_raw_parts(data, size, size) });
                let mut buffer = DefaultFrameBuffer {
                    buf: buf,
                    planes: Vec::new(),
                };
                for &component in video.format.iter() {
                    if let Some(c) = component {
                        let planesize = c.get_data_size(video.width, video.height, ALIGNMENT);
                        let linesize = c.get_linesize(video.width, ALIGNMENT);
                        buffer.planes.push(Plane {
                            buf: buffer.buf.split_to(planesize),
                            linesize: linesize,
                        });
                    }
                }
                buffer
            },
            &Audio(ref audio) => {
                let size = audio.size(ALIGNMENT);
                let data = unsafe {
                    Heap.alloc(Layout::from_size_align(size, ALIGNMENT).unwrap())
                        .unwrap()
                };
                let buf = BytesMut::from(unsafe { Vec::from_raw_parts(data, size, size) });
                let mut buffer = DefaultFrameBuffer {
                    buf: buf,
                    planes: Vec::new(),
                };
                for _ in 0..audio.map.len() {
                    let size = audio.format.get_audio_size(audio.samples, ALIGNMENT);
                    buffer.planes.push(Plane {
                        buf: buffer.buf.split_to(size),
                        linesize: size,
                    });
                }
                buffer
            },
        }
    }
}

pub fn new_default_frame<T>(kind: T, t: Option<TimeInfo>) -> Frame
where T: Into<MediaKind> + Clone
{
    let k = kind.into();
    let buf = DefaultFrameBuffer::new(&k);

    Frame {
        kind: k,
        buf: box buf,
        t: t,
    }
}

use std::ptr::copy_nonoverlapping;

fn copy_plane(
    dst: &mut [u8],
    dst_linesize: usize,
    src: &[u8],
    src_linesize: usize,
    w: usize,
    h: usize,
) {
    let dst_chunks = dst.chunks_mut(dst_linesize);
    let src_chunks = src.chunks(src_linesize);

    for (d, s) in dst_chunks.zip(src_chunks).take(h) {
        unsafe {
            copy_nonoverlapping(s.as_ptr(), d.as_mut_ptr(), w);
        }
    }
}

fn copy_image<'a, IM, IU, I>(
    dst: IM,
    dst_linesizes: IU,
    src: I,
    src_linesizes: IU,
    w: usize,
    h: usize,
    fmt: &Formaton,
) where
    IM: Iterator<Item = &'a mut [u8]>,
    I: Iterator<Item = &'a [u8]>,
    IU: Iterator<Item = usize>,
{
    let dst_iter = dst.zip(dst_linesizes);
    let src_iter = src.zip(src_linesizes);
    let iter = dst_iter.zip(src_iter).zip(fmt.iter());

    for (((d, d_linesize), (s, s_linesize)), c) in iter {
        copy_plane(
            d,
            d_linesize,
            s,
            s_linesize,
            c.unwrap().get_width(w),
            c.unwrap().get_height(h),
        );
    }
}

// TODO: Add proper tests
fn copy_to_frame<'a, I, IU>(dst: &mut Frame, mut src: I, mut src_linesize: IU, w: usize, h: usize)
where
    I: Iterator<Item = &'a [u8]>,
    IU: Iterator<Item = usize>,
{
    if let MediaKind::Video(ref fmt) = dst.kind {
        let mut f_iter = fmt.format.iter();
        let w = fmt.width;
        let h = fmt.height;
        for i in 0..dst.buf.count() {
            let d_linesize = dst.buf.linesize(i).unwrap();
            let s_linesize = src_linesize.next().unwrap();
            let d = dst.buf.as_mut_slice(i).unwrap();
            let s = src.next().unwrap();
            let c = f_iter.next().unwrap();
            copy_plane(
                d,
                d_linesize,
                s,
                s_linesize,
                c.unwrap().get_width(w),
                c.unwrap().get_height(h),
            );
        }
    } else {
        unimplemented!();
    }
}

use std::mem;
use std::slice;

// TODO make it a separate trait
impl Frame {
    pub fn copy_from_slice<'a, I, IU>(&mut self, mut src: I, mut src_linesize: IU)
    where
        I: Iterator<Item = &'a [u8]>,
        IU: Iterator<Item = usize>,
    {
        if let MediaKind::Video(ref fmt) = self.kind {
            let mut f_iter = fmt.format.iter();
            let w = fmt.width;
            let h = fmt.height;
            for i in 0..self.buf.count() {
                let d_linesize = self.buf.linesize(i).unwrap();
                let s_linesize = src_linesize.next().unwrap();
                let d = self.buf.as_mut_slice(i).unwrap();
                let s = src.next().unwrap();
                let c = f_iter.next().unwrap();
                copy_plane(
                    d,
                    d_linesize,
                    s,
                    s_linesize,
                    c.unwrap().get_width(w),
                    c.unwrap().get_height(h),
                );
            }
        } else {
            unimplemented!();
        }
    }

    pub fn copy_from_raw_parts<'a, I, IU>(&mut self, mut src: I, mut src_linesize: IU)
    where
        I: Iterator<Item = *const u8>,
        IU: Iterator<Item = usize>,
    {
        if let MediaKind::Video(ref fmt) = self.kind {
            let mut f_iter = fmt.format.iter();
            let w = fmt.width;
            let h = fmt.height;
            for i in 0..self.buf.count() {
                let d_linesize = self.buf.linesize(i).unwrap();
                let s_linesize = src_linesize.next().unwrap();
                let d = self.buf.as_mut_slice(i).unwrap();
                let c = f_iter.next().unwrap();
                let r = src.next().unwrap();
                let hb = c.unwrap().get_height(h);
                let s = unsafe { slice::from_raw_parts(r, hb * s_linesize) };
                copy_plane(d, d_linesize, s, s_linesize, c.unwrap().get_width(w), hb);
                mem::forget(s);
            }
        } else {
            unimplemented!();
        }
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use serialize_packed::{PackedOutputStream, PackedInputStream};
use io;

pub fn expect_packs_to(unpacked : &[u8],
                       packed : &[u8]) {

    use std::io::{Reader, Writer};

    // --------
    // write

    let mut bytes : std::vec::Vec<u8> = ::std::iter::repeat(0u8).take(packed.len()).collect();
    {
        let mut writer = io::ArrayOutputStream::new(bytes.as_mut_slice());
        let mut packed_output_stream = PackedOutputStream {inner : &mut writer};
        packed_output_stream.write(unpacked).unwrap();
        packed_output_stream.flush().unwrap();
    }

    assert!(bytes.as_slice().eq(packed),
            "expected: {}, got: {}", packed, bytes);

    // --------
    // read

    let mut reader = io::ArrayInputStream::new(packed);
    let mut packed_input_stream = PackedInputStream {inner : &mut reader};

    let bytes = packed_input_stream.read_exact(unpacked.len()).unwrap();

//    assert!(packed_input_stream.eof());
    assert!(bytes.slice(0, bytes.len()).eq(unpacked),
            "expected: {}, got: {}", unpacked, bytes);

}

static ZEROES : &'static[u8] = &[0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0];

#[test]
pub fn simple_packing() {
    expect_packs_to(&[], &[]);
    expect_packs_to(ZEROES.slice(0, 8), &[0,0]);
    expect_packs_to(&[0,0,12,0,0,34,0,0], &[0x24,12,34]);
    expect_packs_to(&[1,3,2,4,5,7,6,8], &[0xff,1,3,2,4,5,7,6,8,0]);
    expect_packs_to(&[0,0,0,0,0,0,0,0,1,3,2,4,5,7,6,8], &[0,0,0xff,1,3,2,4,5,7,6,8,0]);
    expect_packs_to(&[0,0,12,0,0,34,0,0,1,3,2,4,5,7,6,8], &[0x24,12,34,0xff,1,3,2,4,5,7,6,8,0]);
    expect_packs_to(&[1,3,2,4,5,7,6,8,8,6,7,4,5,2,3,1], &[0xff,1,3,2,4,5,7,6,8,1,8,6,7,4,5,2,3,1]);

    expect_packs_to(
        &[1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 0,2,4,0,9,0,5,1],
        &[0xff,1,2,3,4,5,6,7,8, 3, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8,
          0xd6,2,4,9,5,1]);
    expect_packs_to(
        &[1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 6,2,4,3,9,0,5,1, 1,2,3,4,5,6,7,8, 0,2,4,0,9,0,5,1],
        &[0xff,1,2,3,4,5,6,7,8, 3, 1,2,3,4,5,6,7,8, 6,2,4,3,9,0,5,1, 1,2,3,4,5,6,7,8,
          0xd6,2,4,9,5,1]);

    expect_packs_to(
        &[8,0,100,6,0,1,1,2, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,1,0,2,0,3,1],
        &[0xed,8,100,6,1,1,2, 0,2, 0xd4,1,2,3,1]);

    expect_packs_to(ZEROES.slice(0,16), &[0,1]);
    expect_packs_to(&[0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0], &[0,2]);

}
<commit_msg>get the tests to build again<commit_after>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use serialize_packed::{PackedOutputStream, PackedInputStream};
use io;

pub fn expect_packs_to(unpacked : &[u8],
                       packed : &[u8]) {

    use std::io::{Reader, Writer};

    // --------
    // write

    let mut bytes : std::vec::Vec<u8> = ::std::iter::repeat(0u8).take(packed.len()).collect();
    {
        let mut writer = io::ArrayOutputStream::new(bytes.as_mut_slice());
        let mut packed_output_stream = PackedOutputStream {inner : &mut writer};
        packed_output_stream.write(unpacked).unwrap();
        packed_output_stream.flush().unwrap();
    }

    assert!(bytes.as_slice().eq(packed));
//            "expected: {}, got: {}", packed, bytes.as_slice());

    // --------
    // read

    let mut reader = io::ArrayInputStream::new(packed);
    let mut packed_input_stream = PackedInputStream {inner : &mut reader};

    let bytes = packed_input_stream.read_exact(unpacked.len()).unwrap();

//    assert!(packed_input_stream.eof());
    assert!(bytes.slice(0, bytes.len()).eq(unpacked));
//            "expected: {}, got: {}", unpacked, bytes.as_slice());

}

static ZEROES : &'static[u8] = &[0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0];

#[test]
pub fn simple_packing() {
    expect_packs_to(&[], &[]);
    expect_packs_to(ZEROES.slice(0, 8), &[0,0]);
    expect_packs_to(&[0,0,12,0,0,34,0,0], &[0x24,12,34]);
    expect_packs_to(&[1,3,2,4,5,7,6,8], &[0xff,1,3,2,4,5,7,6,8,0]);
    expect_packs_to(&[0,0,0,0,0,0,0,0,1,3,2,4,5,7,6,8], &[0,0,0xff,1,3,2,4,5,7,6,8,0]);
    expect_packs_to(&[0,0,12,0,0,34,0,0,1,3,2,4,5,7,6,8], &[0x24,12,34,0xff,1,3,2,4,5,7,6,8,0]);
    expect_packs_to(&[1,3,2,4,5,7,6,8,8,6,7,4,5,2,3,1], &[0xff,1,3,2,4,5,7,6,8,1,8,6,7,4,5,2,3,1]);

    expect_packs_to(
        &[1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 0,2,4,0,9,0,5,1],
        &[0xff,1,2,3,4,5,6,7,8, 3, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8,
          0xd6,2,4,9,5,1]);
    expect_packs_to(
        &[1,2,3,4,5,6,7,8, 1,2,3,4,5,6,7,8, 6,2,4,3,9,0,5,1, 1,2,3,4,5,6,7,8, 0,2,4,0,9,0,5,1],
        &[0xff,1,2,3,4,5,6,7,8, 3, 1,2,3,4,5,6,7,8, 6,2,4,3,9,0,5,1, 1,2,3,4,5,6,7,8,
          0xd6,2,4,9,5,1]);

    expect_packs_to(
        &[8,0,100,6,0,1,1,2, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,1,0,2,0,3,1],
        &[0xed,8,100,6,1,1,2, 0,2, 0xd4,1,2,3,1]);

    expect_packs_to(ZEROES.slice(0,16), &[0,1]);
    expect_packs_to(&[0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0], &[0,2]);

}
<|endoftext|>


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

//! Standardized isolation of logic as a `process` wrapper around a `thread`.
//!
//! Basic wrapper around threads to provide some standard API. All long-living Hat threads run
//! as a `process`. The main difference between a normal Rust `thread` and a Hat `process` is
//! currently that the `process` has a bounded input-channel and a standard implementation of a
//! synchronous `send_reply()`.

/// A long-living `thread` is promoted to a standard `process`.
///
/// To create a new `process`, simply define a `Msg` type, a `Reply` type and a state struct
/// that implements the `MsgHandler` trait. You can then create a new process with:
///
/// ```rust,ignore
/// enum Msg {Ping, Foo}
/// enum Reply {Pong, Bar}
///
/// struct MyHandler {
///   fn handle(&mut self, msg: Msg, reply: |Reply|) {
///     match Msg {
///       Ping => return reply(Pong),
///       Foo => return reply(Bar),
///     }
///   }
/// }
///
/// fn main() {
///   let p = Process::new(MyHandler::new());
///   let pong = p.send_reply(Ping);
/// }
/// ```
///
/// The handler must always call `reply()` to release any call to `send_reply()`. It is allowed to
/// call `reply()` "early" - before all work has been done - but calling `reply()` is usually the
/// last thing a handler does. The idiom `return reply(...)` is a simple way of guaranteeing this.
///
/// A handler is allowed to call `reply()` **exactly** once. Not calling it or calling it multiple
/// times will likely cause runtime panicure when using `send_reply()`.

use std::fmt;
use std::thread;
use std::sync::{mpsc, Arc, Mutex};

pub struct Process<Msg, Reply, E> {
    sender: mpsc::SyncSender<(Msg, mpsc::Sender<Result<Reply, E>>)>,
    shutdown_after: Option<i64>,
}

/// When cloning a `process` we clone the input-channel, allowing multiple threads to share the same
/// `process`.
impl<Msg: Send, Reply: Send, E> Clone for Process<Msg, Reply, E> {
    fn clone(&self) -> Process<Msg, Reply, E> {
        Process {
            sender: self.sender.clone(),
            shutdown_after: None,
        }
    }
}

pub trait MsgHandler<Msg, Reply> {
    type Err;

    fn handle(&mut self,
              msg: Msg,
              callback: Box<Fn(Result<Reply, Self::Err>)>)
              -> Result<(), Self::Err>;
}

impl<Msg: 'static + Send, Reply: 'static + Send, E> Process<Msg, Reply, E> {
    /// Create and start a new process using `handler`.
    pub fn new<H, F>(handler_proc: F) -> Result<Process<Msg, Reply, E>, E>
        where H: 'static + MsgHandler<Msg, Reply, Err = E> + Send,
              E: 'static + From<mpsc::RecvError> + fmt::Debug + Send,
              F: FnOnce() -> Result<H, E>
    {
        Process::new_with_shutdown(handler_proc, None)
    }

    pub fn new_with_shutdown<H, F>(handler_proc: F,
                                   shutdown_after: Option<i64>)
                                   -> Result<Process<Msg, Reply, E>, E>
        where H: 'static + MsgHandler<Msg, Reply, Err = E> + Send,
              E: 'static + From<mpsc::RecvError> + fmt::Debug + Send,
              F: FnOnce() -> Result<H, E>
    {
        let (sender, receiver) = mpsc::sync_channel(10);
        let p = Process {
            sender: sender,
            shutdown_after: shutdown_after,
        };

        try!(p.start(receiver, handler_proc));

        Ok(p)
    }

    fn start<H, F>(&self,
                   receiver: mpsc::Receiver<(Msg, mpsc::Sender<Result<Reply, E>>)>,
                   handler_proc: F)
                   -> Result<(), E>
        where H: 'static + MsgHandler<Msg, Reply, Err = E> + Send,
              E: 'static + From<mpsc::RecvError> + fmt::Debug + Send,
              F: FnOnce() -> Result<H, E>
    {
        let shutdown_after = self.shutdown_after;
        let mut my_handler = try!(handler_proc());
        thread::spawn(move || {
            // fork handler
            let mut handle_msg = |msg, rep: mpsc::Sender<Result<Reply, E>>| {
                let did_reply = Arc::new(Mutex::new(false));
                let local_did_reply = did_reply.clone();
                let local_rep = rep.clone();
                let ret = my_handler.handle(msg,
                                            Box::new(move |r| {
                                                *local_did_reply.lock().unwrap() = true;
                                                local_rep.send(r)
                                                    .expect("Message reply ignored");
                                            }));
                match ret {
                    Ok(()) => assert_eq!(*did_reply.lock().unwrap(), true),
                    Err(e) => {
                        if *did_reply.lock().expect("Message reply ignored") {
                            return Err(format!("Encountered unrecoverable error: {:?}", e));
                        } else {
                            rep.send(Err(e)).expect("Message reply ignored");
                        }
                    }
                }

                Ok(())
            };
            let mut do_loop = || {
                match shutdown_after {
                    Some(msg_count) => {
                        for _ in 0..msg_count {
                            if let Ok((msg, rep)) = receiver.recv() {
                                try!(handle_msg(msg, rep))
                            } else {
                                return Ok(());  // Clean shutdown.
                            }
                        }
                        return Err("Process has reached message limit".to_string());
                    }
                    None => {
                        while let Ok((msg, rep)) = receiver.recv() {
                            try!(handle_msg(msg, rep))
                        }
                    }
                };
                Ok(())
            };
            if let Err(e) = do_loop() {
                error!("Shutting down process: {:?}", e);
            }
        });

        Ok(())
    }

    /// Synchronous send.
    ///
    /// Will always wait for a reply from the receiving `process`.
    pub fn send_reply(&self, msg: Msg) -> Result<Reply, E> where E: From<String> {
        let (sender, receiver) = mpsc::channel();
        if let Err(e) = self.sender.send((msg, sender)) {
            return Err(From::from("Could not send message: process is dead".to_string()));
        }

        match receiver.recv() {
            Err(e) => Err(From::from("Could not read reply: process has died".to_string())),
            Ok(reply) => reply,
        }
    }
}
<commit_msg>Let processes become poisoned instead of shutting down.<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Standardized isolation of logic as a `process` wrapper around a `thread`.
//!
//! Basic wrapper around threads to provide some standard API. All long-living Hat threads run
//! as a `process`. The main difference between a normal Rust `thread` and a Hat `process` is
//! currently that the `process` has a bounded input-channel and a standard implementation of a
//! synchronous `send_reply()`.

/// A long-living `thread` is promoted to a standard `process`.
///
/// To create a new `process`, simply define a `Msg` type, a `Reply` type and a state struct
/// that implements the `MsgHandler` trait. You can then create a new process with:
///
/// ```rust,ignore
/// enum Msg {Ping, Foo}
/// enum Reply {Pong, Bar}
///
/// struct MyHandler {
///   fn handle(&mut self, msg: Msg, reply: |Reply|) {
///     match Msg {
///       Ping => return reply(Pong),
///       Foo => return reply(Bar),
///     }
///   }
/// }
///
/// fn main() {
///   let p = Process::new(MyHandler::new());
///   let pong = p.send_reply(Ping);
/// }
/// ```
///
/// The handler must always call `reply()` to release any call to `send_reply()`. It is allowed to
/// call `reply()` "early" - before all work has been done - but calling `reply()` is usually the
/// last thing a handler does. The idiom `return reply(...)` is a simple way of guaranteeing this.
///
/// A handler is allowed to call `reply()` **exactly** once. Not calling it or calling it multiple
/// times will likely cause runtime panicure when using `send_reply()`.

use std::fmt;
use std::thread;
use std::sync::{mpsc, Arc, Mutex};

pub struct Process<Msg, Reply, E> {
    sender: mpsc::SyncSender<(Msg, mpsc::Sender<Result<Reply, E>>)>,
    poison_after: Option<i64>,
    poisoned: Arc<Mutex<bool>>,
}

/// When cloning a `process` we clone the input-channel, allowing multiple threads to share the same
/// `process`.
impl<Msg: Send, Reply: Send, E> Clone for Process<Msg, Reply, E> {
    fn clone(&self) -> Process<Msg, Reply, E> {
        Process {
            sender: self.sender.clone(),
            poison_after: None,
            poisoned: self.poisoned.clone(),
        }
    }
}

pub trait MsgHandler<Msg, Reply> {
    type Err;

    fn handle(&mut self,
              msg: Msg,
              callback: Box<Fn(Result<Reply, Self::Err>)>)
              -> Result<(), Self::Err>;
}

impl<Msg: 'static + Send, Reply: 'static + Send, E> Process<Msg, Reply, E> {
    /// Create and start a new process using `handler`.
    pub fn new<H, F>(handler_proc: F) -> Result<Process<Msg, Reply, E>, E>
        where H: 'static + MsgHandler<Msg, Reply, Err = E> + Send,
              E: 'static + From<mpsc::RecvError> + fmt::Debug + Send,
              F: FnOnce() -> Result<H, E>
    {
        Process::new_with_shutdown(handler_proc, None)
    }

    pub fn new_with_shutdown<H, F>(handler_proc: F,
                                   poison_after: Option<i64>)
                                   -> Result<Process<Msg, Reply, E>, E>
        where H: 'static + MsgHandler<Msg, Reply, Err = E> + Send,
              E: 'static + From<mpsc::RecvError> + fmt::Debug + Send,
              F: FnOnce() -> Result<H, E>
    {
        let (sender, receiver) = mpsc::sync_channel(10);
        let p = Process {
            sender: sender,
            poison_after: poison_after,
            poisoned: Arc::new(Mutex::new(false)),
        };

        try!(p.start(receiver, handler_proc));

        Ok(p)
    }

    fn start<H, F>(&self,
                   receiver: mpsc::Receiver<(Msg, mpsc::Sender<Result<Reply, E>>)>,
                   handler_proc: F)
                   -> Result<(), E>
        where H: 'static + MsgHandler<Msg, Reply, Err = E> + Send,
              E: 'static + From<mpsc::RecvError> + fmt::Debug + Send,
              F: FnOnce() -> Result<H, E>
    {
        let poison_after = self.poison_after;
        let mut my_handler = try!(handler_proc());

        let poisoned = self.poisoned.clone();

        thread::spawn(move || {
            // fork handler
            let mut handle_msg = |msg, rep: mpsc::Sender<Result<Reply, E>>| {
                let did_reply = Arc::new(Mutex::new(false));
                let local_did_reply = did_reply.clone();
                let local_rep = rep.clone();
                let ret = my_handler.handle(msg,
                                            Box::new(move |r| {
                                                *local_did_reply.lock().unwrap() = true;
                                                local_rep.send(r)
                                                    .expect("Message reply ignored");
                                            }));
                match ret {
                    Ok(()) => assert_eq!(*did_reply.lock().unwrap(), true),
                    Err(e) => {
                        if *did_reply.lock().expect("Message reply ignored") {
                            return Err(format!("Encountered unrecoverable error: {:?}", e));
                        } else {
                            rep.send(Err(e)).expect("Message reply ignored");
                        }
                    }
                }

                Ok(())
            };
            let mut do_loop = || {
                match poison_after {
                    Some(msg_count) => {
                        for _ in 0..msg_count {
                            if let Ok((msg, rep)) = receiver.recv() {
                                try!(handle_msg(msg, rep))
                            } else {
                                return Ok(());  // Clean shutdown.
                            }
                        }
                        return Err("Process has reached message limit".to_string());
                    }
                    None => {
                        while let Ok((msg, rep)) = receiver.recv() {
                            try!(handle_msg(msg, rep))
                        }
                    }
                };
                Ok(())
            };
            loop {
                if let Err(e) = do_loop() {
                    *poisoned.lock().unwrap() = true;  // This process is in a polluted state.
                    error!("Poisoned process: {:?}", e);
                } else { break }
            }
        });

        Ok(())
    }

    /// Synchronous send.
    ///
    /// Will always wait for a reply from the receiving `process`.
    pub fn send_reply(&self, msg: Msg) -> Result<Reply, E> where E: From<String> {
        let (sender, receiver) = mpsc::channel();

        if *self.poisoned.lock().unwrap() {
            return Err(From::from("Refusing to forward message to poisoned process".to_string()));
        }

        if let Err(e) = self.sender.send((msg, sender)) {
            return Err(From::from("Could not send message: process is dead".to_string()));
        }

        match receiver.recv() {
            Err(e) => Err(From::from("Could not read reply".to_string())),
            Ok(reply) => reply,
        }
    }
}
<|endoftext|>


<commit_before>use std::net::SocketAddr;
use std::collections::HashMap;
use std::io;

use solicit::http::session::StreamState;
use solicit::http::connection::HttpConnection;
use solicit::http::connection::EndStream;
use solicit::http::connection::SendFrame;
use solicit::http::frame::*;
use solicit::http::StreamId;
use solicit::http::HttpScheme;
use solicit::http::HttpError;
use solicit::http::StaticHeader;
use hpack;

use futures;
use futures::Future;
use futures::stream::Stream;

use tokio_core;
use tokio_core::net::TcpStream;
use tokio_core::io as tokio_io;
use tokio_core::io::Io;
use tokio_core::io::ReadHalf;
use tokio_core::io::WriteHalf;
use tokio_core::reactor;

use tokio_tls;

use futures_misc::*;

use solicit_async::*;
use solicit_misc::*;

use http_common::*;

struct GrpcHttpClientStream {
    state: StreamState,
    response_handler: Option<tokio_core::channel::Sender<ResultOrEof<HttpStreamPart, HttpError>>>,
}

impl GrpcHttpClientStream {

    fn _close(&mut self) {
        self.set_state(StreamState::Closed);
    }

    fn _close_local(&mut self) {
        let next = match self.state() {
            StreamState::HalfClosedRemote => StreamState::Closed,
            _ => StreamState::HalfClosedLocal,
        };
        self.set_state(next);
    }

    fn close_remote(&mut self) {
        let next = match self.state() {
            StreamState::HalfClosedLocal => StreamState::Closed,
            _ => StreamState::HalfClosedRemote,
        };
        self.set_state(next);
        if let Some(response_handler) = self.response_handler.take() {
            response_handler.send(ResultOrEof::Eof).unwrap();
        }
    }

    fn _is_closed(&self) -> bool {
        self.state() == StreamState::Closed
    }

    fn _is_closed_local(&self) -> bool {
        match self.state() {
            StreamState::HalfClosedLocal | StreamState::Closed => true,
            _ => false,
        }
    }

    fn _is_closed_remote(&self) -> bool {
        match self.state() {
            StreamState::HalfClosedRemote | StreamState::Closed => true,
            _ => false,
        }
    }

    fn set_state(&mut self, state: StreamState) {
        self.state = state;
    }

    fn state(&self) -> StreamState {
        self.state
    }
}

struct GrpcHttpClientSessionState {
    streams: HashMap<StreamId, GrpcHttpClientStream>,
    next_stream_id: StreamId,
    decoder: hpack::Decoder<'static>,
}

impl GrpcHttpClientSessionState {

    fn insert_stream(&mut self, stream: GrpcHttpClientStream) -> StreamId {
        let id = self.next_stream_id;
        self.streams.insert(id, stream);
        self.next_stream_id += 2;
        id
    }

    fn _get_stream_ref(&self, stream_id: StreamId) -> Option<&GrpcHttpClientStream> {
        self.streams.get(&stream_id)
    }

    fn get_stream_mut(&mut self, stream_id: StreamId) -> Option<&mut GrpcHttpClientStream> {
        self.streams.get_mut(&stream_id)
    }

    fn remove_stream(&mut self, stream_id: StreamId) -> Option<GrpcHttpClientStream> {
        self.streams.remove(&stream_id)
    }
}

struct ClientInner {
    conn: HttpConnection,
    to_write_tx: tokio_core::channel::Sender<ClientToWriteMessage>,
    session_state: GrpcHttpClientSessionState,
}

pub struct HttpClientConnectionAsync {
    call_tx: tokio_core::channel::Sender<ClientToWriteMessage>,
    remote: reactor::Remote,
}

unsafe impl  Sync for HttpClientConnectionAsync {}

struct StartRequestMessage {
    headers: Vec<StaticHeader>,
    body: HttpStreamSend<Vec<u8>>,
    response_handler: tokio_core::channel::Sender<ResultOrEof<HttpStreamPart, HttpError>>,
}

struct BodyChunkMessage {
    stream_id: StreamId,
    chunk: Vec<u8>,
}

struct EndRequestMessage {
    stream_id: StreamId,
}

struct ClientReadToWriteMessage {
    buf: Vec<u8>,
}

enum ClientToWriteMessage {
    Start(StartRequestMessage),
    BodyChunk(BodyChunkMessage),
    End(EndRequestMessage),
    FromRead(ClientReadToWriteMessage),
}

struct ClientWriteLoop<I : Io + Send + 'static> {
    write: WriteHalf<I>,
    inner: TaskRcMut<ClientInner>,
}

impl<I : Io + Send + 'static> ClientWriteLoop<I> {
    fn write_all(self, buf: Vec<u8>) -> HttpFuture<Self> {
        let ClientWriteLoop { write, inner } = self;

        Box::new(tokio_io::write_all(write, buf)
            .map(move |(write, _)| ClientWriteLoop { write: write, inner: inner })
            .map_err(HttpError::from))
    }

    fn process_from_read(self, message: ClientReadToWriteMessage) -> HttpFuture<Self> {
        self.write_all(message.buf)
    }

    fn process_start(self, start: StartRequestMessage) -> HttpFuture<Self> {
        let StartRequestMessage { headers, body, response_handler } = start;
        let (buf, stream_id) = self.inner.with(move |inner: &mut ClientInner| {

            let stream = GrpcHttpClientStream {
                state: StreamState::Open,
                response_handler: Some(response_handler),
            };
            let stream_id = inner.session_state.insert_stream(stream);

            let send_buf = inner.conn.send_headers_to_vec(stream_id, &headers, EndStream::No).unwrap();

            (send_buf, stream_id)
        });

        Box::new(self.write_all(buf)
            .and_then(move |wl: ClientWriteLoop<_>| {
                body.fold(wl, move |wl: ClientWriteLoop<_>, chunk| {
                    wl.inner.with(|inner: &mut ClientInner| {
                        inner.to_write_tx.send(ClientToWriteMessage::BodyChunk(BodyChunkMessage {
                            stream_id: stream_id,
                            chunk: chunk,
                        })).unwrap()
                    });
                    futures::finished::<_, HttpError>(wl)
                })
            })
            .and_then(move |wl: ClientWriteLoop<_>| {
                wl.inner.with(|inner: &mut ClientInner| {
                    inner.to_write_tx.send(ClientToWriteMessage::End(EndRequestMessage {
                        stream_id: stream_id,
                    })).unwrap()
                });
                futures::finished::<_, HttpError>(wl)
            }))
    }

    fn process_body_chunk(self, body_chunk: BodyChunkMessage) -> HttpFuture<Self> {
        let BodyChunkMessage { stream_id, chunk } = body_chunk;

        let buf = self.inner.with(move |inner: &mut ClientInner| {
            {
                let _stream = inner.session_state.get_stream_mut(stream_id);
                // TODO: check stream state
            }

            inner.conn.send_data_frames_to_vec(stream_id, &chunk, EndStream::No).unwrap()
        });

        self.write_all(buf)
    }

    fn process_end(self, end: EndRequestMessage) -> HttpFuture<Self> {
        let EndRequestMessage { stream_id } = end;

        let buf = self.inner.with(move |inner: &mut ClientInner| {
            inner.conn.send_end_of_stream_to_vec(stream_id).unwrap()
        });

        self.write_all(buf)
    }

    fn process_message(self, message: ClientToWriteMessage) -> HttpFuture<Self> {
        match message {
            ClientToWriteMessage::Start(start) => self.process_start(start),
            ClientToWriteMessage::BodyChunk(body_chunk) => self.process_body_chunk(body_chunk),
            ClientToWriteMessage::End(end) => self.process_end(end),
            ClientToWriteMessage::FromRead(from_read) => self.process_from_read(from_read),
        }
    }

    fn run(self, requests: HttpStreamSend<ClientToWriteMessage>) -> HttpFuture<()> {
        let requests = requests.map_err(HttpError::from);
        Box::new(requests
            .fold(self, move |wl, message: ClientToWriteMessage| {
                wl.process_message(message)
            })
            .map(|_| ()))
    }
}

struct ClientReadLoop<I : Io + Send + 'static> {
    read: ReadHalf<I>,
    inner: TaskRcMut<ClientInner>,
}

impl<I : Io + Send + 'static> HttpReadLoop for ClientReadLoop<I> {
    fn recv_raw_frame(self) -> HttpFuture<(Self, RawFrame<'static>)> {
        let ClientReadLoop { read, inner } = self;
        Box::new(recv_raw_frame(read)
            .map(|(read, frame)| (ClientReadLoop { read: read, inner: inner}, frame))
            .map_err(HttpError::from))
    }

    fn process_data_frame(self, frame: DataFrame) -> HttpFuture<Self> {
        self.inner.with(|inner: &mut ClientInner| {
            let mut stream: &mut GrpcHttpClientStream = match inner.session_state.get_stream_mut(frame.get_stream_id()) {
                None => {
                    // TODO(mlalic): This can currently indicate two things:
                    //                 1) the stream was idle => PROTOCOL_ERROR
                    //                 2) the stream was closed => STREAM_CLOSED (stream error)
                    return;
                }
                Some(stream) => stream,
            };

            if let Some(ref mut response_handler) = stream.response_handler {
                response_handler.send(ResultOrEof::Item(HttpStreamPart {
                    content: HttpStreamPartContent::Data(frame.data.clone().into_owned()),
                    last: frame.is_end_of_stream(),
                })).unwrap();
            }

            if frame.is_end_of_stream() {
                stream.close_remote();
                // TODO: GC streams
            }
        });
        Box::new(futures::finished(self))
    }

    fn process_headers_frame(self, frame: HeadersFrame) -> HttpFuture<Self> {
        self.inner.with(|inner: &mut ClientInner| {
            let headers = inner.session_state.decoder
                                   .decode(&frame.header_fragment())
                                   .map_err(HttpError::CompressionError).unwrap();
            let headers: Vec<StaticHeader> = headers.into_iter().map(|h| h.into()).collect();

            let mut stream: &mut GrpcHttpClientStream = match inner.session_state.get_stream_mut(frame.get_stream_id()) {
                None => {
                    // TODO(mlalic): This means that the server's header is not associated to any
                    //               request made by the client nor any server-initiated stream (pushed)
                    return;
                }
                Some(stream) => stream,
            };
            // TODO: hack
            if headers.len() != 0 {

                if let Some(ref mut response_handler) = stream.response_handler {
                    response_handler.send(ResultOrEof::Item(HttpStreamPart {
                        content: HttpStreamPartContent::Headers(headers),
                        last: frame.is_end_of_stream(),
                    })).unwrap();
                }
            }

            if frame.is_end_of_stream() {
                stream.close_remote();
                // TODO: GC streams
            }
        });
        Box::new(futures::finished(self))
    }

    fn process_rst_stream_frame(self, frame: RstStreamFrame) -> HttpFuture<Self> {
        self.inner.with(move |inner: &mut ClientInner| {
            // TODO: check actually removed
            inner.session_state.remove_stream(frame.get_stream_id());
        });

        Box::new(futures::finished(self))
    }

    fn process_window_update_frame(self, _frame: WindowUpdateFrame) -> HttpFuture<Self> {
        // TODO
        Box::new(futures::finished(self))
    }

    fn process_settings_global(self, _frame: SettingsFrame) -> HttpFuture<Self> {
        // TODO: apply settings
        // TODO: send ack
        Box::new(futures::finished(self))
    }

    fn send_frame<R : FrameIR>(self, frame: R) -> HttpFuture<Self> {
        self.inner.with(|inner: &mut ClientInner| {
            let mut send_buf = VecSendFrame(Vec::new());
            send_buf.send_frame(frame).unwrap();
            inner.to_write_tx.send(ClientToWriteMessage::FromRead(ClientReadToWriteMessage { buf: send_buf.0 }))
                .expect("read to write");
        });

        Box::new(futures::finished(self))
    }

    fn process_conn_window_update(self, _frame: WindowUpdateFrame) -> HttpFuture<Self> {
        // TODO
        Box::new(futures::finished(self))
    }
}

impl HttpClientConnectionAsync {
    fn connected<I : Io + Send + 'static>(lh: reactor::Handle, connect: HttpFutureSend<I>) -> (Self, HttpFuture<()>) {
        let (to_write_tx, to_write_rx) = tokio_core::channel::channel(&lh).unwrap();

        let to_write_rx = Box::new(to_write_rx.map_err(HttpError::from));

        let c = HttpClientConnectionAsync {
            remote: lh.remote().clone(),
            call_tx: to_write_tx.clone(),
        };

        let handshake = connect.and_then(client_handshake);

        let future = handshake.and_then(move |conn| {
            trace!("handshake done");
            let (read, write) = conn.split();

            let inner = TaskRcMut::new(ClientInner {
                conn: HttpConnection::new(HttpScheme::Http),
                to_write_tx: to_write_tx.clone(),
                session_state: GrpcHttpClientSessionState {
                    next_stream_id: 1,
                    streams: HashMap::new(),
                    decoder: hpack::Decoder::new(),
                }
            });

            let write_inner = inner.clone();
            let run_write = ClientWriteLoop { write: write, inner: write_inner }.run(to_write_rx);
            let run_read = ClientReadLoop { read: read, inner: inner.clone() }.run();

            run_write.join(run_read).map(|_| ())
        });

        (c, Box::new(future))
    }

    pub fn new_plain(lh: reactor::Handle, addr: &SocketAddr) -> (Self, HttpFuture<()>) {
        let connect = TcpStream::connect(&addr, &lh)
            .map_err(|e| e.into());

        HttpClientConnectionAsync::connected(lh, Box::new(connect))
    }

    pub fn new_tls(lh: reactor::Handle, addr: &SocketAddr) -> (Self, HttpFuture<()>) {
        let connect = TcpStream::connect(&addr, &lh)
            .map_err(|e| e.into());

        let tls_conn = connect.and_then(|conn| {
            tokio_tls::ClientContext::new()
                .unwrap()
                .handshake("localhost", conn)
        });

        let tls_conn = tls_conn.map_err(HttpError::from);

        HttpClientConnectionAsync::connected(lh, Box::new(tls_conn))
    }

    pub fn start_request(
        &self,
        headers: Vec<StaticHeader>,
        body: HttpStreamSend<Vec<u8>>)
            -> HttpStreamStreamSend
    {
        let (tx, rx) = futures::oneshot();

        let call_tx = self.call_tx.clone();

        self.remote.spawn(move |handle| {
            let (req_tx, req_rx) = tokio_core::channel::channel(&handle).unwrap();

            call_tx.send(ClientToWriteMessage::Start(StartRequestMessage {
                headers: headers,
                body: body,
                response_handler: req_tx,
            })).unwrap();

            let req_rx = req_rx.map_err(HttpError::from);

            tx.complete(stream_with_eof_and_error(req_rx, || HttpError::from(io::Error::new(io::ErrorKind::Other, "unexpected eof"))));

            Ok(())
        });

        let rx = rx.map_err(|_| HttpError::from(io::Error::new(io::ErrorKind::Other, "oneshot canceled")));

        Box::new(rx.flatten_stream())
    }
}
<commit_msg>logging<commit_after>use std::net::SocketAddr;
use std::collections::HashMap;
use std::io;

use solicit::http::session::StreamState;
use solicit::http::connection::HttpConnection;
use solicit::http::connection::EndStream;
use solicit::http::connection::SendFrame;
use solicit::http::frame::*;
use solicit::http::StreamId;
use solicit::http::HttpScheme;
use solicit::http::HttpError;
use solicit::http::StaticHeader;
use hpack;

use futures;
use futures::Future;
use futures::stream::Stream;

use tokio_core;
use tokio_core::net::TcpStream;
use tokio_core::io as tokio_io;
use tokio_core::io::Io;
use tokio_core::io::ReadHalf;
use tokio_core::io::WriteHalf;
use tokio_core::reactor;

use tokio_tls;

use futures_misc::*;

use solicit_async::*;
use solicit_misc::*;

use http_common::*;

struct GrpcHttpClientStream {
    state: StreamState,
    response_handler: Option<tokio_core::channel::Sender<ResultOrEof<HttpStreamPart, HttpError>>>,
}

impl GrpcHttpClientStream {

    fn _close(&mut self) {
        self.set_state(StreamState::Closed);
    }

    fn _close_local(&mut self) {
        let next = match self.state() {
            StreamState::HalfClosedRemote => StreamState::Closed,
            _ => StreamState::HalfClosedLocal,
        };
        self.set_state(next);
    }

    fn close_remote(&mut self) {
        let next = match self.state() {
            StreamState::HalfClosedLocal => StreamState::Closed,
            _ => StreamState::HalfClosedRemote,
        };
        self.set_state(next);
        if let Some(response_handler) = self.response_handler.take() {
            response_handler.send(ResultOrEof::Eof).unwrap();
        }
    }

    fn _is_closed(&self) -> bool {
        self.state() == StreamState::Closed
    }

    fn _is_closed_local(&self) -> bool {
        match self.state() {
            StreamState::HalfClosedLocal | StreamState::Closed => true,
            _ => false,
        }
    }

    fn _is_closed_remote(&self) -> bool {
        match self.state() {
            StreamState::HalfClosedRemote | StreamState::Closed => true,
            _ => false,
        }
    }

    fn set_state(&mut self, state: StreamState) {
        self.state = state;
    }

    fn state(&self) -> StreamState {
        self.state
    }
}

struct GrpcHttpClientSessionState {
    streams: HashMap<StreamId, GrpcHttpClientStream>,
    next_stream_id: StreamId,
    decoder: hpack::Decoder<'static>,
}

impl GrpcHttpClientSessionState {

    fn insert_stream(&mut self, stream: GrpcHttpClientStream) -> StreamId {
        let id = self.next_stream_id;
        self.streams.insert(id, stream);
        self.next_stream_id += 2;
        id
    }

    fn _get_stream_ref(&self, stream_id: StreamId) -> Option<&GrpcHttpClientStream> {
        self.streams.get(&stream_id)
    }

    fn get_stream_mut(&mut self, stream_id: StreamId) -> Option<&mut GrpcHttpClientStream> {
        self.streams.get_mut(&stream_id)
    }

    fn remove_stream(&mut self, stream_id: StreamId) -> Option<GrpcHttpClientStream> {
        self.streams.remove(&stream_id)
    }
}

struct ClientInner {
    conn: HttpConnection,
    to_write_tx: tokio_core::channel::Sender<ClientToWriteMessage>,
    session_state: GrpcHttpClientSessionState,
}

pub struct HttpClientConnectionAsync {
    call_tx: tokio_core::channel::Sender<ClientToWriteMessage>,
    remote: reactor::Remote,
}

unsafe impl  Sync for HttpClientConnectionAsync {}

struct StartRequestMessage {
    headers: Vec<StaticHeader>,
    body: HttpStreamSend<Vec<u8>>,
    response_handler: tokio_core::channel::Sender<ResultOrEof<HttpStreamPart, HttpError>>,
}

struct BodyChunkMessage {
    stream_id: StreamId,
    chunk: Vec<u8>,
}

struct EndRequestMessage {
    stream_id: StreamId,
}

struct ClientReadToWriteMessage {
    buf: Vec<u8>,
}

enum ClientToWriteMessage {
    Start(StartRequestMessage),
    BodyChunk(BodyChunkMessage),
    End(EndRequestMessage),
    FromRead(ClientReadToWriteMessage),
}

struct ClientWriteLoop<I : Io + Send + 'static> {
    write: WriteHalf<I>,
    inner: TaskRcMut<ClientInner>,
}

impl<I : Io + Send + 'static> ClientWriteLoop<I> {
    fn write_all(self, buf: Vec<u8>) -> HttpFuture<Self> {
        let ClientWriteLoop { write, inner } = self;

        Box::new(tokio_io::write_all(write, buf)
            .map(move |(write, _)| ClientWriteLoop { write: write, inner: inner })
            .map_err(HttpError::from))
    }

    fn process_from_read(self, message: ClientReadToWriteMessage) -> HttpFuture<Self> {
        self.write_all(message.buf)
    }

    fn process_start(self, start: StartRequestMessage) -> HttpFuture<Self> {
        let StartRequestMessage { headers, body, response_handler } = start;
        let (buf, stream_id) = self.inner.with(move |inner: &mut ClientInner| {

            let stream = GrpcHttpClientStream {
                state: StreamState::Open,
                response_handler: Some(response_handler),
            };
            let stream_id = inner.session_state.insert_stream(stream);

            let send_buf = inner.conn.send_headers_to_vec(stream_id, &headers, EndStream::No).unwrap();

            (send_buf, stream_id)
        });

        Box::new(self.write_all(buf)
            .and_then(move |wl: ClientWriteLoop<_>| {
                body.fold(wl, move |wl: ClientWriteLoop<_>, chunk| {
                    wl.inner.with(|inner: &mut ClientInner| {
                        inner.to_write_tx.send(ClientToWriteMessage::BodyChunk(BodyChunkMessage {
                            stream_id: stream_id,
                            chunk: chunk,
                        })).unwrap()
                    });
                    futures::finished::<_, HttpError>(wl)
                })
            })
            .and_then(move |wl: ClientWriteLoop<_>| {
                wl.inner.with(|inner: &mut ClientInner| {
                    inner.to_write_tx.send(ClientToWriteMessage::End(EndRequestMessage {
                        stream_id: stream_id,
                    })).unwrap()
                });
                futures::finished::<_, HttpError>(wl)
            }))
    }

    fn process_body_chunk(self, body_chunk: BodyChunkMessage) -> HttpFuture<Self> {
        let BodyChunkMessage { stream_id, chunk } = body_chunk;

        let buf = self.inner.with(move |inner: &mut ClientInner| {
            {
                let _stream = inner.session_state.get_stream_mut(stream_id);
                // TODO: check stream state
            }

            inner.conn.send_data_frames_to_vec(stream_id, &chunk, EndStream::No).unwrap()
        });

        self.write_all(buf)
    }

    fn process_end(self, end: EndRequestMessage) -> HttpFuture<Self> {
        let EndRequestMessage { stream_id } = end;

        let buf = self.inner.with(move |inner: &mut ClientInner| {
            inner.conn.send_end_of_stream_to_vec(stream_id).unwrap()
        });

        self.write_all(buf)
    }

    fn process_message(self, message: ClientToWriteMessage) -> HttpFuture<Self> {
        match message {
            ClientToWriteMessage::Start(start) => self.process_start(start),
            ClientToWriteMessage::BodyChunk(body_chunk) => self.process_body_chunk(body_chunk),
            ClientToWriteMessage::End(end) => self.process_end(end),
            ClientToWriteMessage::FromRead(from_read) => self.process_from_read(from_read),
        }
    }

    fn run(self, requests: HttpStreamSend<ClientToWriteMessage>) -> HttpFuture<()> {
        let requests = requests.map_err(HttpError::from);
        Box::new(requests
            .fold(self, move |wl, message: ClientToWriteMessage| {
                wl.process_message(message)
            })
            .map(|_| ()))
    }
}

struct ClientReadLoop<I : Io + Send + 'static> {
    read: ReadHalf<I>,
    inner: TaskRcMut<ClientInner>,
}

impl<I : Io + Send + 'static> HttpReadLoop for ClientReadLoop<I> {
    fn recv_raw_frame(self) -> HttpFuture<(Self, RawFrame<'static>)> {
        let ClientReadLoop { read, inner } = self;
        Box::new(recv_raw_frame(read)
            .map(|(read, frame)| (ClientReadLoop { read: read, inner: inner}, frame))
            .map_err(HttpError::from))
    }

    fn process_data_frame(self, frame: DataFrame) -> HttpFuture<Self> {
        self.inner.with(|inner: &mut ClientInner| {
            let mut stream: &mut GrpcHttpClientStream = match inner.session_state.get_stream_mut(frame.get_stream_id()) {
                None => {
                    // TODO(mlalic): This can currently indicate two things:
                    //                 1) the stream was idle => PROTOCOL_ERROR
                    //                 2) the stream was closed => STREAM_CLOSED (stream error)
                    return;
                }
                Some(stream) => stream,
            };

            if let Some(ref mut response_handler) = stream.response_handler {
                response_handler.send(ResultOrEof::Item(HttpStreamPart {
                    content: HttpStreamPartContent::Data(frame.data.clone().into_owned()),
                    last: frame.is_end_of_stream(),
                })).unwrap();
            }

            if frame.is_end_of_stream() {
                stream.close_remote();
                // TODO: GC streams
            }
        });
        Box::new(futures::finished(self))
    }

    fn process_headers_frame(self, frame: HeadersFrame) -> HttpFuture<Self> {
        self.inner.with(|inner: &mut ClientInner| {
            let headers = inner.session_state.decoder
                                   .decode(&frame.header_fragment())
                                   .map_err(HttpError::CompressionError).unwrap();
            let headers: Vec<StaticHeader> = headers.into_iter().map(|h| h.into()).collect();

            let mut stream: &mut GrpcHttpClientStream = match inner.session_state.get_stream_mut(frame.get_stream_id()) {
                None => {
                    // TODO(mlalic): This means that the server's header is not associated to any
                    //               request made by the client nor any server-initiated stream (pushed)
                    return;
                }
                Some(stream) => stream,
            };
            // TODO: hack
            if headers.len() != 0 {

                if let Some(ref mut response_handler) = stream.response_handler {
                    response_handler.send(ResultOrEof::Item(HttpStreamPart {
                        content: HttpStreamPartContent::Headers(headers),
                        last: frame.is_end_of_stream(),
                    })).unwrap();
                }
            }

            if frame.is_end_of_stream() {
                stream.close_remote();
                // TODO: GC streams
            }
        });
        Box::new(futures::finished(self))
    }

    fn process_rst_stream_frame(self, frame: RstStreamFrame) -> HttpFuture<Self> {
        self.inner.with(move |inner: &mut ClientInner| {
            // TODO: check actually removed
            inner.session_state.remove_stream(frame.get_stream_id());
        });

        Box::new(futures::finished(self))
    }

    fn process_window_update_frame(self, _frame: WindowUpdateFrame) -> HttpFuture<Self> {
        // TODO
        Box::new(futures::finished(self))
    }

    fn process_settings_global(self, _frame: SettingsFrame) -> HttpFuture<Self> {
        // TODO: apply settings
        // TODO: send ack
        Box::new(futures::finished(self))
    }

    fn send_frame<R : FrameIR>(self, frame: R) -> HttpFuture<Self> {
        self.inner.with(|inner: &mut ClientInner| {
            let mut send_buf = VecSendFrame(Vec::new());
            send_buf.send_frame(frame).unwrap();
            inner.to_write_tx.send(ClientToWriteMessage::FromRead(ClientReadToWriteMessage { buf: send_buf.0 }))
                .expect("read to write");
        });

        Box::new(futures::finished(self))
    }

    fn process_conn_window_update(self, _frame: WindowUpdateFrame) -> HttpFuture<Self> {
        // TODO
        Box::new(futures::finished(self))
    }
}

impl HttpClientConnectionAsync {
    fn connected<I : Io + Send + 'static>(lh: reactor::Handle, connect: HttpFutureSend<I>) -> (Self, HttpFuture<()>) {
        let (to_write_tx, to_write_rx) = tokio_core::channel::channel(&lh).unwrap();

        let to_write_rx = Box::new(to_write_rx.map_err(HttpError::from));

        let c = HttpClientConnectionAsync {
            remote: lh.remote().clone(),
            call_tx: to_write_tx.clone(),
        };

        let handshake = connect.and_then(client_handshake);

        let future = handshake.and_then(move |conn| {
            trace!("handshake done");
            let (read, write) = conn.split();

            let inner = TaskRcMut::new(ClientInner {
                conn: HttpConnection::new(HttpScheme::Http),
                to_write_tx: to_write_tx.clone(),
                session_state: GrpcHttpClientSessionState {
                    next_stream_id: 1,
                    streams: HashMap::new(),
                    decoder: hpack::Decoder::new(),
                }
            });

            let write_inner = inner.clone();
            let run_write = ClientWriteLoop { write: write, inner: write_inner }.run(to_write_rx);
            let run_read = ClientReadLoop { read: read, inner: inner.clone() }.run();

            run_write.join(run_read).map(|_| ())
        });

        (c, Box::new(future))
    }

    pub fn new_plain(lh: reactor::Handle, addr: &SocketAddr) -> (Self, HttpFuture<()>) {
        let addr = addr.clone();

        let connect = TcpStream::connect(&addr, &lh)
            .map(move |c| { info!("connected to {}", addr); c })
            .map_err(|e| e.into());

        HttpClientConnectionAsync::connected(lh, Box::new(connect))
    }

    pub fn new_tls(lh: reactor::Handle, addr: &SocketAddr) -> (Self, HttpFuture<()>) {
        let addr = addr.clone();

        let connect = TcpStream::connect(&addr, &lh)
            .map(move |c| { info!("connected to {}", addr); c })
            .map_err(|e| e.into());

        let tls_conn = connect.and_then(|conn| {
            tokio_tls::ClientContext::new()
                .unwrap()
                .handshake("localhost", conn)
        });

        let tls_conn = tls_conn.map_err(HttpError::from);

        HttpClientConnectionAsync::connected(lh, Box::new(tls_conn))
    }

    pub fn start_request(
        &self,
        headers: Vec<StaticHeader>,
        body: HttpStreamSend<Vec<u8>>)
            -> HttpStreamStreamSend
    {
        let (tx, rx) = futures::oneshot();

        let call_tx = self.call_tx.clone();

        self.remote.spawn(move |handle| {
            let (req_tx, req_rx) = tokio_core::channel::channel(&handle).unwrap();

            call_tx.send(ClientToWriteMessage::Start(StartRequestMessage {
                headers: headers,
                body: body,
                response_handler: req_tx,
            })).unwrap();

            let req_rx = req_rx.map_err(HttpError::from);

            tx.complete(stream_with_eof_and_error(req_rx, || HttpError::from(io::Error::new(io::ErrorKind::Other, "unexpected eof"))));

            Ok(())
        });

        let rx = rx.map_err(|_| HttpError::from(io::Error::new(io::ErrorKind::Other, "oneshot canceled")));

        Box::new(rx.flatten_stream())
    }
}
<|endoftext|>


<commit_before>// Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

#![deny(warnings)]
// Enable all clippy lints except for many of the pedantic ones. It's a shame this needs to be copied and pasted across crates, but there doesn't appear to be a way to include inner attributes from a common source.
#![deny(
  clippy::all,
  clippy::default_trait_access,
  clippy::expl_impl_clone_on_copy,
  clippy::if_not_else,
  clippy::needless_continue,
  clippy::single_match_else,
  clippy::unseparated_literal_suffix,
  clippy::used_underscore_binding
)]
// It is often more clear to show that nothing is being moved.
#![allow(clippy::match_ref_pats)]
// Subjective style.
#![allow(
  clippy::len_without_is_empty,
  clippy::redundant_field_names,
  clippy::too_many_arguments
)]
// Default isn't as big a deal as people seem to think it is.
#![allow(clippy::new_without_default, clippy::new_ret_no_self)]
// Arc<Mutex> can be more clear than needing to grok Orderings:
#![allow(clippy::mutex_atomic)]

use futures::task_local;
use parking_lot::Mutex;
use rand::thread_rng;
use rand::Rng;
use std::collections::HashSet;
use std::sync::Arc;
use time::Timespec;

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct WorkUnit {
  pub name: String,
  pub start_timestamp: Timespec,
  pub end_timestamp: Timespec,
  pub span_id: String,
  pub parent_id: Option<String>,
}

#[derive(Clone)]
pub struct WorkUnitStore {
  workunits: Arc<Mutex<HashSet<WorkUnit>>>,
}

impl WorkUnitStore {
  pub fn new() -> WorkUnitStore {
    WorkUnitStore {
      workunits: Arc::new(Mutex::new(HashSet::new())),
    }
  }

  pub fn get_workunits(&self) -> Arc<Mutex<HashSet<WorkUnit>>> {
    self.workunits.clone()
  }

  pub fn add_workunit(&self, workunit: WorkUnit) {
    self.workunits.lock().insert(workunit);
  }
}

pub fn generate_random_64bit_string() -> String {
  let mut rng = thread_rng();
  let random_u64: u64 = rng.gen();
  format!("{:16.x}", random_u64)
}

pub fn workunits_with_constant_span_id(workunit_store: &WorkUnitStore) -> HashSet<WorkUnit> {
  //  This function is for the test purpose.

  workunit_store
    .get_workunits()
    .lock()
    .iter()
    .map(|workunit| WorkUnit {
      span_id: String::from("ignore"),
      ..workunit.clone()
    })
    .collect()
}

task_local! {
  static TASK_PARENT_ID: Mutex<Option<String>> = Mutex::new(None)
}

pub fn set_parent_id(parent_id: String) {
  TASK_PARENT_ID.with(|task_parent_id| {
    *task_parent_id.lock() = Some(parent_id);
  })
}

pub fn get_parent_id() -> Option<String> {
  TASK_PARENT_ID.with(|task_parent_id| {
    let task_parent_id = task_parent_id.lock();
    (*task_parent_id).clone()
  })
}
<commit_msg>Fix span id bug (#8129)<commit_after>// Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

#![deny(warnings)]
// Enable all clippy lints except for many of the pedantic ones. It's a shame this needs to be copied and pasted across crates, but there doesn't appear to be a way to include inner attributes from a common source.
#![deny(
  clippy::all,
  clippy::default_trait_access,
  clippy::expl_impl_clone_on_copy,
  clippy::if_not_else,
  clippy::needless_continue,
  clippy::single_match_else,
  clippy::unseparated_literal_suffix,
  clippy::used_underscore_binding
)]
// It is often more clear to show that nothing is being moved.
#![allow(clippy::match_ref_pats)]
// Subjective style.
#![allow(
  clippy::len_without_is_empty,
  clippy::redundant_field_names,
  clippy::too_many_arguments
)]
// Default isn't as big a deal as people seem to think it is.
#![allow(clippy::new_without_default, clippy::new_ret_no_self)]
// Arc<Mutex> can be more clear than needing to grok Orderings:
#![allow(clippy::mutex_atomic)]

use futures::task_local;
use parking_lot::Mutex;
use rand::thread_rng;
use rand::Rng;
use std::collections::HashSet;
use std::sync::Arc;
use time::Timespec;

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct WorkUnit {
  pub name: String,
  pub start_timestamp: Timespec,
  pub end_timestamp: Timespec,
  pub span_id: String,
  pub parent_id: Option<String>,
}

#[derive(Clone)]
pub struct WorkUnitStore {
  workunits: Arc<Mutex<HashSet<WorkUnit>>>,
}

impl WorkUnitStore {
  pub fn new() -> WorkUnitStore {
    WorkUnitStore {
      workunits: Arc::new(Mutex::new(HashSet::new())),
    }
  }

  pub fn get_workunits(&self) -> Arc<Mutex<HashSet<WorkUnit>>> {
    self.workunits.clone()
  }

  pub fn add_workunit(&self, workunit: WorkUnit) {
    self.workunits.lock().insert(workunit);
  }
}

pub fn generate_random_64bit_string() -> String {
  let mut rng = thread_rng();
  let random_u64: u64 = rng.gen();
  hex_16_digit_string(random_u64)
}

fn hex_16_digit_string(number: u64) -> String {
  format!("{:016.x}", number)
}

pub fn workunits_with_constant_span_id(workunit_store: &WorkUnitStore) -> HashSet<WorkUnit> {
  //  This function is for the test purpose.

  workunit_store
    .get_workunits()
    .lock()
    .iter()
    .map(|workunit| WorkUnit {
      span_id: String::from("ignore"),
      ..workunit.clone()
    })
    .collect()
}

task_local! {
  static TASK_PARENT_ID: Mutex<Option<String>> = Mutex::new(None)
}

pub fn set_parent_id(parent_id: String) {
  TASK_PARENT_ID.with(|task_parent_id| {
    *task_parent_id.lock() = Some(parent_id);
  })
}

pub fn get_parent_id() -> Option<String> {
  TASK_PARENT_ID.with(|task_parent_id| {
    let task_parent_id = task_parent_id.lock();
    (*task_parent_id).clone()
  })
}

#[cfg(test)]
mod tests {
  use crate::hex_16_digit_string;

  #[test]
  fn workunit_span_id_has_16_digits_len_hex_format() {
    let number: u64 = 1;
    let hex_string = hex_16_digit_string(number);
    assert_eq!(16, hex_string.len());
    for ch in hex_string.chars() {
      assert!(ch.is_ascii_hexdigit())
    }
  }
}
<|endoftext|>


<commit_before>use arch;
use core::cell::UnsafeCell;
use core::marker::Sync;
use core::ops::{Drop, Deref, DerefMut};
use core::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};

/// This type provides MUTual EXclusion based on spinning.
///
/// # Description
///
/// This structure behaves a lot like a normal Mutex. There are some differences:
///
/// - It may be used outside the runtime.
///   - A normal mutex will fail when used without the runtime, this will just lock
///   - When the runtime is present, it will call the deschedule function when appropriate
/// - No lock poisoning. When a fail occurs when the lock is held, no guarantees are made
///
/// When calling rust functions from bare threads, such as C `pthread`s, this lock will be very
/// helpful. In other cases however, you are encouraged to use the locks from the standard
/// library.
///
/// # Simple example
///
/// ```
/// use spin;
/// let spin_mutex = spin::Mutex::new(0);
///
/// // Modify the data
/// {
///     let mut data = spin_mutex.lock();
///     *data = 2;
/// }
///
/// // Read the data
/// let answer =
/// {
///     let data = spin_mutex.lock();
///     *data
/// };
///
/// assert_eq!(answer, 2);
/// ```
///
/// # Thread-safety example
///
/// ```
/// use spin;
/// use core::sync::{Arc, Barrier};
///
/// let numthreads = 1000;
/// let spin_mutex = Arc::new(spin::Mutex::new(0));
///
/// // We use a barrier to ensure the readout happens after all writing
/// let barrier = Arc::new(Barrier::new(numthreads + 1));
///
/// for _ in (0..numthreads)
/// {
///     let my_barrier = barrier.clone();
///     let my_lock = spin_mutex.clone();
///     std::thread::spawn(move||
///     {
///         let mut guard = my_lock.lock();
///         *guard += 1;
///
///         // Release the lock to prevent a deadlock
///         drop(guard);
///         my_barrier.wait();
///     });
/// }
///
/// barrier.wait();
///
/// let answer = { *spin_mutex.lock() };
/// assert_eq!(answer, numthreads);
/// ```
pub struct Mutex<T>
{
    lock: AtomicUsize,
    data: UnsafeCell<T>,
}

/// A guard to which the protected data can be accessed
///
/// When the guard falls out of scope it will release the lock.
pub struct MutexGuard<'a, T:'a>
{
    lock: &'a AtomicUsize,
    token: usize,
    data: &'a mut T,
}

unsafe impl<T> Sync for Mutex<T> {}

/// A Mutex which may be used statically.
///
/// ```
/// use mutex::{self, STATIC_MUTEX_INIT};
///
/// static SPLCK: spin::StaticMutex = STATIC_MUTEX_INIT;
///
/// fn demo() {
///     let lock = SPLCK.lock();
///     // do something with lock
///     drop(lock);
/// }
/// ```
pub type StaticMutex = Mutex<()>;

/// A initializer for StaticMutex, containing no data.
pub const STATIC_MUTEX_INIT: StaticMutex = Mutex {
    lock: ATOMIC_USIZE_INIT,
    data: UnsafeCell::new(()),
};

impl<T> Mutex<T>
{
    /// Creates a new spinlock wrapping the supplied data.
    pub fn new(user_data: T) -> Mutex<T>
    {
        Mutex
        {
            lock: ATOMIC_USIZE_INIT,
            data: UnsafeCell::new(user_data),
        }
    }

    fn obtain_lock(&self, file_line: *const (&'static str, u32)) -> usize
    {
        let token = arch::disable_interrupts();
        let null = 0;
        while self.lock.compare_and_swap(null, file_line as usize, Ordering::SeqCst) != null
        {
            // Do nothing
        }
        token
    }

    /// Locks the spinlock and returns a guard.
    ///
    /// The returned value may be dereferenced for data access
    /// and the lock will be dropped when the guard falls out of scope.
    ///
    /// ```
    /// let mylock = spin::Mutex::new(0);
    /// {
    ///     let mut data = mylock.lock();
    ///     // The lock is now locked and the data can be accessed
    ///     *data += 1;
    ///     // The lock is implicitly dropped
    /// }
    ///
    /// ```
    pub fn lock(&self, file_line: &'static (&'static str, u32)) -> MutexGuard<T>
    {
        let token = self.obtain_lock(file_line);
        MutexGuard
        {
            lock: &self.lock,
            token: token,
            data: unsafe { &mut *self.data.get() },
        }
    }

    /// Deallocates resources associated with this static mutex.
    ///
    /// This method is unsafe because it provides no guarantees that there are
    /// no active users of this mutex, and safety is not guaranteed if there are
    /// active users of this mutex.
    ///
    /// This method is required to ensure that there are no memory leaks on
    /// *all* platforms. It may be the case that some platforms do not leak
    /// memory if this method is not called, but this is not guaranteed to be
    /// true on all platforms.
    pub unsafe fn destroy(&'static self) {
        // nothing to do
    }
}

impl<'a, T> Deref for MutexGuard<'a, T>
{
    type Target = T;
    fn deref<'b>(&'b self) -> &'b T { &*self.data }
}

impl<'a, T> DerefMut for MutexGuard<'a, T>
{
    fn deref_mut<'b>(&'b mut self) -> &'b mut T { &mut *self.data }
}

impl<'a, T> Drop for MutexGuard<'a, T>
{
    /// The dropping of the MutexGuard will release the lock it was created from.
    fn drop(&mut self)
    {
        self.lock.store(0, Ordering::SeqCst);
        arch::restore_interrupts(self.token);
    }
}

#[macro_export]
macro_rules! lock {
    ($mutex:expr) => ({
        static FILE_LINE: (&'static str, u32) = (file!(), line!());
        $mutex.lock(&FILE_LINE)
    })
}
<commit_msg>Unsafe Sync/Send impls are consistent with std::sync::Mutex<commit_after>use arch;
use core::cell::UnsafeCell;
use core::marker::Sync;
use core::ops::{Drop, Deref, DerefMut};
use core::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};

/// This type provides MUTual EXclusion based on spinning.
///
/// # Description
///
/// This structure behaves a lot like a normal Mutex. There are some differences:
///
/// - It may be used outside the runtime.
///   - A normal mutex will fail when used without the runtime, this will just lock
///   - When the runtime is present, it will call the deschedule function when appropriate
/// - No lock poisoning. When a fail occurs when the lock is held, no guarantees are made
///
/// When calling rust functions from bare threads, such as C `pthread`s, this lock will be very
/// helpful. In other cases however, you are encouraged to use the locks from the standard
/// library.
///
/// # Simple example
///
/// ```
/// use spin;
/// let spin_mutex = spin::Mutex::new(0);
///
/// // Modify the data
/// {
///     let mut data = spin_mutex.lock();
///     *data = 2;
/// }
///
/// // Read the data
/// let answer =
/// {
///     let data = spin_mutex.lock();
///     *data
/// };
///
/// assert_eq!(answer, 2);
/// ```
///
/// # Thread-safety example
///
/// ```
/// use spin;
/// use core::sync::{Arc, Barrier};
///
/// let numthreads = 1000;
/// let spin_mutex = Arc::new(spin::Mutex::new(0));
///
/// // We use a barrier to ensure the readout happens after all writing
/// let barrier = Arc::new(Barrier::new(numthreads + 1));
///
/// for _ in (0..numthreads)
/// {
///     let my_barrier = barrier.clone();
///     let my_lock = spin_mutex.clone();
///     std::thread::spawn(move||
///     {
///         let mut guard = my_lock.lock();
///         *guard += 1;
///
///         // Release the lock to prevent a deadlock
///         drop(guard);
///         my_barrier.wait();
///     });
/// }
///
/// barrier.wait();
///
/// let answer = { *spin_mutex.lock() };
/// assert_eq!(answer, numthreads);
/// ```
pub struct Mutex<T>
{
    lock: AtomicUsize,
    data: UnsafeCell<T>,
}

/// A guard to which the protected data can be accessed
///
/// When the guard falls out of scope it will release the lock.
pub struct MutexGuard<'a, T:'a>
{
    lock: &'a AtomicUsize,
    token: usize,
    data: &'a mut T,
}

unsafe impl<T: Send> Send for Mutex<T> { }
unsafe impl<T: Send> Sync for Mutex<T> { }

/// A Mutex which may be used statically.
///
/// ```
/// use mutex::{self, STATIC_MUTEX_INIT};
///
/// static SPLCK: spin::StaticMutex = STATIC_MUTEX_INIT;
///
/// fn demo() {
///     let lock = SPLCK.lock();
///     // do something with lock
///     drop(lock);
/// }
/// ```
pub type StaticMutex = Mutex<()>;

/// A initializer for StaticMutex, containing no data.
pub const STATIC_MUTEX_INIT: StaticMutex = Mutex {
    lock: ATOMIC_USIZE_INIT,
    data: UnsafeCell::new(()),
};

impl<T> Mutex<T>
{
    /// Creates a new spinlock wrapping the supplied data.
    pub fn new(user_data: T) -> Mutex<T>
    {
        Mutex
        {
            lock: ATOMIC_USIZE_INIT,
            data: UnsafeCell::new(user_data),
        }
    }

    fn obtain_lock(&self, file_line: *const (&'static str, u32)) -> usize
    {
        let token = arch::disable_interrupts();
        let null = 0;
        while self.lock.compare_and_swap(null, file_line as usize, Ordering::SeqCst) != null
        {
            // Do nothing
        }
        token
    }

    /// Locks the spinlock and returns a guard.
    ///
    /// The returned value may be dereferenced for data access
    /// and the lock will be dropped when the guard falls out of scope.
    ///
    /// ```
    /// let mylock = spin::Mutex::new(0);
    /// {
    ///     let mut data = mylock.lock();
    ///     // The lock is now locked and the data can be accessed
    ///     *data += 1;
    ///     // The lock is implicitly dropped
    /// }
    ///
    /// ```
    pub fn lock(&self, file_line: &'static (&'static str, u32)) -> MutexGuard<T>
    {
        let token = self.obtain_lock(file_line);
        MutexGuard
        {
            lock: &self.lock,
            token: token,
            data: unsafe { &mut *self.data.get() },
        }
    }

    /// Deallocates resources associated with this static mutex.
    ///
    /// This method is unsafe because it provides no guarantees that there are
    /// no active users of this mutex, and safety is not guaranteed if there are
    /// active users of this mutex.
    ///
    /// This method is required to ensure that there are no memory leaks on
    /// *all* platforms. It may be the case that some platforms do not leak
    /// memory if this method is not called, but this is not guaranteed to be
    /// true on all platforms.
    pub unsafe fn destroy(&'static self) {
        // nothing to do
    }
}

impl<'a, T> Deref for MutexGuard<'a, T>
{
    type Target = T;
    fn deref<'b>(&'b self) -> &'b T { &*self.data }
}

impl<'a, T> DerefMut for MutexGuard<'a, T>
{
    fn deref_mut<'b>(&'b mut self) -> &'b mut T { &mut *self.data }
}

impl<'a, T> Drop for MutexGuard<'a, T>
{
    /// The dropping of the MutexGuard will release the lock it was created from.
    fn drop(&mut self)
    {
        self.lock.store(0, Ordering::SeqCst);
        arch::restore_interrupts(self.token);
    }
}

#[macro_export]
macro_rules! lock {
    ($mutex:expr) => ({
        static FILE_LINE: (&'static str, u32) = (file!(), line!());
        $mutex.lock(&FILE_LINE)
    })
}
<|endoftext|>