<commit_before>use std::from_str::FromStr;

use super::ScanCursor;
use super::{ScanError, OtherScanError};

macro_rules! from_str_scanner {
	($scan_fn:path -> $T:ty as $name:expr) => {
		impl<'a> Scanner<'a> for $T {
			fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<($T, Cur), ScanError> {
				let err = |cur:&Cur| Err(OtherScanError(format!(concat!("expected ",$name,", got `{}`"), cur.tail_str()), cursor.consumed()));

				let end = match $scan_fn(cursor.tail_str()) {
					Some(i) => i,
					None => return err(cursor)
				};

				let s = cursor.str_slice_to(end);
				let cursor = cursor.slice_from(end);

				FromStr::from_str(s)
					.map(|i| Ok((i, cursor.clone())))
					.unwrap_or_else(|| err(&cursor))
			}
		}
	};
}

pub trait Scanner<'a> {
	fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<(Self, Cur), ScanError>;
}

impl<'a> Scanner<'a> for bool {
	fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<(bool, Cur), ScanError> {
		cursor.expect_tok("true").map(|c| (true, c))
			.or_else(|_| cursor.expect_tok("false").map(|c| (false, c)))
			.or_else(|_| Err(OtherScanError(format!("expected `true` or `false`, got `{}`", cursor.tail_str()), cursor.consumed())))
	}
}

impl<'a> Scanner<'a> for char {
	fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<(char, Cur), ScanError> {
		let s = cursor.tail_str();
		if s.len() == 0 {
			Err(OtherScanError("expected character".into_string(), cursor.consumed()))
		} else {
			let ::std::str::CharRange { ch, next } = s.char_range_at(0);
			Ok((ch, cursor.slice_from(next)))
		}
	}
}

impl<'a> Scanner<'a> for &'a str {
	fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<(&'a str, Cur), ScanError> {
		cursor.pop_token().map(|sc| Ok(sc))
			.unwrap_or_else(|| Err(OtherScanError("expected token".into_string(), cursor.consumed())))
	}
}

from_str_scanner! { scan_float -> f32 as "real number" }
from_str_scanner! { scan_float -> f64 as "real number" }
from_str_scanner! { scan_int -> i8 as "8-bit integer" }
from_str_scanner! { scan_int -> i16 as "16-bit integer" }
from_str_scanner! { scan_int -> i32 as "32-bit integer" }
from_str_scanner! { scan_int -> i64 as "64-bit integer" }
from_str_scanner! { scan_int -> int as "integer" }
from_str_scanner! { scan_uint -> u8 as "8-bit unsigned integer" }
from_str_scanner! { scan_uint -> u16 as "16-bit unsigned integer" }
from_str_scanner! { scan_uint -> u32 as "32-bit unsigned integer" }
from_str_scanner! { scan_uint -> u64 as "64-bit unsigned integer" }
from_str_scanner! { scan_uint -> uint as "unsigned integer" }

fn next_char_at(s: &str, at: uint) -> uint {
	let ::std::str::CharRange { ch: _, next } = s.char_range_at(at);
	next
}

fn scan_float(s: &str) -> Option<uint> {
	enum State {
		Start,
		Whole,
		Suffix,
		ExponentStart,
		Exponent,
	}

	let mut state = Start;

	for (i,c) in s.char_indices() {
		match state {
			Start => {
				assert!(i == 0);
				match c {
					'0'...'9' | '-' => state = Whole,
					_ => return None
				}
			},
			Whole => match c {
				'0'...'9' => (),
				'.' => state = Suffix,
				'e' | 'E' => state = ExponentStart,
				_ => return Some(i)
			},
			Suffix => match c {
				'0'...'9' => (),
				'e' | 'E' => state = ExponentStart,
				_ => return Some(i),
			},
			ExponentStart => match c {
				'+' | '-' | '0'...'9' => state = Exponent,
				_ => return Some(i)
			},
			Exponent => match c {
				'0'...'9' => (),
				_ => return Some(i)
			}
		}
	}

	return Some(s.len())
}

fn scan_uint<'a>(s: &'a str) -> Option<uint> {
	s.char_indices()
		.take_while(|&(_,c)| '0' <= c && c <= '9')
		.map(|(i,_)| Some(next_char_at(s, i)))
		.last().unwrap_or(None)
}

fn scan_int<'a>(s: &'a str) -> Option<uint> {
	if s.len() == 0 { return None }

	let (s, off) = if s.char_at(0) == '-' {
		(s.slice_from(1), 1)
	} else {
		(s, 0)
	};

	scan_uint(s).map(|end| end+off)
}

#[cfg(test)]
mod test {
	use Cursor;
	use ScanError;
	use super::Scanner;
	use tokenizer::WordsAndInts;
	use whitespace::Ignore;

	fn cur<'a>(s: &'a str) -> Cursor<'a, WordsAndInts, Ignore> {
		Cursor::new(s, WordsAndInts, Ignore)
	}

	fn scan_a<'a, T: Scanner<'a, T>>(s: &'a str) -> Result<(T, Cursor<'a, WordsAndInts, Ignore>), ScanError> {
		Scanner::scan(&cur(s))
	}

	#[test]
	fn test_bool() {
		assert!(scan_a::<bool>("").err().is_some());
		assert!(scan_a::<bool>("true").unwrap().0 == true);
		assert!(scan_a::<bool>("false").unwrap().0 == false);
		assert!(scan_a::<bool>("yes").err().is_some());
		assert!(scan_a::<bool>("no").err().is_some());
		assert!(scan_a::<bool>("on").err().is_some());
		assert!(scan_a::<bool>("off").err().is_some());
		assert!(scan_a::<bool>("1").err().is_some());
		assert!(scan_a::<bool>("0").err().is_some());
	}

	#[test]
	fn test_char() {
		assert!(scan_a::<char>("").err().is_some());
		assert!(scan_a::<char>("x").unwrap().0 == 'x');
		assert!(scan_a::<char>("xy").unwrap().0 == 'x');
		assert!(scan_a::<char>("日本語").unwrap().0 == '日');
	}

	#[test]
	fn test_floats() {
		use std::from_str::FromStr;
		use std::num::{Float, Zero};

		fn test<'a, F: Float + Scanner<'a, F> + FromStr>() {
			let f = |v:F| v;
			let fs = |s:&str| -> F FromStr::from_str(s).unwrap();

			assert!(scan_a::<F>("").err().is_some());
			assert!(scan_a::<F>("0").ok().unwrap().0 == f(Zero::zero()));
			assert!(scan_a::<F>("0.0").ok().unwrap().0 == f(Zero::zero()));
			assert!(scan_a::<F>("-0").ok().unwrap().0 == -f(Zero::zero()));
			assert!(scan_a::<F>("1.0").ok().unwrap().0 == fs("1.0"));
			assert!(scan_a::<F>("1.00").ok().unwrap().0 == fs("1.0"));
			assert!(scan_a::<F>("1.0e0").ok().unwrap().0 == fs("1.0"));
			assert!(scan_a::<F>("1.0e1").ok().unwrap().0 == fs("10.0"));
		}

		test::<f32>();
		test::<f64>();
	}

	#[test]
	fn test_sized_ints() {
		use std::fmt::Show;
		use std::num::{Bounded, Int, Zero};

		fn test<'a, I: Int + Scanner<'a, I> + Show>(check_past: bool) {
			let zero: I = Zero::zero();
			let min: I = Bounded::min_value();
			let max: I = Bounded::max_value();

			assert!(scan_a::<I>("").err().is_some());
			assert!(scan_a::<I>("0").ok().unwrap().0 == zero);
			assert!(scan_a::<I>(format!("{}", min).as_slice()).ok().unwrap().0 == min);
			assert!(scan_a::<I>(format!("{}", max).as_slice()).ok().unwrap().0 == max);

			if check_past {
				let past_min: i64 = (min.to_i64()).unwrap() - 1;
				let past_max: u64 = (max.to_u64()).unwrap() + 1;
				assert!(scan_a::<I>(format!("{}", past_min).as_slice()).err().is_some());
				assert!(scan_a::<I>(format!("{}", past_max).as_slice()).err().is_some());
			}
		}

		test::<i8>(true);
		test::<i16>(true);
		test::<i32>(true);
		test::<i64>(false);
		test::<u8>(true);
		test::<u16>(true);
		test::<u32>(true);
		test::<u64>(false);
	}

	#[test]
	fn test_int() {
		assert!(scan_a::<int>("").err().is_some());
		assert!(scan_a::<int>("0").ok().unwrap().0 == 0);
		assert!(scan_a::<int>("42").ok().unwrap().0 == 42);
		assert!(scan_a::<int>("1_234").ok().unwrap().0 == 1);
		assert!(scan_a::<int>("x").err().is_some());
		assert!(scan_a::<int>("0x").ok().unwrap().0 == 0);
		assert!(scan_a::<int>("42x").ok().unwrap().0 == 42);
		assert!(scan_a::<int>("-").err().is_some());
		assert!(scan_a::<int>("-0").ok().unwrap().0 == 0);
		assert!(scan_a::<int>("-42").ok().unwrap().0 == -42);
		assert!(scan_a::<int>("-1_234").ok().unwrap().0 == -1);
	}

	#[test]
	fn test_uint() {
		assert!(scan_a::<uint>("").err().is_some());
		assert!(scan_a::<uint>("0").ok().unwrap().0 == 0);
		assert!(scan_a::<uint>("42").ok().unwrap().0 == 42);
		assert!(scan_a::<uint>("1_234").ok().unwrap().0 == 1);
		assert!(scan_a::<uint>("x").err().is_some());
		assert!(scan_a::<uint>("0x").ok().unwrap().0 == 0);
		assert!(scan_a::<uint>("42x").ok().unwrap().0 == 42);
		assert!(scan_a::<uint>("-").err().is_some());
		assert!(scan_a::<uint>("-0").err().is_some());
		assert!(scan_a::<uint>("-42").err().is_some());
		assert!(scan_a::<uint>("-1_234").err().is_some());
	}

	#[test]
	fn test_str() {
		assert!(scan_a::<&str>("").err().is_some());
		assert!(scan_a::<&str>("a").ok().unwrap().0 == "a");
		assert!(scan_a::<&str>("a b").ok().unwrap().0 == "a");
		assert!(scan_a::<&str>("abc").ok().unwrap().0 == "abc");
		assert!(scan_a::<&str>("ab-c").ok().unwrap().0 == "ab");
	}
}
<commit_msg>Fixed scanner tests being broken.<commit_after>use std::from_str::FromStr;

use super::ScanCursor;
use super::{ScanError, OtherScanError};

