cloneable_service.rs

use std::future::Future;
use std::mem;
use std::pin::Pin;
use std::sync::Arc;
use futures::task::{Context, Poll};
use tokio::sync::{Mutex, OwnedMutexGuard};
use tower_service::service::Service;

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

impl<S> CloneableService<S> {
    pub fn new(inner: S) -> CloneableService<S> {
        let inner = Arc::new(Mutex::new(inner));
        CloneableService {
            inner,
            state: State::None,
        }
    }
}

impl<S> Clone for CloneableService<S> {
    fn clone(&self) -> CloneableService<S> {
        CloneableService {
            inner: self.inner.clone(),
            state: State::None,
        }
    }
}

enum State<S> {
    None,
    Waiting(Pin<Box<dyn Future<Output = OwnedMutexGuard<S>> + Send + Sync>>),
    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(inner) => {
                    let poll = (*inner).poll_ready(cx);
                    self.state = State::Ready(inner);
                    return poll;
                }
                State::None => {
                    self.state = State::Waiting(Box::pin(self.inner.clone().lock_owned()));
                }
                State::Waiting(fut) => match fut.as_mut().poll(cx) {
                    Poll::Ready(lock) => self.state = State::Ready(lock),
                    Poll::Pending => {
                        self.state = State::Waiting(fut);
                        return Poll::Pending;
                    }
                },
            }
        }
    }

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