use super::super::types::{RemoteStream, RemoteUpdate, StreamEventFrame};
use indicatrix_net::{
client::ApplyOutcome,
framing::{self, LEN_PREFIX_BYTES, MAX_CONTROL_FRAME_LEN, MAX_FRAME_LEN},
messages::{ErrorMsg, NetError, StreamEvent, error_codes},
};
use std::io::Read;
pub(super) fn to_remote_update(
request_id: u32,
event: &StreamEvent,
outcome: ApplyOutcome,
) -> Option<RemoteUpdate> {
match outcome {
ApplyOutcome::FrameSummed { samples_done } => Some(RemoteUpdate::Frame {
request_id,
samples_done,
}),
ApplyOutcome::PreviewReplaced => Some(RemoteUpdate::Preview { request_id }),
ApplyOutcome::Progress { samples_done } => Some(RemoteUpdate::Progress {
request_id,
samples_done,
}),
ApplyOutcome::Done { cancelled } => Some(RemoteUpdate::Done {
request_id,
cancelled,
}),
ApplyOutcome::WorkerError => {
let StreamEvent::Error(e) = event else {
unreachable!("WorkerError only ever comes from applying an Error event")
};
Some(worker_error_update(request_id, e))
}
ApplyOutcome::DisplayFrameReplaced => {
let StreamEvent::DisplayFrame(h) = event else {
unreachable!("DisplayFrameReplaced only comes from a DisplayFrame event")
};
Some(RemoteUpdate::DisplayFrame {
request_id,
samples_done: h.samples_done,
})
}
ApplyOutcome::FinalImageReceived => Some(RemoteUpdate::FinalImage { request_id }),
ApplyOutcome::CapabilityChanged => {
let StreamEvent::CapabilityChanged { render } = event else {
unreachable!("CapabilityChanged only comes from a CapabilityChanged event")
};
Some(RemoteUpdate::CapabilityChanged {
request_id,
render: render.clone(),
})
}
ApplyOutcome::StaleDropped | ApplyOutcome::Pong { .. } | ApplyOutcome::NeedAsset { .. } => {
None
}
}
}
const ALL_WORKERS_LOST_NOTE: &str = "All remote workers were lost";
pub(super) fn worker_error_update(request_id: u32, error: &ErrorMsg) -> RemoteUpdate {
match error.code {
error_codes::UNSUPPORTED_REQUEST => RemoteUpdate::Unsupported {
request_id,
message: error.message.clone(),
},
error_codes::ALL_WORKERS_LOST => RemoteUpdate::Failed {
request_id,
message: if error.message.is_empty() {
ALL_WORKERS_LOST_NOTE.to_string()
} else {
format!("{ALL_WORKERS_LOST_NOTE} ({})", error.message)
},
},
_ => RemoteUpdate::Failed {
request_id,
message: error.message.clone(),
},
}
}
fn try_read_one_frame(
stream: &mut RemoteStream,
poll_timeout: std::time::Duration,
) -> Result<Option<Vec<u8>>, NetError> {
stream
.sock
.set_read_timeout(Some(poll_timeout))
.map_err(|e| NetError::Framing(framing::FramingError::Io(e)))?;
let mut len_bytes = [0u8; LEN_PREFIX_BYTES];
let n = match stream.read(&mut len_bytes) {
Ok(0) => {
return Err(NetError::Framing(framing::FramingError::Io(
std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"worker closed the connection",
),
)));
}
Ok(n) => n,
Err(e)
if e.kind() == std::io::ErrorKind::WouldBlock
|| e.kind() == std::io::ErrorKind::TimedOut =>
{
return Ok(None);
}
Err(e) => return Err(NetError::Framing(framing::FramingError::Io(e))),
};
stream
.sock
.set_read_timeout(Some(super::FRAME_REMAINDER_TIMEOUT))
.map_err(|e| NetError::Framing(framing::FramingError::Io(e)))?;
let payload = framing::read_frame_continuing(stream, &len_bytes[..n], MAX_CONTROL_FRAME_LEN)
.map_err(NetError::Framing)?;
Ok(Some(payload))
}
pub(super) fn try_read_stream_event(
stream: &mut RemoteStream,
poll_timeout: std::time::Duration,
) -> Result<Option<StreamEventFrame>, NetError> {
let Some(header_bytes) = try_read_one_frame(stream, poll_timeout)? else {
return Ok(None);
};
let event: StreamEvent = postcard::from_bytes(&header_bytes)?;
let payload = match event.payload_len() {
Some(expected) => {
let bytes = framing::read_frame_bounded(stream, expected.min(MAX_FRAME_LEN))
.map_err(NetError::Framing)?;
if bytes.len() as u32 != expected {
return Err(NetError::FramePayloadLenMismatch {
declared: expected,
actual: bytes.len(),
});
}
Some(bytes)
}
None => None,
};
Ok(Some((event, payload)))
}