macro_rules! from_str_scanner {
	($scan_fn:path -> $T:ty as $name:expr) => {
		impl<'a> Scanner<'a> for $T {
			fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<($T, Cur), ScanError> {
				let err = |cur:&Cur| Err(OtherScanError(format!(concat!("expected ",$name,", got `{}`"), cur.tail_str()), cursor.consumed()));

				let end = match $scan_fn(cursor.tail_str()) {
					Some(i) => i,
					None => return err(cursor)
				};

				let s = cursor.str_slice_to(end);
				let cursor = cursor.slice_from(end);

				FromStr::from_str(s)
					.map(|i| Ok((i, cursor.clone())))
					.unwrap_or_else(|| err(&cursor))
			}
		}
	};
}

pub trait Scanner<'a> {
	fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<(Self, Cur), ScanError>;
}

impl<'a> Scanner<'a> for bool {
	fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<(bool, Cur), ScanError> {
		cursor.expect_tok("true").map(|c| (true, c))
			.or_else(|_| cursor.expect_tok("false").map(|c| (false, c)))
			.or_else(|_| Err(OtherScanError(format!("expected `true` or `false`, got `{}`", cursor.tail_str()), cursor.consumed())))
	}
}

impl<'a> Scanner<'a> for char {
	fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<(char, Cur), ScanError> {
		let s = cursor.tail_str();
		if s.len() == 0 {
			Err(OtherScanError("expected character".into_string(), cursor.consumed()))
		} else {
			let ::std::str::CharRange { ch, next } = s.char_range_at(0);
			Ok((ch, cursor.slice_from(next)))
		}
	}
}

impl<'a> Scanner<'a> for &'a str {
	fn scan<Cur: ScanCursor<'a>>(cursor: &Cur) -> Result<(&'a str, Cur), ScanError> {
		cursor.pop_token().map(|sc| Ok(sc))
			.unwrap_or_else(|| Err(OtherScanError("expected token".into_string(), cursor.consumed())))
	}
}

from_str_scanner! { scan_float -> f32 as "real number" }
from_str_scanner! { scan_float -> f64 as "real number" }
from_str_scanner! { scan_int -> i8 as "8-bit integer" }
from_str_scanner! { scan_int -> i16 as "16-bit integer" }
from_str_scanner! { scan_int -> i32 as "32-bit integer" }
from_str_scanner! { scan_int -> i64 as "64-bit integer" }
from_str_scanner! { scan_int -> int as "integer" }
from_str_scanner! { scan_uint -> u8 as "8-bit unsigned integer" }
from_str_scanner! { scan_uint -> u16 as "16-bit unsigned integer" }
from_str_scanner! { scan_uint -> u32 as "32-bit unsigned integer" }
from_str_scanner! { scan_uint -> u64 as "64-bit unsigned integer" }
from_str_scanner! { scan_uint -> uint as "unsigned integer" }

fn next_char_at(s: &str, at: uint) -> uint {
	let ::std::str::CharRange { ch: _, next } = s.char_range_at(at);
	next
}

fn scan_float(s: &str) -> Option<uint> {
	enum State {
		Start,
		Whole,
		Suffix,
		ExponentStart,
		Exponent,
	}

	let mut state = Start;

	for (i,c) in s.char_indices() {
		match state {
			Start => {
				assert!(i == 0);
				match c {
					'0'...'9' | '-' => state = Whole,
					_ => return None
				}
			},
			Whole => match c {
				'0'...'9' => (),
				'.' => state = Suffix,
				'e' | 'E' => state = ExponentStart,
				_ => return Some(i)
			},
			Suffix => match c {
				'0'...'9' => (),
				'e' | 'E' => state = ExponentStart,
				_ => return Some(i),
			},
			ExponentStart => match c {
				'+' | '-' | '0'...'9' => state = Exponent,
				_ => return Some(i)
			},
			Exponent => match c {
				'0'...'9' => (),
				_ => return Some(i)
			}
		}
	}

	return Some(s.len())
}

fn scan_uint<'a>(s: &'a str) -> Option<uint> {
	s.char_indices()
		.take_while(|&(_,c)| '0' <= c && c <= '9')
		.map(|(i,_)| Some(next_char_at(s, i)))
		.last().unwrap_or(None)
}

fn scan_int<'a>(s: &'a str) -> Option<uint> {
	if s.len() == 0 { return None }

	let (s, off) = if s.char_at(0) == '-' {
		(s.slice_from(1), 1)
	} else {
		(s, 0)
	};

	scan_uint(s).map(|end| end+off)
}

#[cfg(test)]
mod test {
	use Cursor;
	use ScanError;
	use super::Scanner;
	use tokenizer::WordsAndInts;
	use whitespace::Ignore;

	fn cur<'a>(s: &'a str) -> Cursor<'a, WordsAndInts, Ignore> {
		Cursor::new(s, WordsAndInts, Ignore)
	}

	fn scan_a<'a, T: Scanner<'a>>(s: &'a str) -> Result<(T, Cursor<'a, WordsAndInts, Ignore>), ScanError> {
		Scanner::scan(&cur(s))
	}

	#[test]
	fn test_bool() {
		assert!(scan_a::<bool>("").err().is_some());
		assert!(scan_a::<bool>("true").unwrap().0 == true);
		assert!(scan_a::<bool>("false").unwrap().0 == false);
		assert!(scan_a::<bool>("yes").err().is_some());
		assert!(scan_a::<bool>("no").err().is_some());
		assert!(scan_a::<bool>("on").err().is_some());
		assert!(scan_a::<bool>("off").err().is_some());
		assert!(scan_a::<bool>("1").err().is_some());
		assert!(scan_a::<bool>("0").err().is_some());
	}

	#[test]
	fn test_char() {
		assert!(scan_a::<char>("").err().is_some());
		assert!(scan_a::<char>("x").unwrap().0 == 'x');
		assert!(scan_a::<char>("xy").unwrap().0 == 'x');
		assert!(scan_a::<char>("日本語").unwrap().0 == '日');
	}

	#[test]
	fn test_floats() {
		use std::from_str::FromStr;
		use std::num::{Float, Zero};

		fn test<'a, F: Float + Scanner<'a> + FromStr>() {
			let f = |v:F| v;
			let fs = |s:&str| -> F FromStr::from_str(s).unwrap();

			assert!(scan_a::<F>("").err().is_some());
			assert!(scan_a::<F>("0").ok().unwrap().0 == f(Zero::zero()));
			assert!(scan_a::<F>("0.0").ok().unwrap().0 == f(Zero::zero()));
			assert!(scan_a::<F>("-0").ok().unwrap().0 == -f(Zero::zero()));
			assert!(scan_a::<F>("1.0").ok().unwrap().0 == fs("1.0"));
			assert!(scan_a::<F>("1.00").ok().unwrap().0 == fs("1.0"));
			assert!(scan_a::<F>("1.0e0").ok().unwrap().0 == fs("1.0"));
			assert!(scan_a::<F>("1.0e1").ok().unwrap().0 == fs("10.0"));
		}

		test::<f32>();
		test::<f64>();
	}

	#[test]
	fn test_sized_ints() {
		use std::fmt::Show;
		use std::num::{Bounded, Int, Zero};

		fn test<'a, I: Int + Scanner<'a> + Show>(check_past: bool) {
			let zero: I = Zero::zero();
			let min: I = Bounded::min_value();
			let max: I = Bounded::max_value();

			assert!(scan_a::<I>("").err().is_some());
			assert!(scan_a::<I>("0").ok().unwrap().0 == zero);
			assert!(scan_a::<I>(format!("{}", min).as_slice()).ok().unwrap().0 == min);
			assert!(scan_a::<I>(format!("{}", max).as_slice()).ok().unwrap().0 == max);

			if check_past {
				let past_min: i64 = (min.to_i64()).unwrap() - 1;
				let past_max: u64 = (max.to_u64()).unwrap() + 1;
				assert!(scan_a::<I>(format!("{}", past_min).as_slice()).err().is_some());
				assert!(scan_a::<I>(format!("{}", past_max).as_slice()).err().is_some());
			}
		}

		test::<i8>(true);
		test::<i16>(true);
		test::<i32>(true);
		test::<i64>(false);
		test::<u8>(true);
		test::<u16>(true);
		test::<u32>(true);
		test::<u64>(false);
	}

	#[test]
	fn test_int() {
		assert!(scan_a::<int>("").err().is_some());
		assert!(scan_a::<int>("0").ok().unwrap().0 == 0);
		assert!(scan_a::<int>("42").ok().unwrap().0 == 42);
		assert!(scan_a::<int>("1_234").ok().unwrap().0 == 1);
		assert!(scan_a::<int>("x").err().is_some());
		assert!(scan_a::<int>("0x").ok().unwrap().0 == 0);
		assert!(scan_a::<int>("42x").ok().unwrap().0 == 42);
		assert!(scan_a::<int>("-").err().is_some());
		assert!(scan_a::<int>("-0").ok().unwrap().0 == 0);
		assert!(scan_a::<int>("-42").ok().unwrap().0 == -42);
		assert!(scan_a::<int>("-1_234").ok().unwrap().0 == -1);
	}

	#[test]
	fn test_uint() {
		assert!(scan_a::<uint>("").err().is_some());
		assert!(scan_a::<uint>("0").ok().unwrap().0 == 0);
		assert!(scan_a::<uint>("42").ok().unwrap().0 == 42);
		assert!(scan_a::<uint>("1_234").ok().unwrap().0 == 1);
		assert!(scan_a::<uint>("x").err().is_some());
		assert!(scan_a::<uint>("0x").ok().unwrap().0 == 0);
		assert!(scan_a::<uint>("42x").ok().unwrap().0 == 42);
		assert!(scan_a::<uint>("-").err().is_some());
		assert!(scan_a::<uint>("-0").err().is_some());
		assert!(scan_a::<uint>("-42").err().is_some());
		assert!(scan_a::<uint>("-1_234").err().is_some());
	}

	#[test]
	fn test_str() {
		assert!(scan_a::<&str>("").err().is_some());
		assert!(scan_a::<&str>("a").ok().unwrap().0 == "a");
		assert!(scan_a::<&str>("a b").ok().unwrap().0 == "a");
		assert!(scan_a::<&str>("abc").ok().unwrap().0 == "abc");
		assert!(scan_a::<&str>("ab-c").ok().unwrap().0 == "ab");
	}
}
<|endoftext|>


<commit_before>//! Futures integration.

use crate::stream::Stream;
use crate::sync::Mutex;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{LocalWaker, Poll, Waker};

/// The storage of a stream future.
#[derive(Debug)]
struct StreamFutureStorage<T> {
    value: Option<T>,
    waker: Option<Waker>,
}

