diff --git a/src/stream.rs b/src/stream.rs index 1f03cab0..4b90ff14 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1,8 +1,10 @@ +use std::marker::PhantomData; + use futures::unsync::mpsc; -use futures::{Async, Future, Poll, Stream}; +use futures::{future, Async, Future, Poll, Stream}; use tokio::executor::current_thread::spawn; -use super::service::{IntoService, Service}; +use super::service::{IntoService, NewService, Service}; pub struct StreamDispatcher { stream: S, @@ -88,3 +90,65 @@ impl Future for StreamDispatcherService { } } } + +/// `NewService` that implements, read one item from the stream. +pub struct TakeItem { + _t: PhantomData, +} + +impl TakeItem { + /// Create new `TakeRequest` instance. + pub fn new() -> Self { + TakeItem { _t: PhantomData } + } +} + +impl NewService for TakeItem { + type Request = T; + type Response = (Option, T); + type Error = T::Error; + type InitError = (); + type Service = TakeItemService; + type Future = future::FutureResult; + + fn new_service(&self) -> Self::Future { + future::ok(TakeItemService { _t: PhantomData }) + } +} + +/// `NewService` that implements, read one request from framed object feature. +pub struct TakeItemService { + _t: PhantomData, +} + +impl Service for TakeItemService { + type Request = T; + type Response = (Option, T); + type Error = T::Error; + type Future = TakeItemServiceResponse; + + fn poll_ready(&mut self) -> Poll<(), Self::Error> { + Ok(Async::Ready(())) + } + + fn call(&mut self, req: Self::Request) -> Self::Future { + TakeItemServiceResponse { stream: Some(req) } + } +} + +#[doc(hidden)] +pub struct TakeItemServiceResponse { + stream: Option, +} + +impl Future for TakeItemServiceResponse { + type Item = (Option, T); + type Error = T::Error; + + fn poll(&mut self) -> Poll { + match self.stream.as_mut().expect("Use after finish").poll()? { + Async::Ready(item) => Ok(Async::Ready((item, self.stream.take().unwrap()))), + Async::NotReady => Ok(Async::NotReady), + } + } +}