pd: use tokio_util::sync::PollSender in consensus svc

pd/Cargo.toml

@@ -42,7 +42,7 @@ csv = "1.1"
 directories = "4.0"
 tokio = { version = "1.16", features = ["full"]}
 tokio-stream = "0.1"
-tokio-util = "0.6"
+tokio-util = "0.7"
 tower = { version = "0.4", features = ["full"]}
 tracing = "0.1"
 regex = "1.5"

pd/src/consensus/service.rs

@@ -4,28 +4,18 @@ use std::{
     task::{Context, Poll},
 };
 
-use futures::{ready, FutureExt};
+use futures::FutureExt;
 use tendermint::abci::{ConsensusRequest, ConsensusResponse};
-use tokio::sync::{
-    mpsc::{self, error::SendError, OwnedPermit},
-    oneshot,
-};
-use tokio_util::sync::ReusableBoxFuture;
+use tokio::sync::{mpsc, oneshot};
+use tokio_util::sync::PollSender;
 use tower_abci::BoxError;
 
 use super::{Message, Worker};
 use crate::{state, RequestExt};
 
-enum State {
-    NoPermit,
-    Waiting,
-    Permit(OwnedPermit<Message>),
-}
-
+#[derive(Clone)]
 pub struct Consensus {
-    queue: mpsc::Sender<Message>,
-    future: ReusableBoxFuture<Result<OwnedPermit<Message>, SendError<()>>>,
-    state: State,
+    queue: PollSender<Message>,
 }
 
 impl Consensus {
@@ -35,60 +25,32 @@ impl Consensus {
         tokio::spawn(Worker::new(state, queue_rx).await?.run());
 
         Ok(Self {
-            queue: queue_tx,
-            state: State::NoPermit,
-            future: ReusableBoxFuture::new(async { unreachable!() }),
+            queue: PollSender::new(queue_tx),
         })
     }
 }
 
-impl Clone for Consensus {
-    fn clone(&self) -> Self {
-        Self {
-            queue: self.queue.clone(),
-            state: State::NoPermit,
-            future: ReusableBoxFuture::new(async { unreachable!() }),
-        }
-    }
-}
-
 impl tower::Service<ConsensusRequest> for Consensus {
     type Response = ConsensusResponse;
     type Error = BoxError;
     type Future =
         Pin<Box<dyn Future<Output = Result<ConsensusResponse, BoxError>> + Send + 'static>>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        match self.state {
-            State::Permit(_) => Poll::Ready(Ok(())),
-            State::NoPermit => {
-                self.future.set(self.queue.clone().reserve_owned());
-                self.state = State::Waiting;
-                self.poll_ready(cx)
-            }
-            State::Waiting => {
-                let permit = ready!(self.future.poll(cx))?;
-                self.state = State::Permit(permit);
-                Poll::Ready(Ok(()))
-            }
-        }
+        self.queue.poll_reserve(cx).map_err(Into::into)
     }
 
     fn call(&mut self, req: ConsensusRequest) -> Self::Future {
-        let permit = if let State::Permit(p) = std::mem::replace(&mut self.state, State::NoPermit) {
-            p
-        } else {
-            panic!("called without poll_ready");
-        };
-
         let span = req.create_span();
         let (tx, rx) = oneshot::channel();
 
-        permit.send(Message {
-            req,
-            rsp_sender: tx,
-            span,
-        });
+        self.queue
+            .send_item(Message {
+                req,
+                rsp_sender: tx,
+                span,
+            })
+            .expect("called without `poll_ready`");
 
         async move { Ok(rx.await.expect("worker error??")) }.boxed()
     }