impl<T> StreamFutureStorage<T> {
    fn new() -> Arc<Mutex<Self>> {
        Arc::new(Mutex::new(StreamFutureStorage {
            value: None,
            waker: None,
        }))
    }
}

/// A future that waits for a stream value.
///
/// This is created by `Stream::next`.
#[derive(Debug)]
pub struct StreamFuture<T> {
    storage: Arc<Mutex<StreamFutureStorage<T>>>,
    stream: Stream<T>,
}

impl<T> StreamFuture<T> {
    /// Creates a future that returns the next value sent to this stream.
    pub(crate) fn new(stream: Stream<T>) -> Self
    where
        T: Clone + Send + 'static,
    {
        let storage = StreamFutureStorage::new();
        let weak = Arc::downgrade(&storage);
        stream.observe(move |val| {
            if let Some(st) = weak.upgrade() {
                let mut storage = st.lock();
                storage.value = Some(val.into_owned());
                if let Some(waker) = &storage.waker {
                    waker.wake();
                }
            }
            false
        });
        StreamFuture { storage, stream }
    }

    /// Obtains the source stream.
    pub fn get_source(&self) -> &Stream<T> {
        &self.stream
    }
}

impl<T> Future for StreamFuture<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
        let mut storage = self.storage.lock();
        if let Some(value) = storage.value.take() {
            Poll::Ready(value)
        } else {
            storage.waker = Some(lw.clone().into_waker());
            Poll::Pending
        }
    }
}

impl<T> Unpin for StreamFuture<T> {}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::stream::Sink;
    use futures::executor::block_on;

    #[test]
    fn basic() {
        let sink = Sink::new();
        let future = StreamFuture::new(sink.stream());

        sink.send(42);
        sink.send(13);
        assert_eq!(block_on(future), 42);
    }
}
<commit_msg>detect invalid poll on future after the value has been taken<commit_after>//! Futures integration.

use crate::stream::Stream;
use crate::sync::Mutex;
use std::future::Future;
use std::mem;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{LocalWaker, Poll, Waker};

/// The state a stream future.
#[derive(Debug)]
enum FutureValue<T> {
    Pending,
    Ready(T),
    Finished,
}

/// The storage of a stream future.
#[derive(Debug)]
struct StreamFutureStorage<T> {
    value: FutureValue<T>,
    waker: Option<Waker>,
}

impl<T> StreamFutureStorage<T> {
    fn new() -> Arc<Mutex<Self>> {
        Arc::new(Mutex::new(StreamFutureStorage {
            value: FutureValue::Pending,
            waker: None,
        }))
    }
}

/// A future that waits for a stream value.
///
/// This is created by `Stream::next`.
#[derive(Debug)]
pub struct StreamFuture<T> {
    storage: Arc<Mutex<StreamFutureStorage<T>>>,
    stream: Stream<T>,
}

impl<T> StreamFuture<T> {
    /// Creates a future that returns the next value sent to this stream.
    pub(crate) fn new(stream: Stream<T>) -> Self
    where
        T: Clone + Send + 'static,
    {
        let storage = StreamFutureStorage::new();
        let weak = Arc::downgrade(&storage);
        stream.observe(move |val| {
            if let Some(st) = weak.upgrade() {
                let mut storage = st.lock();
                storage.value = FutureValue::Ready(val.into_owned());
                if let Some(waker) = &storage.waker {
                    waker.wake();
                }
            }
            false
        });
        StreamFuture { storage, stream }
    }

    /// Obtains the source stream.
    pub fn get_source(&self) -> &Stream<T> {
        &self.stream
    }
}

impl<T> Future for StreamFuture<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
        let mut storage = self.storage.lock();
        match mem::replace(&mut storage.value, FutureValue::Pending) {
            FutureValue::Ready(value) => {
                storage.value = FutureValue::Finished;
                Poll::Ready(value)
            }
            FutureValue::Pending => {
                storage.waker = Some(lw.clone().into_waker());
                Poll::Pending
            }
            FutureValue::Finished => {
                storage.value = FutureValue::Finished;
                panic!("future polled again after completion");
            }
        }
    }
}

impl<T> Unpin for StreamFuture<T> {}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::stream::Sink;
    use futures::executor::block_on;

    #[test]
    fn basic() {
        let sink = Sink::new();
        let future = StreamFuture::new(sink.stream());

        sink.send(42);
        sink.send(13);
        assert_eq!(block_on(future), 42);
    }

    #[test]
    #[should_panic]
    fn invalid_poll() {
        let sink = Sink::new();
        let mut future = StreamFuture::new(sink.stream());

        sink.send(42);
        let _a = block_on(&mut future);
        let _b = block_on(&mut future);
    }
}
<|endoftext|>


<commit_before>use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    DISABLED,
    NOW,
    CH1_PE,
    CH1_NE,
    CH2_PE,
    CH2_NE,
    EXT_PE,
    EXT_NE,
    AWG_PE,
    AWG_NE,
}

impl ::std::fmt::Display for Source {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Source::DISABLED => "DISABLED",
            &Source::NOW => "NOW",
            &Source::CH1_PE => "CH1_PE",
            &Source::CH1_NE => "CH1_NE",
            &Source::CH2_PE => "CH2_PE",
            &Source::CH2_NE => "CH2_NE",
            &Source::EXT_PE => "EXT_PE",
            &Source::EXT_NE => "EXT_NE",
            &Source::AWG_PE => "AWG_PE",
            &Source::AWG_NE => "AWG_NE",
        };

        write!(f, "{}", display)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum State {
    WAIT,
    TD,
    UNKNOW,
}

impl ::std::convert::From<String> for State {
    fn from(s: String) -> Self {
        match s.as_str() {
            "WAIT" => State::WAIT,
            "TD" => State::TD,
            _ => State::UNKNOW,
        }
    }
}

#[derive(Clone)]
pub struct Trigger {
    socket: ::std::cell::RefCell<Socket>,
}

impl Trigger {
    pub fn new(socket: Socket) -> Self {
        Trigger {
            socket: ::std::cell::RefCell::new(socket),
        }
    }

    /**
     * Trigger immediately or set trigger source & edge.
     */
    pub fn enable(&self, source: Source) {
        self.send(format!("ACQ:TRIG {}", source));
    }

    /**
     * Disable triggering.
     */
    pub fn disable(&self) {
        self.enable(Source::DISABLED);
    }

    /**
     *  Get trigger status.
     *
     *  If DISABLED -> TD else WAIT.
     */
    pub fn get_state(&self) -> State {
        self.send("ACQ:TRIG:STAT?");

        self.receive()
            .into()
    }

    /**
     * Set trigger delay in samples.
     */
    pub fn set_delay(&self, delay: u16) {
        self.send(format!("ACQ:TRIG:DLY {}", delay));
    }

    /**
     * Get trigger delay in samples.
     */
    pub fn get_delay(&self) -> u16 {
        self.send("ACQ:TRIG:DLY?");

        self.receive()
            .parse()
            .unwrap()
    }

    /**
     * Set trigger delay in ns.
     */
    pub fn set_delay_in_ns(&self, delay: u8) {
        self.send(format!("ACQ:TRIG:DLY:NS {}", delay));
    }

    /**
     * Get trigger delay in ns.
     */
    pub fn get_delay_in_ns(&self) -> u8 {
        self.send("ACQ:TRIG:DLY:NS?");

        self.receive()
            .replace("ns", "")
            .parse()
            .unwrap()
    }

    /**
     * Set trigger level in mV.
     */
    pub fn set_level(&self, level: u8) {
        self.send(format!("ACQ:TRIG:LEV {}", level));
    }

    /**
     * Get trigger level in mV.
     */
    pub fn get_level(&self) -> u8 {
        self.send("ACQ:TRIG:LEV?");

        self.receive()
            .replace("mV", "")
            .parse()
            .unwrap()
    }

    fn send<D>(&self, message: D)
        where D: ::std::fmt::Display
    {
        let mut socket = self.socket.borrow_mut();

        socket.send(message);
    }

    fn receive(&self) -> String {
        let mut socket = self.socket.borrow_mut();

        socket.receive()
    }
}

