use std::future::{Future, poll_fn};
use std::pin::Pin;
use std::task::{Context, Poll};
use crate::teardown::{Deferred, Teardown};
use crate::wait::WaitCell;
use super::receiver::AdmitAll;
use super::{Receiver, RecvError, SendError, Sender, TryRecvError};
impl<T, P: Teardown> Sender<T, P> {
pub fn send_async(&mut self, value: T) -> SendFuture<'_, T, P> {
SendFuture {
sender: self,
value: Some(value),
}
}
pub fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), SendError<()>>> {
if !self.is_disconnected() && self.producer.is_full() {
self.signal.space_waiter.register_async(cx.waker());
if !self.is_disconnected() && self.producer.is_full() {
return Poll::Pending;
}
}
self.cancel_send_wait();
Poll::Ready(if self.is_disconnected() {
Err(SendError(()))
} else {
Ok(())
})
}
pub fn cancel_send_wait(&mut self) {
self.signal.space_waiter.cancel_async();
}
}
pub struct SendFuture<'a, T, P: Teardown = Deferred> {
sender: &'a mut Sender<T, P>,
value: Option<T>,
}
impl<T, P: Teardown> std::fmt::Debug for SendFuture<'_, T, P> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SendFuture")
.field("pending_value", &self.value.is_some())
.finish_non_exhaustive()
}
}
impl<T, P: Teardown> Unpin for SendFuture<'_, T, P> {}
impl<T, P: Teardown> Future for SendFuture<'_, T, P> {
type Output = Result<(), SendError<T>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
if this.sender.poll_ready(cx).is_pending() {
return Poll::Pending;
}
let value = this.value.take().expect("send polled after completion");
Poll::Ready(
this.sender
.try_send(value)
.map_err(|error| SendError(error.into_inner())),
)
}
}
impl<T, P: Teardown> Drop for SendFuture<'_, T, P> {
fn drop(&mut self) {
self.sender.cancel_send_wait();
}
}
struct Registration<'a>(&'a WaitCell);
impl Drop for Registration<'_> {
fn drop(&mut self) {
self.0.cancel_async();
}
}
impl<T, P: Teardown> Receiver<T, P> {
pub async fn recv_async(&mut self) -> Result<T, RecvError> {
let shared = self.shared.clone();
let _registration = Registration(&shared.data_waiter);
poll_fn(|cx| self.poll_recv(cx)).await
}
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> {
let mut result = self.try_recv();
if matches!(result, Err(TryRecvError::Empty)) {
self.release_consumed();
self.shared.data_waiter.register_async(cx.waker());
result = self.try_recv();
}
self.release_consumed();
match result {
Ok(value) => {
self.cancel_recv_wait();
Poll::Ready(Ok(value))
}
Err(TryRecvError::Disconnected) => {
self.cancel_recv_wait();
Poll::Ready(Err(RecvError))
}
Err(TryRecvError::Empty) => Poll::Pending,
}
}
pub fn cancel_recv_wait(&mut self) {
self.shared.data_waiter.cancel_async();
}
pub async fn recv_batch_into_async(
&mut self,
output: &mut Vec<T>,
limit: usize,
) -> Result<usize, RecvError> {
if limit == 0 {
self.release_consumed();
return Ok(0);
}
let mut received = self.drain_into(output, limit, &mut AdmitAll).received;
if received == 0 {
match self.recv_async().await {
Ok(first) => {
output.push(first);
received = 1 + self.drain_into(output, limit - 1, &mut AdmitAll).received;
}
Err(error) => {
self.release_consumed();
return Err(error);
}
}
}
self.release_consumed();
Ok(received)
}
}