use std::{
env,
path::PathBuf,
process::{ExitStatus, Stdio},
sync::Once,
time::Duration,
};
use futures_util::{
SinkExt, StreamExt,
stream::{SplitSink, SplitStream},
};
use monty_proto::{FrameError, MAX_FRAME_LEN, decode_frame, encode_framed_into, encode_to_capped_vec, pb};
use rustls::crypto::aws_lc_rs::default_provider;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpStream,
process::{Child, ChildStdin, ChildStdout, Command},
runtime::{Handle, RuntimeFlavor},
sync::mpsc,
task::{JoinHandle, block_in_place},
time::timeout,
};
use tokio_tungstenite::{
MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
tungstenite::{
Bytes, Error as WsError, Message,
client::IntoClientRequest,
http::{HeaderName, HeaderValue, header::USER_AGENT},
protocol::WebSocketConfig,
},
};
#[cfg(feature = "telemetry")]
use crate::telemetry::{TelemetryContext, metrics::TurnMetrics, tracing::Recorder};
use crate::{MontyTransport, PoolConfig, PoolError};
type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
pub(crate) struct Worker {
kind: WorkerKind,
pub(crate) id: u64,
pub(crate) checkouts_served: u32,
#[cfg(feature = "telemetry")]
recorder: Option<Box<Recorder>>,
#[cfg(feature = "telemetry")]
metrics: Option<Box<TurnMetrics>>,
}
enum WorkerKind {
Subprocess(Box<SubprocessWorker>),
WebSocket(Box<WebSocketWorker>),
}
struct SubprocessWorker {
child: Child,
stdin: ChildStdin,
recv: FrameRecv,
send_buf: Vec<u8>,
}
struct WebSocketWorker {
sink: Option<SplitSink<WsStream, Message>>,
events: mpsc::Receiver<Result<Bytes, FrameError>>,
reader: JoinHandle<()>,
}
impl WebSocketWorker {
async fn close(&mut self) {
if let Some(sink) = &mut self.sink {
send_close(sink).await;
}
self.reader.abort();
self.sink = None;
}
}
impl Drop for WebSocketWorker {
fn drop(&mut self) {
self.reader.abort();
if let Some(mut sink) = self.sink.take()
&& let Ok(handle) = Handle::try_current()
{
handle.spawn(async move { send_close(&mut sink).await });
}
}
}
async fn send_close(sink: &mut SplitSink<WsStream, Message>) {
let _ = timeout(CLOSE_WRITE_TIMEOUT, sink.close()).await;
}
impl Worker {
pub(crate) async fn new(config: &PoolConfig, connect_headers: &[(String, String)]) -> Result<Self, PoolError> {
let worker = match &config.transport {
MontyTransport::Subprocess(binary_path) => Self::subprocess(binary_path),
MontyTransport::Websocket(url) => {
Self::websocket(
url,
config.request_timeout.unwrap_or(DEFAULT_DIAL_TIMEOUT),
connect_headers,
)
.await
}
};
#[cfg(feature = "telemetry")]
let worker = worker.map(|mut worker| {
worker.metrics = config
.metrics
.clone()
.map(|metrics| Box::new(TurnMetrics::new(metrics)));
worker
});
worker
}
fn subprocess(binary_path: &PathBuf) -> Result<Self, PoolError> {
let mut command = Command::new(binary_path);
command
.arg("subprocess")
.env_clear()
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.kill_on_drop(true);
if cfg!(windows)
&& let Ok(system_root) = env::var("SystemRoot")
{
command.env("SystemRoot", system_root);
}
let mut child = command
.spawn()
.map_err(|err| PoolError::Spawn(format!("{}: {err}", binary_path.display())))?;
let stdin = child.stdin.take().expect("piped stdin");
let stdout = child.stdout.take().expect("piped stdout");
Ok(Self {
id: 0,
kind: WorkerKind::Subprocess(Box::new(SubprocessWorker {
child,
stdin,
recv: FrameRecv::new(stdout),
send_buf: Vec::with_capacity(SEND_BUF_CAPACITY),
})),
checkouts_served: 0,
#[cfg(feature = "telemetry")]
recorder: None,
#[cfg(feature = "telemetry")]
metrics: None,
})
}
async fn websocket(
url: &str,
dial_timeout: Duration,
connect_headers: &[(String, String)],
) -> Result<Self, PoolError> {
install_crypto_provider();
let ws_config = WebSocketConfig::default()
.max_frame_size(Some(MAX_FRAME_LEN as usize))
.max_message_size(Some(MAX_FRAME_LEN as usize));
let mut request = url
.into_client_request()
.map_err(|err| PoolError::Spawn(format!("{url}: {err}")))?;
request
.headers_mut()
.insert(USER_AGENT, HeaderValue::from_static(USER_AGENT_VALUE));
for (name, value) in connect_headers {
let name = HeaderName::from_bytes(name.as_bytes())
.map_err(|err| PoolError::Spawn(format!("{url}: connect header {name:?}: {err}")))?;
let value = HeaderValue::from_str(value)
.map_err(|err| PoolError::Spawn(format!("{url}: connect header {name:?} value: {err}")))?;
request.headers_mut().insert(name, value);
}
let dial = connect_async_tls_with_config(request, Some(ws_config), true, None);
let (stream, _response) = timeout(dial_timeout, dial)
.await
.map_err(|_| PoolError::Spawn(format!("{url}: connect timed out after {dial_timeout:?}")))?
.map_err(|err| PoolError::Spawn(format!("{url}: {err}")))?;
let (sink, read_half) = stream.split();
let (tx, events) = mpsc::channel(WS_EVENT_CHANNEL_DEPTH);
let reader = tokio::spawn(read_ws_events(read_half, tx));
Ok(Self {
id: 0,
kind: WorkerKind::WebSocket(Box::new(WebSocketWorker {
sink: Some(sink),
events,
reader,
})),
checkouts_served: 0,
#[cfg(feature = "telemetry")]
recorder: None,
#[cfg(feature = "telemetry")]
metrics: None,
})
}
#[cfg(feature = "telemetry")]
pub(crate) fn callback_context(&self) -> opentelemetry::Context {
self.recorder
.as_ref()
.map_or_else(opentelemetry::Context::new, |recorder| recorder.callback_context())
}
pub(crate) async fn send(&mut self, request: &pb::ParentRequest) -> Result<(), FrameError> {
let result = match &mut self.kind {
WorkerKind::Subprocess(w) => {
encode_framed_into(request, &mut w.send_buf)?;
w.stdin.write_all(&w.send_buf).await?;
let len = w.send_buf.len() - 4;
if w.send_buf.capacity() > RETAIN_BUF_MAX {
w.send_buf = Vec::with_capacity(SEND_BUF_CAPACITY);
}
Ok(len)
}
WorkerKind::WebSocket(w) => {
let body = encode_to_capped_vec(request)?;
let len = body.len();
match &mut w.sink {
Some(sink) => sink
.send(Message::Binary(body.into()))
.await
.map(|()| len)
.map_err(ws_to_frame_error),
None => Err(FrameError::Truncated),
}
}
};
#[cfg(feature = "telemetry")]
if let Ok(len) = result {
if let Some(recorder) = &mut self.recorder {
recorder.begin_turn(request);
}
if let Some(metrics) = &mut self.metrics {
metrics.frame("sent", len);
metrics.begin_turn(request);
}
if matches!(
request.kind,
Some(pb::parent_request::Kind::Reset(_) | pb::parent_request::Kind::Shutdown(_))
) {
self.recorder = None;
}
}
result.map(drop)
}
#[cfg(feature = "telemetry")]
pub(crate) fn with_adapter_context(mut self, context: Option<TelemetryContext>) -> Self {
if let Some(context) = context {
let worker_pid = self.pid();
self.recorder
.get_or_insert_with(|| Box::new(Recorder::new(worker_pid)))
.set_adapter_context(context);
}
self
}
#[cfg(feature = "telemetry")]
pub(crate) fn mark_resumed(&mut self) {
if let Some(recorder) = &mut self.recorder {
recorder.mark_resumed();
}
}
pub(crate) async fn recv(&mut self) -> Result<pb::ChildEvent, FrameError> {
let (event, len) = match &mut self.kind {
WorkerKind::Subprocess(w) => match w.recv.recv().await? {
Some(framed) => Ok(framed),
None => Err(FrameError::Truncated), },
WorkerKind::WebSocket(w) => {
if w.sink.is_none() {
return Err(FrameError::Truncated);
}
match w.events.recv().await {
Some(Ok(data)) => decode_event(&data).map(|event| (event, data.len())),
Some(Err(err)) => Err(err),
None => Err(FrameError::Truncated),
}
}
}?;
#[cfg(feature = "telemetry")]
if let Some(recorder) = &mut self.recorder {
recorder.event(&event);
}
#[cfg(feature = "telemetry")]
if let Some(metrics) = &mut self.metrics {
metrics.frame("received", len);
metrics.event(&event);
}
#[cfg(not(feature = "telemetry"))]
let _ = len;
Ok(event)
}
pub(crate) fn pid(&self) -> Option<u32> {
match &self.kind {
WorkerKind::Subprocess(w) => w.child.id(),
WorkerKind::WebSocket(_) => None,
}
}
pub(crate) fn is_dead(&mut self) -> bool {
match &mut self.kind {
WorkerKind::Subprocess(w) => w.child.try_wait().is_ok_and(|status| status.is_some()),
WorkerKind::WebSocket(_) => false,
}
}
pub(crate) async fn reap_or_kill(&mut self, grace: Duration) -> Option<ExitStatus> {
if let WorkerKind::Subprocess(w) = &mut self.kind
&& let Ok(Ok(status)) = timeout(grace, w.child.wait()).await
{
return Some(status);
}
self.kill_and_reap().await
}
pub(crate) async fn kill_and_reap(&mut self) -> Option<ExitStatus> {
match &mut self.kind {
WorkerKind::Subprocess(w) => {
let _ = w.child.start_kill();
w.child.wait().await.ok()
}
WorkerKind::WebSocket(w) => {
w.close().await;
None
}
}
}
pub(crate) async fn close_transport(&mut self) {
if let WorkerKind::WebSocket(w) = &mut self.kind {
w.close().await;
}
}
}
const CLOSE_WRITE_TIMEOUT: Duration = Duration::from_secs(1);
const READ_CHUNK: usize = 8 * 1024;
const RETAIN_BUF_MAX: usize = 64 * 1024;
const SEND_BUF_CAPACITY: usize = 1024;
const DECODE_OFFLOAD_MIN: usize = 1024 * 1024;
struct FrameRecv {
stdout: ChildStdout,
buf: Vec<u8>,
consumed: usize,
filled: usize,
}
impl FrameRecv {
fn new(stdout: ChildStdout) -> Self {
Self {
stdout,
buf: Vec::new(),
consumed: 0,
filled: 0,
}
}
async fn recv(&mut self) -> Result<Option<(pb::ChildEvent, usize)>, FrameError> {
if self.consumed > 0 {
self.buf.copy_within(self.consumed..self.filled, 0);
self.filled -= self.consumed;
self.consumed = 0;
}
loop {
if self.filled >= 4 {
let len = u32::from_le_bytes(self.buf[..4].try_into().expect("4 bytes"));
if len > MAX_FRAME_LEN {
return Err(FrameError::FrameTooLarge {
len,
max: MAX_FRAME_LEN,
});
}
let end = 4 + len as usize;
if self.filled >= end {
let event = decode_event(&self.buf[4..end])?;
self.consumed = end;
if self.buf.len() > RETAIN_BUF_MAX && self.filled == end {
self.buf = Vec::new();
self.consumed = 0;
self.filled = 0;
}
return Ok(Some((event, len as usize)));
}
if self.buf.len() < end {
self.buf.resize(end, 0);
}
} else if self.buf.len() < READ_CHUNK {
self.buf.resize(READ_CHUNK, 0);
}
let n = self.stdout.read(&mut self.buf[self.filled..]).await?;
if n == 0 {
return if self.filled == 0 {
Ok(None)
} else {
Err(FrameError::Truncated)
};
}
self.filled += n;
}
}
}
fn decode_event(body: &[u8]) -> Result<pb::ChildEvent, FrameError> {
if body.len() < DECODE_OFFLOAD_MIN || Handle::current().runtime_flavor() != RuntimeFlavor::MultiThread {
decode_frame(body)
} else {
block_in_place(|| decode_frame(body))
}
}
const WS_EVENT_CHANNEL_DEPTH: usize = 1;
async fn read_ws_events(mut read_half: SplitStream<WsStream>, events: mpsc::Sender<Result<Bytes, FrameError>>) {
while let Ok(permit) = events.reserve().await {
let result = loop {
match read_half.next().await {
Some(Ok(Message::Binary(data))) => break Ok(data),
Some(Ok(Message::Ping(_) | Message::Pong(_))) => {}
Some(Ok(Message::Close(_) | Message::Text(_) | Message::Frame(_))) | None => {
break Err(FrameError::Truncated);
}
Some(Err(WsError::Io(err))) => break Err(FrameError::Io(err)),
Some(Err(_)) => break Err(FrameError::Truncated),
}
};
let stop = result.is_err();
permit.send(result);
if stop {
return;
}
}
}
fn ws_to_frame_error(err: WsError) -> FrameError {
match err {
WsError::Io(err) => FrameError::Io(err),
_ => FrameError::Truncated,
}
}
const USER_AGENT_VALUE: &str = concat!("monty-pool/", env!("CARGO_PKG_VERSION"));
const DEFAULT_DIAL_TIMEOUT: Duration = Duration::from_secs(30);
fn install_crypto_provider() {
static INSTALL: Once = Once::new();
INSTALL.call_once(|| {
let _ = default_provider().install_default();
});
}