cloneable_service.rs

use std::mem;
use std::sync::Arc;
use futures::task::{Context, Poll};
use tokio::sync::{Mutex, OwnedMutexGuard};
use tower_service::Service;
use tokio_util::sync::ReusableBoxFuture;

pub struct CloneableService<S> {
    inner: Arc<Mutex<S>>,
    mutex_fut: ReusableBoxFuture<OwnedMutexGuard<S>>,
    state: State<S>,
}

impl<S> CloneableService<S>
where S: Send + Sync + 'static
{
    pub fn new(inner: S) -> CloneableService<S> {
        let inner = Arc::new(Mutex::new(inner));
        let fut = Arc::clone(&inner).lock_owned();

        CloneableService {
            inner,
            mutex_fut: ReusableBoxFuture::new(fut),
            state: State::None,
        }
    }
}

impl<S> Clone for CloneableService<S>
where S: Send + Sync + 'static
{
    fn clone(&self) -> CloneableService<S> {

        let inner = self.inner.clone();
        let fut = Arc::clone(&inner).lock_owned();

        CloneableService {
            inner,
            mutex_fut: ReusableBoxFuture::new(fut),
            state: State::None,
        }
    }
}

enum State<S> {
    None,
    Waiting,
    Ready(OwnedMutexGuard<S>),
}

impl<S, Request> Service<Request> for CloneableService<S>
where
    S: Service<Request> + Send + Sync + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = S::Future;

    fn poll_ready(&mut self, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
        loop {
            match mem::replace(&mut self.state, State::None) {
                State::Ready(mut inner) => {
                    let poll = (*inner).poll_ready(cx);
                    self.state = State::Ready(inner);
                    return poll;
                }
                State::None => {
                    self.mutex_fut.set(self.inner.clone().lock_owned());
                    self.state = State::Waiting;
                }
                State::Waiting => match self.mutex_fut.poll(cx) {
                    Poll::Ready(lock) => self.state = State::Ready(lock),
                    Poll::Pending => {
                        self.state = State::Waiting;
                        return Poll::Pending;
                    }
                },
            }
        }
    }

    fn call(&mut self, request: Request) -> Self::Future {
        match mem::replace(&mut self.state, State::None) {
            State::Ready(mut svc) => svc.call(request),
            _ => panic!("called before ready"),
        }
    }
}