feat(websocket): WS upgrade route + channels session (server producer half)
- src/websocket/byte_adapter.rs: production WsByteStream from the POC — inbound bounded mpsc (64 slots, backpressure), outbound chunk parser emitting one WS message per chunk with 1 MiB split; write-side backpressure now uses futures mpsc poll_ready (POC spin-wait fixed); text messages closed with 1002; close mapping per websocket.md - src/websocket/upgrade.rs: /alk/channels upgrade route — bearer auth (401 unresolvable), identity attached to the channels Connection, ChannelsAdapter + install_channel_zero running Dispatcher::run_loop_single_stream - test_support module (feature test-support): WsClient, chunk/frame assemblers; shared with from_wss consumer path (ADR-070) - tests/ws_upgrade_session.rs: 10 integration tests — call round-trip, services/list ACL-filtered, 3 MiB split, interleaved calls, ACL 403, internal-op NOT_FOUND, text->1002 close, disconnect mid-call no-hang Verified: cargo test (95), cargo test --all-features (95+10), clippy -D warnings (default + all-features), fmt.
This commit is contained in:
@@ -0,0 +1,255 @@
|
||||
//! The WS ↔ byte-stream adapter (OQ-01) — the single seam between axum's
|
||||
//! message-oriented `WebSocket` and alkcall's byte-oriented channels
|
||||
//! machinery (`AsyncRead` + `AsyncWrite`).
|
||||
//!
|
||||
//! Validated by the ws-byte-adapter POC (`/workspace/ws-byte-adapter-poc/`,
|
||||
//! OQ-01 GO); this is the production shape.
|
||||
//!
|
||||
//! Inbound: a WS read task pushes binary-message bytes into a bounded
|
||||
//! mpsc (64 slots); the `AsyncRead` half drains it. Backpressure = mpsc
|
||||
//! capacity (OQ-01a): the read task awaits `send` when full.
|
||||
//!
|
||||
//! Outbound: the `AsyncWrite` half queues byte spans; a writer task
|
||||
//! parses the pending bytes for complete chunks (8-byte header → payload
|
||||
//! length) and emits one WS binary message per chunk, splitting chunks
|
||||
//! over the 1 MiB message cap (legal — the receiver's boundary is the
|
||||
//! chunk header, not the message; a chunk may span messages). Chunk
|
||||
//! parsing is required because a logical write above the mux (channel
|
||||
//! 0's `write_frame` issues prefix+body separately) surfaces as multiple
|
||||
//! mux payloads. Write-side backpressure uses `futures::channel::mpsc`
|
||||
//! `poll_ready` — the production fix for the POC's spin-wait.
|
||||
//!
|
||||
//! Text WS messages are rejected with a protocol-level close (code
|
||||
//! 1002); all frames are binary (websocket.md §Framing).
|
||||
//!
|
||||
//! Close mapping: WS close (either side) → read EOF → the demux clears
|
||||
//! all channels (REQ-CH-02) and the dispatch loop fails outstanding
|
||||
//! pendings. `AsyncWrite::shutdown` closes the WS sink after the queued
|
||||
//! bytes drain (the mux's EOF sentinels ride the same queue).
|
||||
//!
|
||||
//! Shared with the `from_wss` consumer path (ADR-070): one
|
||||
//! implementation, both directions.
|
||||
|
||||
use std::{
|
||||
io,
|
||||
pin::Pin,
|
||||
task::{Context, Poll},
|
||||
};
|
||||
|
||||
use axum::extract::ws::{CloseFrame, Message, WebSocket};
|
||||
use futures::channel::mpsc as futures_mpsc;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use tokio::io::{AsyncRead, AsyncWrite};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// Practical per-WS-message cap. Chunks larger than this are split
|
||||
/// across multiple WS messages — legal, since the receiver's boundary
|
||||
/// is the chunk header, not the message. alkcall's MAX_CHUNK_LEN is
|
||||
/// 16 MiB.
|
||||
pub const WS_MESSAGE_CAP: usize = 1024 * 1024;
|
||||
|
||||
/// Inbound buffer: slots × in-flight message bytes. The WS read task
|
||||
/// awaits `send` when full — the backpressure mechanism (OQ-01a).
|
||||
pub const READ_SLOTS: usize = 64;
|
||||
|
||||
const WRITE_SLOTS: usize = 64;
|
||||
|
||||
/// Protocol-error close code for text messages (websocket.md §Framing).
|
||||
pub const WS_PROTOCOL_ERROR: u16 = 1002;
|
||||
|
||||
pub(crate) enum WriteMsg {
|
||||
Bytes(Vec<u8>),
|
||||
/// Close with the given code (e.g. the 1002 text-rejection).
|
||||
CloseWith(u16),
|
||||
}
|
||||
|
||||
/// The `AsyncRead + AsyncWrite` view of a `WebSocket` handed to
|
||||
/// alkcall's channels machinery (demux/mux).
|
||||
pub struct WsByteStream {
|
||||
read_rx: mpsc::Receiver<Vec<u8>>,
|
||||
read_buf: Vec<u8>,
|
||||
read_pos: usize,
|
||||
eof: bool,
|
||||
write_tx: futures_mpsc::Sender<WriteMsg>,
|
||||
write_open: bool,
|
||||
}
|
||||
|
||||
/// The WS pump tasks. Dropping this guard detaches them (tokio
|
||||
/// semantics); they end on their own when the socket halves close.
|
||||
/// `abort()` is available for forced teardown.
|
||||
pub struct WsPumps {
|
||||
read_task: tokio::task::JoinHandle<()>,
|
||||
write_task: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl WsPumps {
|
||||
pub fn abort(&self) {
|
||||
self.read_task.abort();
|
||||
self.write_task.abort();
|
||||
}
|
||||
}
|
||||
|
||||
/// Split a `WebSocket` into the byte stream + the pump tasks. The
|
||||
/// adapter is the single seam between axum's WS and alkcall's
|
||||
/// byte-oriented channels machinery; shared with `from_wss`.
|
||||
pub fn split_ws_to_bytes(socket: WebSocket) -> (WsByteStream, WsPumps) {
|
||||
let (mut ws_sink, mut ws_stream) = socket.split();
|
||||
let (read_tx, read_rx) = mpsc::channel::<Vec<u8>>(READ_SLOTS);
|
||||
let (write_tx, mut write_rx) = futures_mpsc::channel::<WriteMsg>(WRITE_SLOTS);
|
||||
|
||||
let write_tx_for_read = write_tx.clone();
|
||||
let read_task = tokio::spawn(async move {
|
||||
while let Some(msg) = ws_stream.next().await {
|
||||
match msg {
|
||||
Ok(Message::Binary(b)) => {
|
||||
if read_tx.send(b.to_vec()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(Message::Text(_)) => {
|
||||
let _ = write_tx_for_read
|
||||
.clone()
|
||||
.send(WriteMsg::CloseWith(WS_PROTOCOL_ERROR))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
Ok(Message::Close(_)) | Err(_) => break,
|
||||
Ok(_) => {}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let write_task = tokio::spawn(async move {
|
||||
let mut pending: Vec<u8> = Vec::new();
|
||||
while let Some(msg) = write_rx.next().await {
|
||||
match msg {
|
||||
WriteMsg::CloseWith(code) => {
|
||||
let _ = ws_sink
|
||||
.send(Message::Close(Some(CloseFrame {
|
||||
code,
|
||||
reason: "text messages not supported".into(),
|
||||
})))
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
WriteMsg::Bytes(b) => pending.extend_from_slice(&b),
|
||||
}
|
||||
loop {
|
||||
if pending.len() < 8 {
|
||||
break;
|
||||
}
|
||||
let len =
|
||||
u32::from_be_bytes([pending[4], pending[5], pending[6], pending[7]]) as usize;
|
||||
let total = 8 + len;
|
||||
if pending.len() < total {
|
||||
break;
|
||||
}
|
||||
let chunk: Vec<u8> = pending.drain(..total).collect();
|
||||
for piece in chunk.chunks(WS_MESSAGE_CAP) {
|
||||
if ws_sink
|
||||
.send(Message::Binary(piece.to_vec().into()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = ws_sink.close().await;
|
||||
});
|
||||
|
||||
(
|
||||
WsByteStream {
|
||||
read_rx,
|
||||
read_buf: Vec::new(),
|
||||
read_pos: 0,
|
||||
eof: false,
|
||||
write_tx,
|
||||
write_open: true,
|
||||
},
|
||||
WsPumps {
|
||||
read_task,
|
||||
write_task,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
impl AsyncRead for WsByteStream {
|
||||
fn poll_read(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &mut tokio::io::ReadBuf<'_>,
|
||||
) -> Poll<io::Result<()>> {
|
||||
let this = self.get_mut();
|
||||
loop {
|
||||
if this.read_pos < this.read_buf.len() {
|
||||
let n = (this.read_buf.len() - this.read_pos).min(buf.remaining());
|
||||
let end = this.read_pos + n;
|
||||
buf.put_slice(&this.read_buf[this.read_pos..end]);
|
||||
this.read_pos = end;
|
||||
if this.read_pos == this.read_buf.len() {
|
||||
this.read_buf.clear();
|
||||
this.read_pos = 0;
|
||||
}
|
||||
return Poll::Ready(Ok(()));
|
||||
}
|
||||
if this.eof {
|
||||
return Poll::Ready(Ok(()));
|
||||
}
|
||||
match this.read_rx.poll_recv(cx) {
|
||||
Poll::Ready(Some(bytes)) => {
|
||||
this.read_buf = bytes;
|
||||
this.read_pos = 0;
|
||||
}
|
||||
Poll::Ready(None) => {
|
||||
this.eof = true;
|
||||
}
|
||||
Poll::Pending => return Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AsyncWrite for WsByteStream {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<io::Result<usize>> {
|
||||
let this = self.get_mut();
|
||||
if !this.write_open {
|
||||
return Poll::Ready(Err(io::Error::new(
|
||||
io::ErrorKind::BrokenPipe,
|
||||
"ws stream shut down",
|
||||
)));
|
||||
}
|
||||
match this.write_tx.poll_ready(cx) {
|
||||
Poll::Ready(Ok(())) => match this.write_tx.try_send(WriteMsg::Bytes(buf.to_vec())) {
|
||||
Ok(()) => Poll::Ready(Ok(buf.len())),
|
||||
Err(_disconnected_or_full_race) => Poll::Ready(Err(io::Error::new(
|
||||
io::ErrorKind::BrokenPipe,
|
||||
"ws writer closed",
|
||||
))),
|
||||
},
|
||||
Poll::Ready(Err(_send_error)) => Poll::Ready(Err(io::Error::new(
|
||||
io::ErrorKind::BrokenPipe,
|
||||
"ws writer closed",
|
||||
))),
|
||||
Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||
let this = self.get_mut();
|
||||
if this.write_open {
|
||||
this.write_open = false;
|
||||
drop(this.write_tx.clone());
|
||||
}
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user