use std::collections::VecDeque;
use std::fmt;
use std::sync::Arc;
use rmux_proto::{encode_frame, FrameDecoder, Request, Response, SdkWaitId, SdkWaitOwnerId};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::sync::{mpsc, oneshot};
use crate::{Result, RmuxError};
mod failure;
mod pending;
mod state;
use failure::TransportFailure;
use pending::PendingCall;
pub(crate) use pending::PendingResponse;
use state::TransportState;
#[cfg(test)]
use state::{allocate_bounded_atomic_id, mix_sdk_wait_owner_id};
const ACTOR_QUEUE_CAPACITY: usize = 64;
const READ_BUFFER_SIZE: usize = 8192;
const TRANSPORT_SHUTDOWN_OPERATION: &str = "shut down rmux SDK transport";
#[derive(Clone)]
pub(crate) struct TransportClient {
commands: mpsc::Sender<ActorMessage>,
state: Arc<TransportState>,
}
impl TransportClient {
pub(crate) fn spawn<S>(stream: S) -> Self
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
let (commands, receiver) = mpsc::channel(ACTOR_QUEUE_CAPACITY);
let state = Arc::new(TransportState::default());
tokio::spawn(run_actor(stream, receiver, state.clone()));
Self { commands, state }
}
pub(crate) async fn request(&self, request: Request) -> Result<Response> {
let operation = request_operation(&request);
if let Some(failure) = self.state.terminal_failure() {
return Err(failure.to_error(&operation));
}
let (reply, response) = oneshot::channel();
self.commands
.send(ActorMessage::Request {
request,
operation: operation.clone(),
reply,
})
.await
.map_err(|_| self.closed_error(&operation))?;
response.await.map_err(|_| self.closed_error(&operation))?
}
pub(crate) async fn armed_request(&self, request: Request) -> Result<PendingResponse> {
let operation = request_operation(&request);
if let Some(failure) = self.state.terminal_failure() {
return Err(failure.to_error(&operation));
}
let (reply, response) = oneshot::channel();
let (armed, armed_response) = oneshot::channel();
self.commands
.send(ActorMessage::ArmedRequest {
request,
operation: operation.clone(),
reply,
armed,
})
.await
.map_err(|_| self.closed_error(&operation))?;
armed_response
.await
.map_err(|_| self.closed_error(&operation))?
.map_err(|failure| failure.to_error(&operation))?;
Ok(PendingResponse::new(operation, response))
}
pub(crate) async fn shutdown(&self) -> Result<()> {
if let Some(failure) = self.state.terminal_failure() {
if failure.is_eof() {
return Ok(());
}
return Err(failure.to_error(TRANSPORT_SHUTDOWN_OPERATION));
}
let (reply, response) = oneshot::channel();
self.commands
.send(ActorMessage::Shutdown { reply })
.await
.map_err(|_| self.closed_error(TRANSPORT_SHUTDOWN_OPERATION))?;
response
.await
.map_err(|_| self.closed_error(TRANSPORT_SHUTDOWN_OPERATION))?
}
fn try_send_best_effort(&self, request: Request) {
if self.state.terminal_failure().is_some() {
return;
}
let _ = self.commands.try_send(ActorMessage::BestEffort { request });
}
pub(crate) fn sdk_wait_owner_id(&self) -> SdkWaitOwnerId {
self.state.sdk_wait_owner_id()
}
pub(crate) fn allocate_sdk_wait_id(&self) -> SdkWaitId {
self.state.allocate_sdk_wait_id()
}
fn closed_error(&self, operation: &str) -> RmuxError {
self.state
.terminal_failure()
.unwrap_or_else(TransportFailure::actor_closed)
.to_error(operation)
}
}
impl fmt::Debug for TransportClient {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("TransportClient")
.finish_non_exhaustive()
}
}
#[derive(Debug, Default)]
pub(crate) struct DropGuard {
action: DropAction,
}
impl DropGuard {
pub(crate) fn noop() -> Self {
Self {
action: DropAction::None,
}
}
pub(crate) fn best_effort(client: TransportClient, request: Request) -> Self {
Self {
action: DropAction::BestEffort {
client,
request: Some(Box::new(request)),
},
}
}
pub(crate) fn disarm(&mut self) {
self.action = DropAction::None;
}
pub(crate) fn trigger(&mut self) {
if let DropAction::BestEffort { client, request } = &mut self.action {
if let Some(request) = request.take() {
client.try_send_best_effort(*request);
}
}
self.action = DropAction::None;
}
}
impl Drop for DropGuard {
fn drop(&mut self) {
self.trigger();
}
}
#[derive(Debug, Default)]
enum DropAction {
#[default]
None,
BestEffort {
client: TransportClient,
request: Option<Box<Request>>,
},
}
enum ActorMessage {
Request {
request: Request,
operation: String,
reply: oneshot::Sender<Result<Response>>,
},
ArmedRequest {
request: Request,
operation: String,
reply: oneshot::Sender<Result<Response>>,
armed: oneshot::Sender<core::result::Result<(), TransportFailure>>,
},
BestEffort {
request: Request,
},
Shutdown {
reply: oneshot::Sender<Result<()>>,
},
}
enum ActorEvent {
Command(ActorMessage),
CommandsClosed,
Response(core::result::Result<Response, TransportFailure>),
}
async fn run_actor<S>(stream: S, commands: mpsc::Receiver<ActorMessage>, state: Arc<TransportState>)
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
let (reader, mut writer) = tokio::io::split(stream);
let (events, mut event_receiver) = mpsc::channel(ACTOR_QUEUE_CAPACITY * 2);
let command_task = tokio::spawn(forward_commands(commands, events.clone()));
let read_task = tokio::spawn(forward_responses(reader, events));
let mut pending = VecDeque::new();
let mut commands_closed = false;
let mut shutdown_reply = None;
while let Some(event) = event_receiver.recv().await {
match event {
ActorEvent::Command(message) => {
if shutdown_reply.is_some() {
reject_command_after_shutdown(message);
continue;
}
match message {
ActorMessage::Request {
request,
operation,
reply,
} => {
let command_name = request.command_name();
let frame = match encode_request(&request) {
Ok(frame) => frame,
Err(failure) => {
let _ = reply.send(Err(failure.to_error(&operation)));
continue;
}
};
pending.push_back(PendingCall::reply(command_name, operation, reply));
if let Err(failure) = write_frame(&mut writer, &frame).await {
fail_transport(&mut pending, &state, failure);
break;
}
}
ActorMessage::ArmedRequest {
request,
operation,
reply,
armed,
} => {
let command_name = request.command_name();
let frame = match encode_request(&request) {
Ok(frame) => frame,
Err(failure) => {
let _ = reply.send(Err(failure.to_error(&operation)));
let _ = armed.send(Err(failure));
continue;
}
};
pending.push_back(PendingCall::reply(
command_name,
operation.clone(),
reply,
));
if let Err(failure) = write_frame(&mut writer, &frame).await {
let _ = armed.send(Err(failure.clone()));
fail_transport(&mut pending, &state, failure);
break;
}
let _ = armed.send(Ok(()));
}
ActorMessage::BestEffort { request } => {
let command_name = request.command_name();
let Ok(frame) = encode_request(&request) else {
continue;
};
pending.push_back(PendingCall::discard(
command_name,
request_operation(&request),
));
if let Err(failure) = write_frame(&mut writer, &frame).await {
fail_transport(&mut pending, &state, failure);
break;
}
}
ActorMessage::Shutdown { reply } => {
match writer.shutdown().await.map_err(TransportFailure::io) {
Ok(()) => {
shutdown_reply = Some(reply);
}
Err(failure) => {
let _ =
reply.send(Err(failure.to_error(TRANSPORT_SHUTDOWN_OPERATION)));
fail_transport(&mut pending, &state, failure);
break;
}
}
}
}
}
ActorEvent::CommandsClosed => {
commands_closed = true;
}
ActorEvent::Response(result) => match result {
Ok(response) => {
let Some(pending_call) = pending.pop_front() else {
let failure = TransportFailure::unsolicited_response(&response);
fail_shutdown(&mut shutdown_reply, &failure);
fail_transport(&mut pending, &state, failure);
break;
};
if let Err(failure) = pending_call.validate_response(&response) {
pending_call.fail(&failure);
fail_shutdown(&mut shutdown_reply, &failure);
fail_transport(&mut pending, &state, failure);
break;
}
pending_call.complete(response);
}
Err(failure) => {
if shutdown_reply.is_some() && pending.is_empty() && failure.is_eof() {
complete_shutdown(&mut shutdown_reply);
break;
}
fail_shutdown(&mut shutdown_reply, &failure);
fail_transport(&mut pending, &state, failure);
break;
}
},
}
if commands_closed && pending.is_empty() && shutdown_reply.is_none() {
let _ = writer.shutdown().await;
break;
}
}
command_task.abort();
read_task.abort();
}
fn reject_command_after_shutdown(message: ActorMessage) {
match message {
ActorMessage::Request {
operation, reply, ..
} => {
let failure = TransportFailure::actor_closed();
let _ = reply.send(Err(failure.to_error(&operation)));
}
ActorMessage::ArmedRequest {
operation,
reply,
armed,
..
} => {
let failure = TransportFailure::actor_closed();
let _ = reply.send(Err(failure.to_error(&operation)));
let _ = armed.send(Err(failure));
}
ActorMessage::BestEffort { .. } => {}
ActorMessage::Shutdown { reply } => {
let failure = TransportFailure::actor_closed();
let _ = reply.send(Err(failure.to_error(TRANSPORT_SHUTDOWN_OPERATION)));
}
}
}
fn complete_shutdown(reply: &mut Option<oneshot::Sender<Result<()>>>) {
if let Some(reply) = reply.take() {
let _ = reply.send(Ok(()));
}
}
fn fail_shutdown(reply: &mut Option<oneshot::Sender<Result<()>>>, failure: &TransportFailure) {
if let Some(reply) = reply.take() {
let _ = reply.send(Err(failure.to_error(TRANSPORT_SHUTDOWN_OPERATION)));
}
}
async fn forward_commands(
mut commands: mpsc::Receiver<ActorMessage>,
events: mpsc::Sender<ActorEvent>,
) {
while let Some(message) = commands.recv().await {
if events.send(ActorEvent::Command(message)).await.is_err() {
return;
}
}
let _ = events.send(ActorEvent::CommandsClosed).await;
}
async fn forward_responses<R>(mut reader: R, events: mpsc::Sender<ActorEvent>)
where
R: AsyncRead + Unpin,
{
let mut decoder = FrameDecoder::new();
loop {
let result = read_response(&mut reader, &mut decoder).await;
let stop = result.is_err();
if events.send(ActorEvent::Response(result)).await.is_err() {
return;
}
if stop {
return;
}
}
}
fn encode_request(request: &Request) -> core::result::Result<Vec<u8>, TransportFailure> {
encode_frame(request).map_err(TransportFailure::frame)
}
async fn write_frame<W>(writer: &mut W, frame: &[u8]) -> core::result::Result<(), TransportFailure>
where
W: AsyncWrite + Unpin,
{
writer
.write_all(frame)
.await
.map_err(TransportFailure::io)?;
writer.flush().await.map_err(TransportFailure::io)
}
async fn read_response<R>(
reader: &mut R,
decoder: &mut FrameDecoder,
) -> core::result::Result<Response, TransportFailure>
where
R: AsyncRead + Unpin,
{
let mut buffer = [0; READ_BUFFER_SIZE];
loop {
if let Some(response) = decoder
.next_frame::<Response>()
.map_err(TransportFailure::frame)?
{
return Ok(response);
}
let read = reader
.read(&mut buffer)
.await
.map_err(TransportFailure::io)?;
if read == 0 {
return Err(TransportFailure::eof());
}
decoder.push_bytes(&buffer[..read]);
}
}
fn fail_all(pending: &mut VecDeque<PendingCall>, failure: &TransportFailure) {
while let Some(call) = pending.pop_front() {
call.fail(failure);
}
}
fn fail_transport(
pending: &mut VecDeque<PendingCall>,
state: &TransportState,
failure: TransportFailure,
) {
state.set_terminal_failure(failure.clone());
fail_all(pending, &failure);
}
fn request_operation(request: &Request) -> String {
format!(
"complete `{}` request/response exchange with rmux daemon",
request.command_name()
)
}
#[cfg(test)]
mod tests;