use bytes::Bytes;
use futures::{SinkExt, StreamExt, future::Either};
use tokio_util::sync::CancellationToken;
use super::super::bridge;
use crate::{
codec::BoxReadStream,
message::stream::guard,
quic::{self, CancelStreamExt, GetStreamIdExt, StopStreamExt},
util::deferred::Deferred,
varint::VarInt,
};
#[remoc::rtc::remote]
pub trait ReadStream: Send {
async fn stream_id(&mut self) -> Result<VarInt, quic::StreamError>;
async fn read(&mut self) -> Result<Option<Bytes>, quic::StreamError>;
async fn stop(&mut self, code: VarInt) -> Result<(), quic::StreamError>;
}
impl ReadStreamClient {
pub async fn into_quic(mut self) -> Result<impl quic::ReadStream, quic::StreamError> {
let stream_id = self.stream_id().await?;
Ok(bridge::ReadBridge::<_, quic::StreamError, _, _, _, _>::new(
stream_id,
self,
|mut client: ReadStreamClient, token: CancellationToken| async move {
tokio::select! {
res = client.read() => Either::Left((client, res.transpose())),
_ = token.cancelled() => Either::Right(client),
}
},
|mut client: ReadStreamClient, code| async move {
let res = client.stop(code).await;
(client, res)
},
))
}
pub fn into_boxed_quic(self) -> guard::GuardedQuicReader {
let raw: BoxReadStream = Box::pin(Deferred::from(self.into_quic()));
guard::GuardedQuicReader::new(raw)
}
}
#[remoc::rtc::remote]
pub trait WriteStream: Send {
async fn stream_id(&mut self) -> Result<VarInt, quic::StreamError>;
async fn write(&mut self, data: Bytes) -> Result<(), quic::StreamError>;
async fn flush(&mut self) -> Result<(), quic::StreamError>;
async fn shutdown(&mut self) -> Result<(), quic::StreamError>;
async fn cancel(&mut self, code: VarInt) -> Result<(), quic::StreamError>;
}
impl WriteStreamClient {
pub async fn into_quic(mut self) -> Result<impl quic::WriteStream, quic::StreamError> {
let stream_id = self.stream_id().await?;
Ok(bridge::WriteBridge::<
_,
quic::StreamError,
_,
_,
_,
_,
_,
_,
_,
_,
>::new(
stream_id,
self,
|mut client: WriteStreamClient, token: CancellationToken, bytes| async move {
tokio::select! {
res = client.write(bytes) => Either::Left((client, res)),
_ = token.cancelled() => Either::Right(client),
}
},
|mut client: WriteStreamClient, token: CancellationToken| async move {
tokio::select! {
res = client.flush() => Either::Left((client, res)),
_ = token.cancelled() => Either::Right(client),
}
},
|mut client: WriteStreamClient, token: CancellationToken| async move {
tokio::select! {
res = client.shutdown() => Either::Left((client, res)),
_ = token.cancelled() => Either::Right(client),
}
},
|mut client: WriteStreamClient, code| async move {
let res = client.cancel(code).await;
(client, res)
},
))
}
pub fn into_boxed_quic(self) -> guard::GuardedQuicWriter {
let raw: crate::codec::BoxWriteStream = Box::pin(Deferred::from(self.into_quic()));
guard::GuardedQuicWriter::new(raw)
}
}
impl<S> ReadStream for S
where
S: quic::ReadStream + Unpin + Send,
{
async fn stream_id(&mut self) -> Result<VarInt, quic::StreamError> {
GetStreamIdExt::stream_id(self).await
}
async fn read(&mut self) -> Result<Option<Bytes>, quic::StreamError> {
StreamExt::next(self).await.transpose()
}
async fn stop(&mut self, code: VarInt) -> Result<(), quic::StreamError> {
StopStreamExt::stop(self, code).await
}
}
impl<S> WriteStream for S
where
S: quic::WriteStream + Unpin + Send,
{
async fn stream_id(&mut self) -> Result<VarInt, quic::StreamError> {
GetStreamIdExt::stream_id(self).await
}
async fn write(&mut self, data: Bytes) -> Result<(), quic::StreamError> {
SinkExt::send(self, data).await
}
async fn flush(&mut self) -> Result<(), quic::StreamError> {
SinkExt::flush(self).await
}
async fn shutdown(&mut self) -> Result<(), quic::StreamError> {
SinkExt::close(self).await
}
async fn cancel(&mut self, code: VarInt) -> Result<(), quic::StreamError> {
CancelStreamExt::cancel(self, code).await
}
}