#[cfg(test)]
mod test {
    macro_rules! trigger_assert {
        ($f:ident, $e:expr) => {
            let (rx, trigger) = create_trigger();

            trigger.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_enable() {
        let (rx, trigger) = create_trigger();

        trigger.enable(::trigger::Source::NOW);
        assert_eq!("ACQ:TRIG NOW\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_disable() {
        trigger_assert!(disable, "ACQ:TRIG DISABLED\r\n");
    }

    #[test]
    fn test_get_state() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_state(), ::trigger::State::WAIT);
    }

    #[test]
    fn test_set_delay() {
        let (rx, trigger) = create_trigger();

        trigger.set_delay(0);
        assert_eq!("ACQ:TRIG:DLY 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_delay() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_delay(), 2314);
    }

    #[test]
    fn test_set_delay_in_ns() {
        let (rx, trigger) = create_trigger();

        trigger.set_delay_in_ns(0);
        assert_eq!("ACQ:TRIG:DLY:NS 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_delay_in_ns() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_delay_in_ns(), 128);
    }

    #[test]
    fn test_set_level() {
        let (rx, trigger) = create_trigger();

        trigger.set_level(0);
        assert_eq!("ACQ:TRIG:LEV 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_level() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_level(), 123);
    }

    fn create_trigger() -> (::std::sync::mpsc::Receiver<String>, ::trigger::Trigger) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::trigger::Trigger::new(socket))
    }
}
<commit_msg>[trigger] Fix level type<commit_after>use socket::Socket;

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum Source {
    DISABLED,
    NOW,
    CH1_PE,
    CH1_NE,
    CH2_PE,
    CH2_NE,
    EXT_PE,
    EXT_NE,
    AWG_PE,
    AWG_NE,
}

impl ::std::fmt::Display for Source {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        let display = match self {
            &Source::DISABLED => "DISABLED",
            &Source::NOW => "NOW",
            &Source::CH1_PE => "CH1_PE",
            &Source::CH1_NE => "CH1_NE",
            &Source::CH2_PE => "CH2_PE",
            &Source::CH2_NE => "CH2_NE",
            &Source::EXT_PE => "EXT_PE",
            &Source::EXT_NE => "EXT_NE",
            &Source::AWG_PE => "AWG_PE",
            &Source::AWG_NE => "AWG_NE",
        };

        write!(f, "{}", display)
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum State {
    WAIT,
    TD,
    UNKNOW,
}

impl ::std::convert::From<String> for State {
    fn from(s: String) -> Self {
        match s.as_str() {
            "WAIT" => State::WAIT,
            "TD" => State::TD,
            _ => State::UNKNOW,
        }
    }
}

#[derive(Clone)]
pub struct Trigger {
    socket: ::std::cell::RefCell<Socket>,
}

impl Trigger {
    pub fn new(socket: Socket) -> Self {
        Trigger {
            socket: ::std::cell::RefCell::new(socket),
        }
    }

    /**
     * Trigger immediately or set trigger source & edge.
     */
    pub fn enable(&self, source: Source) {
        self.send(format!("ACQ:TRIG {}", source));
    }

    /**
     * Disable triggering.
     */
    pub fn disable(&self) {
        self.enable(Source::DISABLED);
    }

    /**
     *  Get trigger status.
     *
     *  If DISABLED -> TD else WAIT.
     */
    pub fn get_state(&self) -> State {
        self.send("ACQ:TRIG:STAT?");

        self.receive()
            .into()
    }

    /**
     * Set trigger delay in samples.
     */
    pub fn set_delay(&self, delay: u16) {
        self.send(format!("ACQ:TRIG:DLY {}", delay));
    }

    /**
     * Get trigger delay in samples.
     */
    pub fn get_delay(&self) -> u16 {
        self.send("ACQ:TRIG:DLY?");

        self.receive()
            .parse()
            .unwrap()
    }

    /**
     * Set trigger delay in ns.
     */
    pub fn set_delay_in_ns(&self, delay: u8) {
        self.send(format!("ACQ:TRIG:DLY:NS {}", delay));
    }

    /**
     * Get trigger delay in ns.
     */
    pub fn get_delay_in_ns(&self) -> u8 {
        self.send("ACQ:TRIG:DLY:NS?");

        self.receive()
            .replace("ns", "")
            .parse()
            .unwrap()
    }

    /**
     * Set trigger level in mV.
     */
    pub fn set_level(&self, level: f32) {
        self.send(format!("ACQ:TRIG:LEV {}", level));
    }

    /**
     * Get trigger level in mV.
     */
    pub fn get_level(&self) -> f32 {
        self.send("ACQ:TRIG:LEV?");

        self.receive()
            .replace("mV", "")
            .parse()
            .unwrap()
    }

    fn send<D>(&self, message: D)
        where D: ::std::fmt::Display
    {
        let mut socket = self.socket.borrow_mut();

        socket.send(message);
    }

    fn receive(&self) -> String {
        let mut socket = self.socket.borrow_mut();

        socket.receive()
    }
}

#[cfg(test)]
mod test {
    macro_rules! trigger_assert {
        ($f:ident, $e:expr) => {
            let (rx, trigger) = create_trigger();

            trigger.$f();
            assert_eq!($e, rx.recv().unwrap());
        }
    }

    #[test]
    fn test_enable() {
        let (rx, trigger) = create_trigger();

        trigger.enable(::trigger::Source::NOW);
        assert_eq!("ACQ:TRIG NOW\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_disable() {
        trigger_assert!(disable, "ACQ:TRIG DISABLED\r\n");
    }

    #[test]
    fn test_get_state() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_state(), ::trigger::State::WAIT);
    }

    #[test]
    fn test_set_delay() {
        let (rx, trigger) = create_trigger();

        trigger.set_delay(0);
        assert_eq!("ACQ:TRIG:DLY 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_delay() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_delay(), 2314);
    }

    #[test]
    fn test_set_delay_in_ns() {
        let (rx, trigger) = create_trigger();

        trigger.set_delay_in_ns(0);
        assert_eq!("ACQ:TRIG:DLY:NS 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_delay_in_ns() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_delay_in_ns(), 128);
    }

    #[test]
    fn test_set_level() {
        let (rx, trigger) = create_trigger();

        trigger.set_level(0.0);
        assert_eq!("ACQ:TRIG:LEV 0\r\n", rx.recv().unwrap());
    }

    #[test]
    fn test_get_level() {
        let (_, trigger) = create_trigger();

        assert_eq!(trigger.get_level(), 123.0);
    }

    fn create_trigger() -> (::std::sync::mpsc::Receiver<String>, ::trigger::Trigger) {
        let (addr, rx) = ::test::launch_server();
        let socket = ::socket::Socket::new(addr);

        (rx, ::trigger::Trigger::new(socket))
    }
}
<|endoftext|>


<commit_before>use std::marker::PhantomData;
use std::io::{self, ErrorKind};
use std::error::Error;
use std::net::{UdpSocket, SocketAddr, ToSocketAddrs};

pub struct Transit<T> {
    socket: UdpSocket,
    packet_type: PhantomData<T>,
}

/// Trait implemented by types that can be read from the network.
pub trait FromTransit {
    fn from_transit(&[u8]) -> io::Result<Self> where Self: Sized;
}

impl FromTransit for String {
    fn from_transit(buf: &[u8]) -> io::Result<String> {
        let vec = Vec::from(buf);
        let res = String::from_utf8(vec);
        match res {
            Err(utf8err) => Err(io::Error::new(ErrorKind::InvalidData, utf8err.description())),
            Ok(string) => Ok(string),
        }
    }
}

impl FromTransit for Vec<u8> {
    fn from_transit(buf: &[u8]) -> io::Result<Vec<u8>> {
        Ok(buf.iter().map(|x| *x).collect())
    }
}

/// Trait implemented by types that can be written to the network.
pub trait IntoTransit {
    fn into_transit(&self) -> &[u8];
}

impl IntoTransit for String {
    fn into_transit(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl<'a> IntoTransit for &'a [u8] {
    fn into_transit(&self) -> &[u8] {
        *self
    }
}

impl<T> Transit<T> {
    pub fn new<A>(addr: A) -> io::Result<Transit<T>> where A: ToSocketAddrs {
        let socket = try!(UdpSocket::bind(addr));
        Ok(Transit {
            socket: socket,
            packet_type: PhantomData,
        })
    }

    /// On success, this function returns the type deserialized using the FromTransit trait
    /// implementation. The trait implementation should be able to detect whether or not the buffer
    /// contains a valid UDP message and emit an error appropriately.
    pub fn recv_from(&self) -> io::Result<(T, SocketAddr)> where T: FromTransit {
        let mut buf = [0; 1024];
        let (n, addr) = try!(self.socket.recv_from(&mut buf));
        assert!(n < 1024);
        let data = try!(T::from_transit(&buf[..n]));
        Ok((data, addr))
    }

    /// Transforms the packet into a byte array and sends it to the associated address.
    pub fn send_to<A>(&self, pkt: &T, addr: A) -> io::Result<()> where T: IntoTransit, A: ToSocketAddrs {
        let buf = pkt.into_transit();
        try!(self.socket.send_to(buf, addr));
        Ok(())
    }

    pub fn local_addr(&self)  -> io::Result<SocketAddr> {
        self.socket.local_addr()
    }
}

#[cfg(test)]
mod test {
    use udp::*;
    use std::slice;
    use std::io::{self, Error, ErrorKind};

    #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Debug)]
    struct Test {
        ten: u8,
    }

    impl IntoTransit for Test {
        fn into_transit(&self) -> &[u8] {
            unsafe { slice::from_raw_parts(&self.ten as *const u8, 1) }
        }
    }

    impl FromTransit for Test {
        fn from_transit(buf: &[u8]) -> io::Result<Test> {
            if buf[0] != 10 {
                Err(Error::new(ErrorKind::InvalidData, "failed to serialize"))
            } else {
                Ok(Test { ten: buf[0] })
            }
        }
    }

    #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Debug)]
    struct Another {
        data: u8,
    }

    impl IntoTransit for Another {
        fn into_transit(&self) -> &[u8] {
            unsafe { slice::from_raw_parts(&self.data as *const u8, 1) }
        }
    }

    impl FromTransit for Another {
        fn from_transit(buf: &[u8]) -> io::Result<Another> {
            Ok(Another { data: buf[0] })
        }
    }

    #[test]
    fn test_send_recv() {
        let addr1 = "127.0.0.1:61001";
        let addr2 = "127.0.0.1:61002";
        let transit1: Transit<Test> = Transit::new(addr1).unwrap();
        let transit2 = Transit::new(addr2).unwrap();
        let test = Test { ten: 10 };

        let res = transit2.send_to(&test, addr1);
        assert!(res.is_ok());
        let res = transit1.recv_from();
        assert!(res.is_ok());
        let (data, _addr) = res.unwrap();
        assert_eq!(data, test);
    }

    #[test]
    fn test_send_recv_string() {
        let addr1 = "127.0.0.1:63001";
        let addr2 = "127.0.0.1:63002";
        let transit1: Transit<String> = Transit::new(addr1).unwrap();
        let transit2 = Transit::new(addr2).unwrap();
        let test = String::from("hello");

        let res = transit2.send_to(&test, addr1);
        assert!(res.is_ok());
        let res = transit1.recv_from();
        assert!(res.is_ok());
        let (data, _addr) = res.unwrap();
        assert_eq!(data, test);
    }

    #[test]
    fn test_send_recv_bytes() {
        let addr1 = "127.0.0.1:64001";
        let addr2 = "127.0.0.1:64002";
        let transit1 = Transit::new(addr1).unwrap();
        let transit2 = Transit::new(addr2).unwrap();
        let vec = vec!(9u8);
        let slice = &vec[..];

        let res = transit2.send_to(&slice, addr1);
        assert!(res.is_ok());
        let res: io::Result<(Vec<u8>, _)> = transit1.recv_from();
        assert!(res.is_ok());
        let (data, _addr) = res.unwrap();
        assert_eq!(data, vec);
    }

    #[test]
    fn test_packet_type() {
        let addr1 = "127.0.0.1:62001";
        let addr2 = "127.0.0.1:62002";
        let transit1: Transit<Another> = Transit::new(addr1).unwrap();
        let transit2: Transit<Test> = Transit::new(addr2).unwrap();
        let test = Another { data: 27 };

        let res = transit1.send_to(&test, addr2);
        assert!(res.is_ok());
        let res = transit2.recv_from();
        assert!(res.is_err());
    }
}
<commit_msg>Add documentation to FromTransit<commit_after>use std::marker::PhantomData;
use std::io::{self, ErrorKind};
use std::error::Error;
use std::net::{UdpSocket, SocketAddr, ToSocketAddrs};

pub struct Transit<T> {
    socket: UdpSocket,
    packet_type: PhantomData<T>,
}

/// Trait implemented by types that can be read from the network.
///
/// # Warnings
///
/// A type should always make sure to check whether the passed in buffer is a legitimate packet
/// sent from a known application. This will be called on any UDP packet received on the socket.
pub trait FromTransit {
    fn from_transit(&[u8]) -> io::Result<Self> where Self: Sized;
}

impl FromTransit for String {
    fn from_transit(buf: &[u8]) -> io::Result<String> {
        let vec = Vec::from(buf);
        let res = String::from_utf8(vec);
        match res {
            Err(utf8err) => Err(io::Error::new(ErrorKind::InvalidData, utf8err.description())),
            Ok(string) => Ok(string),
        }
    }
}

impl FromTransit for Vec<u8> {
    fn from_transit(buf: &[u8]) -> io::Result<Vec<u8>> {
        Ok(buf.iter().map(|x| *x).collect())
    }
}

/// Trait implemented by types that can be written to the network.
pub trait IntoTransit {
    fn into_transit(&self) -> &[u8];
}

impl IntoTransit for String {
    fn into_transit(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl<'a> IntoTransit for &'a [u8] {
    fn into_transit(&self) -> &[u8] {
        *self
    }
}

impl<T> Transit<T> {
    pub fn new<A>(addr: A) -> io::Result<Transit<T>> where A: ToSocketAddrs {
        let socket = try!(UdpSocket::bind(addr));
        Ok(Transit {
            socket: socket,
            packet_type: PhantomData,
        })
    }

    /// On success, this function returns the type deserialized using the FromTransit trait
    /// implementation. The trait implementation should be able to detect whether or not the buffer
    /// contains a valid UDP message and emit an error appropriately.
    pub fn recv_from(&self) -> io::Result<(T, SocketAddr)> where T: FromTransit {
        let mut buf = [0; 1024];
        let (n, addr) = try!(self.socket.recv_from(&mut buf));
        assert!(n < 1024);
        let data = try!(T::from_transit(&buf[..n]));
        Ok((data, addr))
    }

    /// Transforms the packet into a byte array and sends it to the associated address.
    pub fn send_to<A>(&self, pkt: &T, addr: A) -> io::Result<()> where T: IntoTransit, A: ToSocketAddrs {
        let buf = pkt.into_transit();
        try!(self.socket.send_to(buf, addr));
        Ok(())
    }

    pub fn local_addr(&self)  -> io::Result<SocketAddr> {
        self.socket.local_addr()
    }
}

#[cfg(test)]
mod test {
    use udp::*;
    use std::slice;
    use std::io::{self, Error, ErrorKind};

    #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Debug)]
    struct Test {
        ten: u8,
    }

    impl IntoTransit for Test {
        fn into_transit(&self) -> &[u8] {
            unsafe { slice::from_raw_parts(&self.ten as *const u8, 1) }
        }
    }

    impl FromTransit for Test {
        fn from_transit(buf: &[u8]) -> io::Result<Test> {
            if buf[0] != 10 {
                Err(Error::new(ErrorKind::InvalidData, "failed to serialize"))
            } else {
                Ok(Test { ten: buf[0] })
            }
        }
    }

    #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Debug)]
    struct Another {
        data: u8,
    }

    impl IntoTransit for Another {
        fn into_transit(&self) -> &[u8] {
            unsafe { slice::from_raw_parts(&self.data as *const u8, 1) }
        }
    }

    impl FromTransit for Another {
        fn from_transit(buf: &[u8]) -> io::Result<Another> {
            Ok(Another { data: buf[0] })
        }
    }

    #[test]
    fn test_send_recv() {
        let addr1 = "127.0.0.1:61001";
        let addr2 = "127.0.0.1:61002";
        let transit1: Transit<Test> = Transit::new(addr1).unwrap();
        let transit2 = Transit::new(addr2).unwrap();
        let test = Test { ten: 10 };

        let res = transit2.send_to(&test, addr1);
        assert!(res.is_ok());
        let res = transit1.recv_from();
        assert!(res.is_ok());
        let (data, _addr) = res.unwrap();
        assert_eq!(data, test);
    }

    #[test]
    fn test_send_recv_string() {
        let addr1 = "127.0.0.1:63001";
        let addr2 = "127.0.0.1:63002";
        let transit1: Transit<String> = Transit::new(addr1).unwrap();
        let transit2 = Transit::new(addr2).unwrap();
        let test = String::from("hello");

        let res = transit2.send_to(&test, addr1);
        assert!(res.is_ok());
        let res = transit1.recv_from();
        assert!(res.is_ok());
        let (data, _addr) = res.unwrap();
        assert_eq!(data, test);
    }

    #[test]
    fn test_send_recv_bytes() {
        let addr1 = "127.0.0.1:64001";
        let addr2 = "127.0.0.1:64002";
        let transit1 = Transit::new(addr1).unwrap();
        let transit2 = Transit::new(addr2).unwrap();
        let vec = vec!(9u8);
        let slice = &vec[..];

        let res = transit2.send_to(&slice, addr1);
        assert!(res.is_ok());
        let res: io::Result<(Vec<u8>, _)> = transit1.recv_from();
        assert!(res.is_ok());
        let (data, _addr) = res.unwrap();
        assert_eq!(data, vec);
    }

    #[test]
    fn test_packet_type() {
        let addr1 = "127.0.0.1:62001";
        let addr2 = "127.0.0.1:62002";
        let transit1: Transit<Another> = Transit::new(addr1).unwrap();
        let transit2: Transit<Test> = Transit::new(addr2).unwrap();
        let test = Another { data: 27 };

        let res = transit1.send_to(&test, addr2);
        assert!(res.is_ok());
        let res = transit2.recv_from();
        assert!(res.is_err());
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::io;
use std::ops::Drop;
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicUsize, Ordering};

use futures::{future, Future, BoxFuture, Sink};
use futures::sync::mpsc::Sender;
use tokio_core::reactor::Remote;
use tokio_service::Service;

use bytes::{BufMut, Bytes, BytesMut};

use admin;
use codec::Packet;
use job::Job;
use packet::PacketMagic;
use queues::{HandleJobStorage, JobQueuePriority, SharedJobStorage};
use worker::{SharedWorkers, Worker, Wake};
use constants::*;

fn new_res(ptype: u32, data: Bytes) -> Packet {
    Packet {
        magic: PacketMagic::RES,
        ptype: ptype,
        psize: data.len() as u32,
        data: data,
    }
}

fn new_noop() -> Packet {
    new_res(NOOP, Bytes::new())
}

type JobWaiters = Arc<Mutex<HashMap<Bytes, Vec<usize>>>>;
type SendersByConnId = Arc<Mutex<HashMap<usize, Sender<Packet>>>>;

pub struct GearmanService {
    pub conn_id: usize,
    pub queues: SharedJobStorage,
    pub workers: SharedWorkers,
    pub worker: Arc<Mutex<Worker>>,
    pub job_count: Arc<AtomicUsize>,
    senders_by_conn_id: SendersByConnId,
    job_waiters: JobWaiters,
    remote: Remote,
}

fn next_field(buf: &mut Bytes) -> Result<Bytes, io::Error> {
    match buf[..].iter().position(|b| *b == b'\0') {
        Some(null_pos) => {
            let value = buf.split_to(null_pos);
            buf.split_to(1);
            Ok(value)
        }
        None => Err(io::Error::new(io::ErrorKind::Other, "Can't find null")),
    }
}

impl Drop for GearmanService {
    fn drop(&mut self) {
        trace!("Dropping conn_id = {}", self.conn_id);
        self.workers.shutdown(self.conn_id);
        debug!("Dropped conn_id = {}", self.conn_id);
    }
}

impl GearmanService {
    /// Things that don't require a body should use this
    fn response_from_packet(&self, packet: &Packet) -> Packet {
        match packet.ptype {
            ADMIN_VERSION => {
                let resp_str = "OK some-rustygear-version\n";
                let mut resp_body = BytesMut::with_capacity(resp_str.len());
                resp_body.put(&resp_str[..]);
                Packet {
                    magic: PacketMagic::TEXT,
                    ptype: packet.ptype,
                    psize: resp_str.len() as u32,
                    data: resp_body.freeze(),
                }
            }
            ADMIN_STATUS => admin::admin_command_status(self.queues.clone(), self.workers.clone()),
            _ => {
                panic!("response_from_packet called with invalid ptype: {}",
                       packet.ptype)
            }
        }
    }

    fn send_to_conn_id(&self, conn_id: usize, packet: Packet) {
        let senders_by_conn_id = self.senders_by_conn_id.lock().unwrap();
        match senders_by_conn_id.get(&conn_id) {
            None => {
                panic!("No connection found for conn_id = {}", &conn_id); // XXX You can do better
            }
            Some(tx) => {
                let tx = tx.clone();
                self.remote.spawn(move |handle| {
                    handle.spawn(tx.send(packet).then(|res| {
                        match res {
                            Ok(_) => {}
                            Err(e) => error!("Send Error! {:?}", e),
                        }
                        Ok(())
                    }));
                    Ok(())
                });
            }
        }
    }

    pub fn new(conn_id: usize,
               queues: SharedJobStorage,
               workers: SharedWorkers,
               job_count: Arc<AtomicUsize>,
               senders_by_conn_id: SendersByConnId,
               job_waiters: JobWaiters,
               remote: Remote)
               -> GearmanService {
        GearmanService {
            conn_id: conn_id,
            queues: queues,
            worker: Arc::new(Mutex::new((Worker::new()))),
            workers: workers,
            job_count: job_count,
            senders_by_conn_id: senders_by_conn_id,
            job_waiters: job_waiters,
            remote: remote,
        }
    }

    fn no_response() -> Packet {
        Packet {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: 0,
            data: Bytes::new(),
        }
    }

    fn handle_can_do(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        let workers = self.workers.clone();
        let conn_id = self.conn_id;
        debug!("CAN_DO fname = {:?}", packet.data);
        let mut worker = worker.lock().unwrap();
        worker.can_do(packet.data.clone());
        workers.clone().wakeup(&mut worker, conn_id);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_cant_do(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        debug!("CANT_DO fname = {:?}", packet.data);
        let mut worker = worker.lock().unwrap();
        worker.cant_do(&packet.data);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_grab_job_all(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(4 + j.handle.len() + j.fname.len() +
                                                       j.unique.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.unique);
                data.put_u8(b'\0');
                // reducer not implemented
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN_ALL, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_grab_job_uniq(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(3 + j.handle.len() + j.fname.len() +
                                                       j.unique.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.unique);
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN_UNIQ, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_grab_job(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(2 + j.handle.len() + j.fname.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_pre_sleep(&self) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        let ref mut w = worker.lock().unwrap();
        self.workers.clone().sleep(w, self.conn_id);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_submit_job(&self,
                         priority: JobQueuePriority,
                         wait: bool,
                         packet: Packet)
                         -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let conn_id = match wait {
            true => Some(self.conn_id),
            false => None,
        };
        let mut workers = self.workers.clone();
        let job_count = self.job_count.clone();
        let remote = self.remote.clone();
        let senders_by_conn_id = self.senders_by_conn_id.clone();
        let mut fields = packet.data.clone();
        trace!("fields = {:?}", fields);
        let fname = next_field(&mut fields).unwrap();
        let unique = next_field(&mut fields).unwrap();
        trace!("  --> fname = {:?} unique = {:?}", fname, unique);
        let mut add = false;
        let handle = match queues.coalesce_unique(&unique, conn_id) {
            Some(handle) => handle,
            None => {
                {
                    for wake in workers.queue_wake(&fname) {
                        let senders_by_conn_id = senders_by_conn_id.lock().unwrap();
                        match senders_by_conn_id.get(&wake) {
                            None => {
                                debug!("No connection found to wake up for conn_id = {}", wake);
                            }
                            Some(tx) => {
                                let tx = tx.clone();
                                remote.spawn(move |handle| {
                                    handle.spawn(tx.send(new_noop()).then(|res| {
                                        match res {
                                            Ok(_) => {}
                                            Err(e) => error!("Send Error! {:?}", e),
                                        }
                                        Ok(())
                                    }));
                                    Ok(())
                                });
                            }
                        }
                    }
                }
                // H:091234567890
                let mut handle = BytesMut::with_capacity(12);
                let job_num = job_count.fetch_add(1, Ordering::Relaxed);
                debug!("job_num = {}", job_num);
                handle.extend(format!("H:{:010}", job_num).as_bytes());
                add = true;
                handle.freeze()
            }
        };
        if add {
            let job = Arc::new(Job::new(fname, unique, fields, handle.clone()));
            info!("Created job {:?}", job);
            queues.add_job(job.clone(), priority, conn_id);
        }
        // If we don't store any senders, the sender will be dropped and the rx
        // stream should end thus releasing the waiter immediately.
        let psize = handle.len() as u32;
        if wait {
            // Fetch our sender
            let mut job_waiters = self.job_waiters.lock().unwrap();
            let mut waiters = job_waiters.entry(handle.clone()).or_insert(Vec::new());
            waiters.push(self.conn_id);
        }
        future::finished(Packet {
                magic: PacketMagic::RES,
                ptype: JOB_CREATED,
                psize: psize,
                data: handle,
            })
            .boxed()
    }

    fn handle_work_complete(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        // Search for handle
        let mut fields = packet.data.clone();
        let handle = next_field(&mut fields).unwrap();
        let worker = self.worker.clone();
        let queues = self.queues.clone();
        info!("Job is complete {:?}", handle);
        let mut worker = worker.lock().unwrap();
        match worker.get_assigned_job(&handle) {
            Some(ref mut j) => {
                let mut queues = queues.lock().unwrap();
                queues.remove_job(&j.unique);
            }
            None => {
                error!("WORK_COMPLETE received but no active jobs");
            }
        }
        worker.unassign_job(&handle);
        let mut job_waiters = self.job_waiters.lock().unwrap();
        // If there are waiters, send the packet to them
        if let Some(waiters) = job_waiters.remove(&handle) {
            for conn_id in waiters.iter() {
                self.send_to_conn_id(*conn_id, packet.clone());
            }
        }
        future::finished(Self::no_response()).boxed()
    }

    fn handle_work_update(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let mut fields = packet.data.clone();
        let handle = next_field(&mut fields).unwrap();
        let job_waiters = self.job_waiters.lock().unwrap();
        if let Some(waiters) = job_waiters.get(&handle) {
            for conn_id in waiters.iter() {
                self.send_to_conn_id(*conn_id, packet.clone());
            }
        }
        future::finished(Self::no_response()).boxed()
    }
}

impl Service for GearmanService {
    type Request = Packet;
    type Response = Packet;
    type Error = io::Error;
    type Future = BoxFuture<Self::Response, Self::Error>;

    fn call(&self, req: Self::Request) -> Self::Future {
        debug!("Got a req {:?}", req);
        match req.ptype {
            ADMIN_VERSION | ADMIN_STATUS => future::ok(self.response_from_packet(&req)).boxed(),
            SUBMIT_JOB => self.handle_submit_job(PRIORITY_NORMAL, true, req),
            SUBMIT_JOB_HIGH => self.handle_submit_job(PRIORITY_HIGH, true, req),
            SUBMIT_JOB_LOW => self.handle_submit_job(PRIORITY_LOW, true, req),
            SUBMIT_JOB_BG => self.handle_submit_job(PRIORITY_NORMAL, false, req),
            SUBMIT_JOB_HIGH_BG => self.handle_submit_job(PRIORITY_HIGH, false, req),
            SUBMIT_JOB_LOW_BG => self.handle_submit_job(PRIORITY_LOW, false, req),
            PRE_SLEEP => self.handle_pre_sleep(),
            CAN_DO => self.handle_can_do(&req),
            CANT_DO => self.handle_cant_do(&req),
            GRAB_JOB => self.handle_grab_job(),
            GRAB_JOB_UNIQ => self.handle_grab_job_uniq(),
            GRAB_JOB_ALL => self.handle_grab_job_all(),
            WORK_COMPLETE => self.handle_work_complete(&req),
            WORK_STATUS | WORK_DATA | WORK_WARNING => self.handle_work_update(&req),/*
            ECHO_REQ => self.handle_echo_req(&req),*/
            _ => {
                error!("Unimplemented: {:?} processing packet", req);
                future::err(io::Error::new(io::ErrorKind::Other,
                                           format!("Invalid packet type {}", req.ptype)))
                    .boxed()
            }
        }
    }
}
<commit_msg>make ECHO work<commit_after>use std::collections::HashMap;
use std::io;
use std::ops::Drop;
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicUsize, Ordering};

use futures::{future, Future, BoxFuture, Sink};
use futures::sync::mpsc::Sender;
use tokio_core::reactor::Remote;
use tokio_service::Service;

use bytes::{BufMut, Bytes, BytesMut};

use admin;
use codec::Packet;
use job::Job;
use packet::PacketMagic;
use queues::{HandleJobStorage, JobQueuePriority, SharedJobStorage};
use worker::{SharedWorkers, Worker, Wake};
use constants::*;

fn new_res(ptype: u32, data: Bytes) -> Packet {
    Packet {
        magic: PacketMagic::RES,
        ptype: ptype,
        psize: data.len() as u32,
        data: data,
    }
}

fn new_noop() -> Packet {
    new_res(NOOP, Bytes::new())
}

type JobWaiters = Arc<Mutex<HashMap<Bytes, Vec<usize>>>>;
type SendersByConnId = Arc<Mutex<HashMap<usize, Sender<Packet>>>>;

pub struct GearmanService {
    pub conn_id: usize,
    pub queues: SharedJobStorage,
    pub workers: SharedWorkers,
    pub worker: Arc<Mutex<Worker>>,
    pub job_count: Arc<AtomicUsize>,
    senders_by_conn_id: SendersByConnId,
    job_waiters: JobWaiters,
    remote: Remote,
}

fn next_field(buf: &mut Bytes) -> Result<Bytes, io::Error> {
    match buf[..].iter().position(|b| *b == b'\0') {
        Some(null_pos) => {
            let value = buf.split_to(null_pos);
            buf.split_to(1);
            Ok(value)
        }
        None => Err(io::Error::new(io::ErrorKind::Other, "Can't find null")),
    }
}

impl Drop for GearmanService {
    fn drop(&mut self) {
        trace!("Dropping conn_id = {}", self.conn_id);
        self.workers.shutdown(self.conn_id);
        debug!("Dropped conn_id = {}", self.conn_id);
    }
}

impl GearmanService {
    /// Things that don't require a body should use this
    fn response_from_packet(&self, packet: &Packet) -> Packet {
        match packet.ptype {
            ADMIN_VERSION => {
                let resp_str = "OK some-rustygear-version\n";
                let mut resp_body = BytesMut::with_capacity(resp_str.len());
                resp_body.put(&resp_str[..]);
                Packet {
                    magic: PacketMagic::TEXT,
                    ptype: packet.ptype,
                    psize: resp_str.len() as u32,
                    data: resp_body.freeze(),
                }
            }
            ADMIN_STATUS => admin::admin_command_status(self.queues.clone(), self.workers.clone()),
            _ => {
                panic!("response_from_packet called with invalid ptype: {}",
                       packet.ptype)
            }
        }
    }

    fn send_to_conn_id(&self, conn_id: usize, packet: Packet) {
        let senders_by_conn_id = self.senders_by_conn_id.lock().unwrap();
        match senders_by_conn_id.get(&conn_id) {
            None => {
                panic!("No connection found for conn_id = {}", &conn_id); // XXX You can do better
            }
            Some(tx) => {
                let tx = tx.clone();
                self.remote.spawn(move |handle| {
                    handle.spawn(tx.send(packet).then(|res| {
                        match res {
                            Ok(_) => {}
                            Err(e) => error!("Send Error! {:?}", e),
                        }
                        Ok(())
                    }));
                    Ok(())
                });
            }
        }
    }

    pub fn new(conn_id: usize,
               queues: SharedJobStorage,
               workers: SharedWorkers,
               job_count: Arc<AtomicUsize>,
               senders_by_conn_id: SendersByConnId,
               job_waiters: JobWaiters,
               remote: Remote)
               -> GearmanService {
        GearmanService {
            conn_id: conn_id,
            queues: queues,
            worker: Arc::new(Mutex::new((Worker::new()))),
            workers: workers,
            job_count: job_count,
            senders_by_conn_id: senders_by_conn_id,
            job_waiters: job_waiters,
            remote: remote,
        }
    }

    fn no_response() -> Packet {
        Packet {
            magic: PacketMagic::TEXT,
            ptype: ADMIN_RESPONSE,
            psize: 0,
            data: Bytes::new(),
        }
    }

    fn handle_can_do(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        let workers = self.workers.clone();
        let conn_id = self.conn_id;
        debug!("CAN_DO fname = {:?}", packet.data);
        let mut worker = worker.lock().unwrap();
        worker.can_do(packet.data.clone());
        workers.clone().wakeup(&mut worker, conn_id);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_cant_do(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        debug!("CANT_DO fname = {:?}", packet.data);
        let mut worker = worker.lock().unwrap();
        worker.cant_do(&packet.data);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_grab_job_all(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(4 + j.handle.len() + j.fname.len() +
                                                       j.unique.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.unique);
                data.put_u8(b'\0');
                // reducer not implemented
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN_ALL, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_grab_job_uniq(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(3 + j.handle.len() + j.fname.len() +
                                                       j.unique.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.unique);
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN_UNIQ, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_grab_job(&self) -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let worker = self.worker.clone();
        let mut worker = worker.lock().unwrap();
        let ref mut worker = worker;
        match queues.get_job(worker) {
            Some(ref j) => {
                let mut data = BytesMut::with_capacity(2 + j.handle.len() + j.fname.len() +
                                                       j.data.len());
                data.extend(&j.handle);
                data.put_u8(b'\0');
                data.extend(&j.fname);
                data.put_u8(b'\0');
                data.extend(&j.data);
                return future::finished(new_res(JOB_ASSIGN, data.freeze())).boxed();
            }
            None => {}
        };
        future::finished(new_res(NO_JOB, Bytes::new())).boxed()
    }

    fn handle_pre_sleep(&self) -> BoxFuture<Packet, io::Error> {
        let worker = self.worker.clone();
        let ref mut w = worker.lock().unwrap();
        self.workers.clone().sleep(w, self.conn_id);
        future::finished(Self::no_response()).boxed()
    }

    fn handle_submit_job(&self,
                         priority: JobQueuePriority,
                         wait: bool,
                         packet: Packet)
                         -> BoxFuture<Packet, io::Error> {
        let mut queues = self.queues.clone();
        let conn_id = match wait {
            true => Some(self.conn_id),
            false => None,
        };
        let mut workers = self.workers.clone();
        let job_count = self.job_count.clone();
        let remote = self.remote.clone();
        let senders_by_conn_id = self.senders_by_conn_id.clone();
        let mut fields = packet.data.clone();
        trace!("fields = {:?}", fields);
        let fname = next_field(&mut fields).unwrap();
        let unique = next_field(&mut fields).unwrap();
        trace!("  --> fname = {:?} unique = {:?}", fname, unique);
        let mut add = false;
        let handle = match queues.coalesce_unique(&unique, conn_id) {
            Some(handle) => handle,
            None => {
                {
                    for wake in workers.queue_wake(&fname) {
                        let senders_by_conn_id = senders_by_conn_id.lock().unwrap();
                        match senders_by_conn_id.get(&wake) {
                            None => {
                                debug!("No connection found to wake up for conn_id = {}", wake);
                            }
                            Some(tx) => {
                                let tx = tx.clone();
                                remote.spawn(move |handle| {
                                    handle.spawn(tx.send(new_noop()).then(|res| {
                                        match res {
                                            Ok(_) => {}
                                            Err(e) => error!("Send Error! {:?}", e),
                                        }
                                        Ok(())
                                    }));
                                    Ok(())
                                });
                            }
                        }
                    }
                }
                // H:091234567890
                let mut handle = BytesMut::with_capacity(12);
                let job_num = job_count.fetch_add(1, Ordering::Relaxed);
                debug!("job_num = {}", job_num);
                handle.extend(format!("H:{:010}", job_num).as_bytes());
                add = true;
                handle.freeze()
            }
        };
        if add {
            let job = Arc::new(Job::new(fname, unique, fields, handle.clone()));
            info!("Created job {:?}", job);
            queues.add_job(job.clone(), priority, conn_id);
        }
        // If we don't store any senders, the sender will be dropped and the rx
        // stream should end thus releasing the waiter immediately.
        let psize = handle.len() as u32;
        if wait {
            // Fetch our sender
            let mut job_waiters = self.job_waiters.lock().unwrap();
            let mut waiters = job_waiters.entry(handle.clone()).or_insert(Vec::new());
            waiters.push(self.conn_id);
        }
        future::finished(Packet {
                magic: PacketMagic::RES,
                ptype: JOB_CREATED,
                psize: psize,
                data: handle,
            })
            .boxed()
    }

    fn handle_work_complete(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        // Search for handle
        let mut fields = packet.data.clone();
        let handle = next_field(&mut fields).unwrap();
        let worker = self.worker.clone();
        let queues = self.queues.clone();
        info!("Job is complete {:?}", handle);
        let mut worker = worker.lock().unwrap();
        match worker.get_assigned_job(&handle) {
            Some(ref mut j) => {
                let mut queues = queues.lock().unwrap();
                queues.remove_job(&j.unique);
            }
            None => {
                error!("WORK_COMPLETE received but no active jobs");
            }
        }
        worker.unassign_job(&handle);
        let mut job_waiters = self.job_waiters.lock().unwrap();
        // If there are waiters, send the packet to them
        if let Some(waiters) = job_waiters.remove(&handle) {
            for conn_id in waiters.iter() {
                self.send_to_conn_id(*conn_id, packet.clone());
            }
        }
        future::finished(Self::no_response()).boxed()
    }

    fn handle_work_update(&self, packet: &Packet) -> BoxFuture<Packet, io::Error> {
        let mut fields = packet.data.clone();
        let handle = next_field(&mut fields).unwrap();
        let job_waiters = self.job_waiters.lock().unwrap();
        if let Some(waiters) = job_waiters.get(&handle) {
            for conn_id in waiters.iter() {
                self.send_to_conn_id(*conn_id, packet.clone());
            }
        }
        future::finished(Self::no_response()).boxed()
    }
}

impl Service for GearmanService {
    type Request = Packet;
    type Response = Packet;
    type Error = io::Error;
    type Future = BoxFuture<Self::Response, Self::Error>;

    fn call(&self, req: Self::Request) -> Self::Future {
        debug!("Got a req {:?}", req);
        match req.ptype {
            ADMIN_VERSION | ADMIN_STATUS => future::ok(self.response_from_packet(&req)).boxed(),
            SUBMIT_JOB => self.handle_submit_job(PRIORITY_NORMAL, true, req),
            SUBMIT_JOB_HIGH => self.handle_submit_job(PRIORITY_HIGH, true, req),
            SUBMIT_JOB_LOW => self.handle_submit_job(PRIORITY_LOW, true, req),
            SUBMIT_JOB_BG => self.handle_submit_job(PRIORITY_NORMAL, false, req),
            SUBMIT_JOB_HIGH_BG => self.handle_submit_job(PRIORITY_HIGH, false, req),
            SUBMIT_JOB_LOW_BG => self.handle_submit_job(PRIORITY_LOW, false, req),
            PRE_SLEEP => self.handle_pre_sleep(),
            CAN_DO => self.handle_can_do(&req),
            CANT_DO => self.handle_cant_do(&req),
            GRAB_JOB => self.handle_grab_job(),
            GRAB_JOB_UNIQ => self.handle_grab_job_uniq(),
            GRAB_JOB_ALL => self.handle_grab_job_all(),
            WORK_COMPLETE => self.handle_work_complete(&req),
            WORK_STATUS | WORK_DATA | WORK_WARNING => self.handle_work_update(&req),
            ECHO_REQ => future::finished(new_res(ECHO_RES, req.data)).boxed(),
            _ => {
                error!("Unimplemented: {:?} processing packet", req);
                future::err(io::Error::new(io::ErrorKind::Other,
                                           format!("Invalid packet type {}", req.ptype)))
                    .boxed()
            }
        }
    }
}
<|endoftext|>


<commit_before>use pest::prelude::*;

impl_rdp! {
    grammar! {
        whitespace = _{ [" "]|["\t"]|["\n"]|["\r"] }

        raw_text = @{ ( !["{{"] ~ any )+ }
        raw_block_text = @{ ( !["{{{{"] ~ any )+ }

// Note: this is not full and strict json literal definition, just for tokenize string,
// array and object types which may contains whitespace. We will use a real json parser
// for real json processing
        literal = _{ string_literal |
                     array_literal |
                     object_literal |
                     number_literal |
                     null_literal |
                     boolean_literal }

        null_literal = { ["null"] }
        boolean_literal = { ["true"]|["false"] }
        number_literal = @{ ["-"]? ~ ['0'..'9']+ ~ ["."]? ~ ['0'..'9']* ~ (["E"] ~ ["-"]? ~ ['0'..'9']+)? }
        string_literal = @{ ["\""] ~ (!["\""] ~ (["\\\""] | any))* ~ ["\""] }
        array_literal = { ["["] ~ literal? ~ ([","] ~ literal)* ~ ["]"] }
        object_literal = { ["{"] ~ (string_literal ~ [":"] ~ literal)? ~ ([","] ~ string_literal ~ [":"] ~ literal)* ~ ["}"] }

// FIXME: a[0], a["b]
        symbol_char = _{ ['a'..'z']|['A'..'Z']|['0'..'9']|["_"]|["."]|["@"]|["$"] }
        path_char = _{ ["/"] }

        identifier = @{ symbol_char ~ ( symbol_char | path_char )* }
        reference = @{ identifier ~ (["["] ~ (string_literal|['0'..'9']+) ~ ["]"])* ~ reference* }
        name = { subexpression | identifier }

        param = { literal | reference | subexpression }
        hash = @{ identifier ~ ["="] ~ param }
        exp_line = _{ identifier ~ (hash|param)* }

        subexpression = { ["("] ~ name ~ (hash|param)* ~ [")"] }

        pre_whitespace_omitter = { ["~"] }
        pro_whitespace_omitter = { ["~"] }

        expression = { ["{{"] ~ pre_whitespace_omitter? ~ name ~
                                pro_whitespace_omitter? ~ ["}}"] }

        html_expression = { ["{{{"] ~ pre_whitespace_omitter? ~ name ~
                                      pro_whitespace_omitter? ~ ["}}}"] }

        helper_expression = { ["{{"] ~ pre_whitespace_omitter? ~ exp_line ~
                                       pro_whitespace_omitter? ~ ["}}"] }

        invert_tag = { ["{{else}}"]|["{{^}}"] }
        helper_block_start = { ["{{"] ~ pre_whitespace_omitter? ~ ["#"] ~ exp_line ~
                                        pro_whitespace_omitter? ~ ["}}"] }
        helper_block_end = { ["{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                      pro_whitespace_omitter? ~ ["}}"] }
        helper_block = { helper_block_start ~ template ~
                         (invert_tag ~ template)? ~
                         helper_block_end }

        raw_block_start = { ["{{{{"] ~ pre_whitespace_omitter? ~ ["#"] ~ exp_line ~
                                       pro_whitespace_omitter? ~ ["}}}}"] }
        raw_block_end = { ["{{{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                     pro_whitespace_omitter? ~ ["}}}}"] }
        raw_block = { raw_block_start ~ raw_block_text ~ raw_block_end }

        comment = { ["{{!"] ~ (!["}}"] ~ any)* ~ ["}}"] }

        template = { (
            raw_text |
            expression |
            html_expression |
            helper_expression |
            helper_block |
            raw_block |
            comment )*
        }

        handlebars = _ { template ~ eoi }
    }
}

#[test]
fn test_raw_text() {
    let mut rdp = Rdp::new(StringInput::new("<h1> helloworld </h1>    "));
    assert!(rdp.raw_text());
    assert!(rdp.end());
}

#[test]
fn test_raw_block_text() {
    let mut rdp = Rdp::new(StringInput::new("<h1> {{hello}} </h1>"));
    assert!(rdp.raw_block_text());
    assert!(rdp.end());
}

#[test]
fn test_reference() {
    let s = vec!["a",
                 "abc",
                 "../a",
                 "a.b",
                 "@abc",
                 "a[\"abc\"]",
                 "aBc[\"abc\"]",
                 "abc[0][\"nice\"]"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.reference());
        assert!(rdp.end());
    }
}

#[test]
fn test_name() {
    let s = vec!["if", "(abc)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.name());
        assert!(rdp.end());
    }
}

#[test]
fn test_param() {
    let s = vec!["hello", "\"json literal\""];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.param());
        assert!(rdp.end());
    }
}

#[test]
fn test_hash() {
    let s = vec!["hello=world", "hello=\"world\"", "hello=(world)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.hash());
        assert!(rdp.end());
    }
}

#[test]
fn test_json_literal() {
    let s = vec!["\"json string\"",
                 "\"quot: \\\"\"",
                 "[]",
                 "[\"hello\"]",
                 "[1,2,3,4,true]",
                 "{\"hello\": \"world\"}",
                 "{}",
                 "{\"a\":1, \"b\":2 }"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.literal());
        assert!(rdp.end());
    }
}

#[test]
fn test_comment() {
    let s = vec!["{{! hello }}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.comment());
        assert!(rdp.end());
    }
}

#[test]
fn test_expression() {
    let s = vec!["{{exp}}", "{{(exp)}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_expression() {
    let s = vec!["{{exp 1}}",
                 "{{exp \"literal\"}}",
                 "{{exp ref}}",
                 "{{exp (sub)}}",
                 "{{exp (sub 123)}}",
                 "{{exp []}}",
                 "{{exp {}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_html_expression() {
    let s = vec!["{{{html}}}", "{{{(html)}}}", "{{{(html)}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.html_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_start() {
    let s = vec!["{{#if hello}}",
                 "{{#if (hello)}}",
                 "{{#if hello=world}}",
                 "{{#if hello hello=world}}",
                 "{{#if []}}",
                 "{{#if {}}}",
                 "{{#if}}",
                 "{{~#if hello~}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block_start());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_end() {
    let s = vec!["{{/if}}", "{{~/if}}", "{{~/if ~}}", "{{/if   ~}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block_end());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_block() {
    let s = vec!["{{#if hello}}hello{{/if}}",
                 "{{#if true}}hello{{/if}}",
                 "{{#if nice ok=1}}hello{{/if}}",
                 "{{#if}}hello{{else}}world{{/if}}",
                 "{{#if}}hello{{^}}world{{/if}}",
                 "{{#if}}{{#if}}hello{{/if}}{{/if}}",
                 "{{#if}}{{/if}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block());
        assert!(rdp.end());
    }
}

#[test]
fn test_raw_block() {
    let s = vec!["{{{{#if hello}}}}good {{hello}}{{{{/if}}}}",
                 "{{{{#if hello}}}}{{#if nice}}{{/if}}{{{{/if}}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.raw_block());
        assert!(rdp.end());
    }
}
<commit_msg>(fix) invert tag grammar<commit_after>use pest::prelude::*;

impl_rdp! {
    grammar! {
        whitespace = _{ [" "]|["\t"]|["\n"]|["\r"] }

        raw_text = @{ ( !["{{"] ~ any )+ }
        raw_block_text = @{ ( !["{{{{"] ~ any )+ }

// Note: this is not full and strict json literal definition, just for tokenize string,
// array and object types which may contains whitespace. We will use a real json parser
// for real json processing
        literal = _{ string_literal |
                     array_literal |
                     object_literal |
                     number_literal |
                     null_literal |
                     boolean_literal }

        null_literal = { ["null"] }
        boolean_literal = { ["true"]|["false"] }
        number_literal = @{ ["-"]? ~ ['0'..'9']+ ~ ["."]? ~ ['0'..'9']* ~ (["E"] ~ ["-"]? ~ ['0'..'9']+)? }
        string_literal = @{ ["\""] ~ (!["\""] ~ (["\\\""] | any))* ~ ["\""] }
        array_literal = { ["["] ~ literal? ~ ([","] ~ literal)* ~ ["]"] }
        object_literal = { ["{"] ~ (string_literal ~ [":"] ~ literal)? ~ ([","] ~ string_literal ~ [":"] ~ literal)* ~ ["}"] }

// FIXME: a[0], a["b]
        symbol_char = _{ ['a'..'z']|['A'..'Z']|['0'..'9']|["_"]|["."]|["@"]|["$"] }
        path_char = _{ ["/"] }

        identifier = @{ symbol_char ~ ( symbol_char | path_char )* }
        reference = @{ identifier ~ (["["] ~ (string_literal|['0'..'9']+) ~ ["]"])* ~ reference* }
        name = _{ subexpression | identifier }

        param = { literal | reference | subexpression }
        hash = @{ identifier ~ ["="] ~ param }
        exp_line = _{ identifier ~ (hash|param)* }

        subexpression = { ["("] ~ name ~ (hash|param)* ~ [")"] }

        pre_whitespace_omitter = { ["~"] }
        pro_whitespace_omitter = { ["~"] }

        expression = { !invert_tag ~ ["{{"] ~ pre_whitespace_omitter? ~ name ~
                        pro_whitespace_omitter? ~ ["}}"] }

        html_expression = { ["{{{"] ~ pre_whitespace_omitter? ~ name ~
                                      pro_whitespace_omitter? ~ ["}}}"] }

        helper_expression = { !invert_tag ~ ["{{"] ~ pre_whitespace_omitter? ~ exp_line ~
                               pro_whitespace_omitter? ~ ["}}"] }

        invert_tag = { ["{{else}}"]|["{{^}}"] }
        helper_block_start = { ["{{"] ~ pre_whitespace_omitter? ~ ["#"] ~ exp_line ~
                                        pro_whitespace_omitter? ~ ["}}"] }
        helper_block_end = { ["{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                      pro_whitespace_omitter? ~ ["}}"] }
        helper_block = { helper_block_start ~ template ~
                         (invert_tag ~ template)? ~
                         helper_block_end }

        raw_block_start = { ["{{{{"] ~ pre_whitespace_omitter? ~ ["#"] ~ exp_line ~
                                       pro_whitespace_omitter? ~ ["}}}}"] }
        raw_block_end = { ["{{{{"] ~ pre_whitespace_omitter? ~ ["/"] ~ name ~
                                     pro_whitespace_omitter? ~ ["}}}}"] }
        raw_block = { raw_block_start ~ raw_block_text ~ raw_block_end }

        comment = { ["{{!"] ~ (!["}}"] ~ any)* ~ ["}}"] }

        template = { (
            raw_text |
            expression |
            html_expression |
            helper_expression |
            helper_block |
            raw_block |
            comment )*
        }

        handlebars = _ { template ~ eoi }
    }
}

#[test]
fn test_raw_text() {
    let mut rdp = Rdp::new(StringInput::new("<h1> helloworld </h1>    "));
    assert!(rdp.raw_text());
    assert!(rdp.end());
}

#[test]
fn test_raw_block_text() {
    let mut rdp = Rdp::new(StringInput::new("<h1> {{hello}} </h1>"));
    assert!(rdp.raw_block_text());
    assert!(rdp.end());
}

#[test]
fn test_reference() {
    let s = vec!["a",
                 "abc",
                 "../a",
                 "a.b",
                 "@abc",
                 "a[\"abc\"]",
                 "aBc[\"abc\"]",
                 "abc[0][\"nice\"]"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.reference());
        assert!(rdp.end());
    }
}

#[test]
fn test_name() {
    let s = vec!["if", "(abc)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.name());
        assert!(rdp.end());
    }
}

#[test]
fn test_param() {
    let s = vec!["hello", "\"json literal\""];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.param());
        assert!(rdp.end());
    }
}

#[test]
fn test_hash() {
    let s = vec!["hello=world", "hello=\"world\"", "hello=(world)"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.hash());
        assert!(rdp.end());
    }
}

#[test]
fn test_json_literal() {
    let s = vec!["\"json string\"",
                 "\"quot: \\\"\"",
                 "[]",
                 "[\"hello\"]",
                 "[1,2,3,4,true]",
                 "{\"hello\": \"world\"}",
                 "{}",
                 "{\"a\":1, \"b\":2 }"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.literal());
        assert!(rdp.end());
    }
}

#[test]
fn test_comment() {
    let s = vec!["{{! hello }}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.comment());
        assert!(rdp.end());
    }
}

#[test]
fn test_expression() {
    let s = vec!["{{exp}}", "{{(exp)}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_expression() {
    let s = vec!["{{exp 1}}",
                 "{{exp \"literal\"}}",
                 "{{exp ref}}",
                 "{{exp (sub)}}",
                 "{{exp (sub 123)}}",
                 "{{exp []}}",
                 "{{exp {}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_html_expression() {
    let s = vec!["{{{html}}}", "{{{(html)}}}", "{{{(html)}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.html_expression());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_start() {
    let s = vec!["{{#if hello}}",
                 "{{#if (hello)}}",
                 "{{#if hello=world}}",
                 "{{#if hello hello=world}}",
                 "{{#if []}}",
                 "{{#if {}}}",
                 "{{#if}}",
                 "{{~#if hello~}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block_start());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_end() {
    let s = vec!["{{/if}}", "{{~/if}}", "{{~/if ~}}", "{{/if   ~}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block_end());
        assert!(rdp.end());
    }
}

#[test]
fn test_helper_block() {
    let s = vec!["{{#if hello}}hello{{/if}}",
                 "{{#if true}}hello{{/if}}",
                 "{{#if nice ok=1}}hello{{/if}}",
                 "{{#if}}hello{{else}}world{{/if}}",
                 "{{#if}}hello{{^}}world{{/if}}",
                 "{{#if}}{{#if}}hello{{/if}}{{/if}}",
                 "{{#if}}{{/if}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.helper_block());
        assert!(rdp.end());
    }
}

#[test]
fn test_raw_block() {
    let s = vec!["{{{{#if hello}}}}good {{hello}}{{{{/if}}}}",
                 "{{{{#if hello}}}}{{#if nice}}{{/if}}{{{{/if}}}}"];
    for i in s.iter() {
        let mut rdp = Rdp::new(StringInput::new(i));
        assert!(rdp.raw_block());
        assert!(rdp.end());
    }
}
<|endoftext|>