implement multiplex

ubw-sward/src/utils/multiplexed.rs

@@ -0,0 +1,75 @@
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+use tokio_util::sync::PollSemaphore;
+use tower::Service;
+
+pub struct MultiplexedSward<S> {
+    inner: S,
+    poll: PollSemaphore,
+    stash_permit: Option<OwnedSemaphorePermit>,
+}
+
+impl<S> MultiplexedSward<S> {
+    pub fn new(inner: S, capacity: usize) -> Self {
+        let semaphore = Arc::new(Semaphore::new(capacity));
+        Self {
+            inner,
+            poll: PollSemaphore::new(semaphore.clone()),
+            stash_permit: None,
+        }
+    }
+}
+
+impl<S, Req> Service<Req> for MultiplexedSward<S>
+where
+    S: Service<Req>,
+    S::Future: Send + 'static,
+{
+    type Response = S::Response;
+    type Error = MultiplexedSwardError<S::Error>;
+    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        if self.stash_permit.is_none() {
+            match self.poll.poll_acquire(cx) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(None) => {
+                    return Poll::Ready(Err(MultiplexedSwardError::SemaphoreClosed));
+                }
+                Poll::Ready(Some(permit)) => {
+                    self.stash_permit = Some(permit);
+                }
+            }
+        }
+        self.inner
+            .poll_ready(cx)
+            .map_err(MultiplexedSwardError::Inner)
+    }
+
+    fn call(&mut self, req: Req) -> Self::Future {
+        let Some(permit) = self.stash_permit.take() else {
+            return Box::pin(async { Err(MultiplexedSwardError::NotReady) });
+        };
+        let fut = self.inner.call(req);
+        Box::pin(async move {
+            let _ = permit;
+            let res = fut.await;
+            res.map_err(MultiplexedSwardError::Inner)
+        })
+    }
+}
+
+/// Error that can occur when multiplexing requests.
+#[derive(Debug, thiserror::Error)]
+pub enum MultiplexedSwardError<E> {
+    #[error("{0}")]
+    Inner(E),
+
+    #[error("Semaphore is closed.")]
+    SemaphoreClosed,
+
+    #[error("call() invoked without ready; call poll_ready first")]
+    NotReady,
+}