603 lines
21 KiB
Rust
603 lines
21 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
use iroh::{Endpoint, PublicKey};
|
|
use tokio::io::AsyncWriteExt;
|
|
|
|
use swactor::actor::{ActorAddress, ActorInterface, Ctx, Down};
|
|
use swactor::runtime::Runtime;
|
|
|
|
use crate::streams::data_plane;
|
|
use crate::streams::handle::{create_stream_handle, StreamHandle};
|
|
use crate::streams::messages::{OneShot, StreamManagerMsg, StreamNotification};
|
|
use crate::streams::types::{StreamConfig, StreamError, StreamId, StreamMode};
|
|
use crate::streams::wire;
|
|
|
|
/// Well-known name for the StreamManager actor in the name registry.
|
|
pub const STREAM_MANAGER_NAME: &str = "StreamManager";
|
|
|
|
/// Accept byte sent back on control stream to indicate stream acceptance.
|
|
const ACCEPT_BYTE: u8 = 0x01;
|
|
/// Reject byte sent back on control stream to indicate stream rejection.
|
|
const REJECT_BYTE: u8 = 0x00;
|
|
|
|
struct StreamState {
|
|
owner: ActorAddress,
|
|
}
|
|
|
|
struct PendingIncoming {
|
|
config: StreamConfig,
|
|
conn: OneShot<iroh::endpoint::Connection>,
|
|
}
|
|
|
|
pub struct StreamManager {
|
|
streams: HashMap<StreamId, StreamState>,
|
|
pending_incoming: HashMap<StreamId, PendingIncoming>,
|
|
listeners: HashMap<StreamMode, Vec<ActorAddress>>,
|
|
endpoint: Endpoint,
|
|
tokio_handle: tokio::runtime::Handle,
|
|
runtime: Arc<Runtime>,
|
|
self_addr: Option<ActorAddress>,
|
|
}
|
|
|
|
impl StreamManager {
|
|
pub fn new(
|
|
endpoint: Endpoint,
|
|
tokio_handle: tokio::runtime::Handle,
|
|
runtime: Arc<Runtime>,
|
|
) -> Self {
|
|
StreamManager {
|
|
streams: HashMap::new(),
|
|
pending_incoming: HashMap::new(),
|
|
listeners: HashMap::new(),
|
|
endpoint,
|
|
tokio_handle,
|
|
runtime,
|
|
self_addr: None,
|
|
}
|
|
}
|
|
|
|
fn handle_open(
|
|
&mut self,
|
|
target_node: [u8; 32],
|
|
mode: StreamMode,
|
|
config: StreamConfig,
|
|
reply_to: ActorAddress,
|
|
) {
|
|
let stream_id = StreamId::new_random();
|
|
let endpoint = self.endpoint.clone();
|
|
let runtime = Arc::clone(&self.runtime);
|
|
let self_addr = self.self_addr.expect("StreamManager not started");
|
|
|
|
self.tokio_handle.spawn(async move {
|
|
let result = open_stream_async(endpoint, target_node, stream_id, mode, &config).await;
|
|
let msg = StreamManagerMsg::OpenCompleted {
|
|
stream_id,
|
|
reply_to,
|
|
result: OneShot::new(result),
|
|
};
|
|
let _ = runtime.send_to(self_addr, msg);
|
|
});
|
|
}
|
|
|
|
fn handle_open_completed(
|
|
&mut self,
|
|
ctx: &Ctx,
|
|
stream_id: StreamId,
|
|
reply_to: ActorAddress,
|
|
result: OneShot<Result<StreamHandle, StreamError>>,
|
|
) {
|
|
match result.take() {
|
|
Some(Ok(handle)) => {
|
|
self.streams.insert(
|
|
stream_id,
|
|
StreamState {
|
|
owner: reply_to,
|
|
},
|
|
);
|
|
let notif = StreamNotification::StreamReady {
|
|
stream_id,
|
|
handle: OneShot::new(handle),
|
|
};
|
|
let _ = ctx.send(reply_to, notif);
|
|
}
|
|
Some(Err(err)) => {
|
|
let notif = StreamNotification::StreamFailed {
|
|
stream_id,
|
|
error: err,
|
|
};
|
|
let _ = ctx.send(reply_to, notif);
|
|
}
|
|
None => {
|
|
// OneShot already consumed — should not happen
|
|
eprintln!("StreamManager: OpenCompleted result already consumed");
|
|
}
|
|
}
|
|
}
|
|
|
|
fn handle_incoming_connection(
|
|
&mut self,
|
|
ctx: &Ctx,
|
|
node_id: [u8; 32],
|
|
stream_id: StreamId,
|
|
mode: StreamMode,
|
|
config: StreamConfig,
|
|
conn: OneShot<iroh::endpoint::Connection>,
|
|
) {
|
|
// Notify listeners for this mode
|
|
if let Some(listeners) = self.listeners.get(&mode) {
|
|
let notif = StreamNotification::StreamOffer {
|
|
stream_id,
|
|
mode,
|
|
metadata: config.metadata.clone(),
|
|
from_node: node_id,
|
|
};
|
|
for listener in listeners {
|
|
let _ = ctx.send(*listener, notif.clone());
|
|
}
|
|
}
|
|
|
|
// Store pending incoming for Accept/Reject
|
|
self.pending_incoming.insert(
|
|
stream_id,
|
|
PendingIncoming {
|
|
config,
|
|
conn,
|
|
},
|
|
);
|
|
}
|
|
|
|
fn handle_accept(&mut self, stream_id: StreamId, reply_to: ActorAddress) {
|
|
let pending = match self.pending_incoming.remove(&stream_id) {
|
|
Some(p) => p,
|
|
None => {
|
|
eprintln!("StreamManager: Accept for unknown stream {stream_id}");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let conn = match pending.conn.take() {
|
|
Some(c) => c,
|
|
None => {
|
|
eprintln!("StreamManager: Accept connection already consumed for {stream_id}");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let config = pending.config;
|
|
let runtime = Arc::clone(&self.runtime);
|
|
let self_addr = self.self_addr.expect("StreamManager not started");
|
|
|
|
self.tokio_handle.spawn(async move {
|
|
let result = accept_stream_async(conn, stream_id, &config).await;
|
|
let msg = StreamManagerMsg::AcceptCompleted {
|
|
stream_id,
|
|
reply_to,
|
|
result: OneShot::new(result),
|
|
};
|
|
let _ = runtime.send_to(self_addr, msg);
|
|
});
|
|
}
|
|
|
|
fn handle_accept_completed(
|
|
&mut self,
|
|
ctx: &Ctx,
|
|
stream_id: StreamId,
|
|
reply_to: ActorAddress,
|
|
result: OneShot<Result<StreamHandle, StreamError>>,
|
|
) {
|
|
match result.take() {
|
|
Some(Ok(handle)) => {
|
|
self.streams.insert(
|
|
stream_id,
|
|
StreamState {
|
|
owner: reply_to,
|
|
},
|
|
);
|
|
let notif = StreamNotification::StreamReady {
|
|
stream_id,
|
|
handle: OneShot::new(handle),
|
|
};
|
|
let _ = ctx.send(reply_to, notif);
|
|
}
|
|
Some(Err(err)) => {
|
|
let notif = StreamNotification::StreamFailed {
|
|
stream_id,
|
|
error: err,
|
|
};
|
|
let _ = ctx.send(reply_to, notif);
|
|
}
|
|
None => {
|
|
eprintln!("StreamManager: AcceptCompleted result already consumed");
|
|
}
|
|
}
|
|
}
|
|
|
|
fn handle_reject(&mut self, stream_id: StreamId) {
|
|
if let Some(pending) = self.pending_incoming.remove(&stream_id) {
|
|
// If we have the connection, send reject and close
|
|
if let Some(conn) = pending.conn.take() {
|
|
let runtime = Arc::clone(&self.runtime);
|
|
self.tokio_handle.spawn(async move {
|
|
// Best-effort: send reject on any open bi-stream, then close
|
|
let _ = reject_stream_async(&conn).await;
|
|
drop(conn);
|
|
drop(runtime);
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
fn handle_listen(&mut self, mode: StreamMode, listener: ActorAddress) {
|
|
self.listeners
|
|
.entry(mode)
|
|
.or_default()
|
|
.push(listener);
|
|
}
|
|
|
|
fn handle_close(&mut self, stream_id: StreamId) {
|
|
// Remove stream state; data-plane tasks terminate when channels drop
|
|
self.streams.remove(&stream_id);
|
|
}
|
|
}
|
|
|
|
impl ActorInterface for StreamManager {
|
|
type Incoming = StreamManagerMsg;
|
|
type Response = ();
|
|
|
|
fn on_start(&mut self, ctx: &Ctx) {
|
|
self.self_addr = Some(ctx.self_addr());
|
|
}
|
|
|
|
fn handle(&mut self, ctx: &Ctx, msg: StreamManagerMsg) {
|
|
match msg {
|
|
StreamManagerMsg::Open {
|
|
target_node,
|
|
mode,
|
|
config,
|
|
reply_to,
|
|
} => self.handle_open(target_node, mode, config, reply_to),
|
|
|
|
StreamManagerMsg::Accept {
|
|
stream_id,
|
|
reply_to,
|
|
} => self.handle_accept(stream_id, reply_to),
|
|
|
|
StreamManagerMsg::Reject { stream_id } => self.handle_reject(stream_id),
|
|
|
|
StreamManagerMsg::Listen { mode, listener } => self.handle_listen(mode, listener),
|
|
|
|
StreamManagerMsg::Close { stream_id } => self.handle_close(stream_id),
|
|
|
|
StreamManagerMsg::IncomingConnection {
|
|
node_id,
|
|
stream_id,
|
|
mode,
|
|
config,
|
|
conn,
|
|
} => self.handle_incoming_connection(ctx, node_id, stream_id, mode, config, conn),
|
|
|
|
StreamManagerMsg::OpenCompleted {
|
|
stream_id,
|
|
reply_to,
|
|
result,
|
|
} => self.handle_open_completed(ctx, stream_id, reply_to, result),
|
|
|
|
StreamManagerMsg::AcceptCompleted {
|
|
stream_id,
|
|
reply_to,
|
|
result,
|
|
} => self.handle_accept_completed(ctx, stream_id, reply_to, result),
|
|
}
|
|
}
|
|
|
|
fn handle_down(&mut self, _ctx: &Ctx, down: Down) {
|
|
// Clean up streams owned by the dead actor
|
|
let dead_addr = down.addr;
|
|
self.streams.retain(|_, state| state.owner != dead_addr);
|
|
|
|
// Remove from listeners
|
|
for listeners in self.listeners.values_mut() {
|
|
listeners.retain(|addr| *addr != dead_addr);
|
|
}
|
|
}
|
|
}
|
|
|
|
// ─── Async helpers (run inside tokio tasks) ─────────────────────────────
|
|
|
|
/// Open a stream to a remote node: connect, send header on control bi-stream,
|
|
/// wait for accept/reject, then spawn data-plane tasks.
|
|
async fn open_stream_async(
|
|
endpoint: Endpoint,
|
|
target_node: [u8; 32],
|
|
stream_id: StreamId,
|
|
mode: StreamMode,
|
|
config: &StreamConfig,
|
|
) -> Result<StreamHandle, StreamError> {
|
|
let key = PublicKey::from_bytes(&target_node)
|
|
.map_err(|e| StreamError::BrokenPipe(format!("invalid public key: {e}")))?;
|
|
|
|
let conn = endpoint
|
|
.connect(key, wire::ALPN)
|
|
.await
|
|
.map_err(|e| StreamError::BrokenPipe(format!("connect failed: {e}")))?;
|
|
|
|
// Open control bi-stream and send header
|
|
let (mut send_ctrl, _recv_ctrl) = conn
|
|
.open_bi()
|
|
.await
|
|
.map_err(|e| StreamError::BrokenPipe(format!("open_bi failed: {e}")))?;
|
|
|
|
let header = wire::StreamHeader {
|
|
stream_id,
|
|
mode,
|
|
config: config.clone(),
|
|
};
|
|
let header_bytes = wire::encode_header(&header);
|
|
send_ctrl
|
|
.write_all(&header_bytes)
|
|
.await
|
|
.map_err(|e| StreamError::BrokenPipe(format!("write header failed: {e}")))?;
|
|
send_ctrl
|
|
.finish()
|
|
.map_err(|e| StreamError::BrokenPipe(format!("finish control send failed: {e}")))?;
|
|
|
|
// Wait for accept/reject response on a uni-stream opened by the acceptor.
|
|
// (The bi-stream's send half was dropped by the accept bridge after reading
|
|
// the header, so the acceptor responds via a separate uni-stream.)
|
|
let mut response_recv = conn
|
|
.accept_uni()
|
|
.await
|
|
.map_err(|e| StreamError::BrokenPipe(format!("accept response stream failed: {e}")))?;
|
|
let mut response = [0u8; 1];
|
|
response_recv
|
|
.read_exact(&mut response)
|
|
.await
|
|
.map_err(|e| StreamError::BrokenPipe(format!("read accept/reject failed: {e}")))?;
|
|
|
|
if response[0] != ACCEPT_BYTE {
|
|
return Err(StreamError::BrokenPipe("stream rejected by remote".into()));
|
|
}
|
|
|
|
// Create StreamHandle and spawn data-plane tasks
|
|
let stripe_count = config.stripe_count as usize;
|
|
let (handle, endpoints) = create_stream_handle(stream_id, config, 32, 16);
|
|
|
|
// Spawn send stripe tasks with QUIC uni-streams
|
|
{
|
|
let pool = endpoints.pool.clone();
|
|
let mut cmd_rx = endpoints.send_cmd_rx;
|
|
let evt_tx = endpoints.send_evt_tx;
|
|
let conn_clone = conn.clone();
|
|
let sc = stripe_count;
|
|
|
|
tokio::spawn(async move {
|
|
// Open uni-streams for each stripe
|
|
let mut writers = Vec::with_capacity(sc);
|
|
for _ in 0..sc {
|
|
match conn_clone.open_uni().await {
|
|
Ok(send_stream) => writers.push(send_stream),
|
|
Err(e) => {
|
|
let _ = evt_tx
|
|
.send(crate::streams::channel::SendEvent::Error(StreamError::BrokenPipe(
|
|
format!("open_uni failed: {e}"),
|
|
)))
|
|
.await;
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Simple single-task approach: round-robin commands across stripes
|
|
let mut stripe_idx = 0;
|
|
while let Some(cmd) = cmd_rx.recv().await {
|
|
match cmd {
|
|
crate::streams::channel::SendCommand::Data(buf) => {
|
|
let frame = wire::encode_data_frame(buf.written());
|
|
let writer = &mut writers[stripe_idx];
|
|
if let Err(e) = writer.write_all(&frame).await {
|
|
pool.checkin(buf);
|
|
let _ = evt_tx
|
|
.send(crate::streams::channel::SendEvent::Error(
|
|
StreamError::BrokenPipe(e.to_string()),
|
|
))
|
|
.await;
|
|
return;
|
|
}
|
|
pool.checkin(buf);
|
|
stripe_idx = (stripe_idx + 1) % sc;
|
|
}
|
|
crate::streams::channel::SendCommand::Flush => {
|
|
for writer in &mut writers {
|
|
let _ = writer.flush().await;
|
|
}
|
|
}
|
|
crate::streams::channel::SendCommand::Close => {
|
|
let sentinel = wire::encode_end_of_stripe();
|
|
for writer in &mut writers {
|
|
let _ = writer.write_all(&sentinel).await;
|
|
let _ = writer.finish();
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
// Spawn recv stripe tasks with QUIC uni-streams (accepted from remote)
|
|
{
|
|
let pool = endpoints.pool.clone();
|
|
let evt_tx = endpoints.recv_evt_tx;
|
|
let conn_clone = conn.clone();
|
|
let sc = stripe_count;
|
|
|
|
tokio::spawn(async move {
|
|
// Accept uni-streams for each recv stripe
|
|
let mut closed_count = 0;
|
|
loop {
|
|
match conn_clone.accept_uni().await {
|
|
Ok(recv_stream) => {
|
|
let pool = pool.clone();
|
|
let tx = evt_tx.clone();
|
|
tokio::spawn(async move {
|
|
let _ =
|
|
data_plane::recv_stripe_task(recv_stream, tx, pool, None).await;
|
|
});
|
|
closed_count += 1;
|
|
if closed_count >= sc {
|
|
// We only expect stripe_count recv streams
|
|
// but keep accepting in case more arrive
|
|
}
|
|
}
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
Ok(handle)
|
|
}
|
|
|
|
/// Accept a stream: send accept byte on control stream, spawn data-plane tasks.
|
|
async fn accept_stream_async(
|
|
conn: iroh::endpoint::Connection,
|
|
stream_id: StreamId,
|
|
config: &StreamConfig,
|
|
) -> Result<StreamHandle, StreamError> {
|
|
// Open a uni-stream to send the accept byte back
|
|
// (The opener reads from the recv side of the bi-stream they opened.
|
|
// We need to open our own bi-stream to send the response.)
|
|
// Actually, the opener opened a bi-stream - we need to accept it and
|
|
// respond on it. But IncomingConnection already accepted the bi-stream
|
|
// and read the header. We need the send half of that bi-stream.
|
|
//
|
|
// Since the accept bridge consumed the bi-stream to read the header,
|
|
// we send the accept response on a new uni-stream that the opener
|
|
// will accept_uni on. But the plan says "1-byte accept/reject response"
|
|
// on the same control bi-stream.
|
|
//
|
|
// The design: the accept bridge reads the header from the bi-stream
|
|
// (recv side), and the StreamManager sends accept/reject on the
|
|
// send side. Since the accept bridge consumed the Connection but not
|
|
// the bi-stream send half, we need a different approach.
|
|
//
|
|
// Simpler: use a uni-stream for the response.
|
|
let mut response_stream = conn
|
|
.open_uni()
|
|
.await
|
|
.map_err(|e| StreamError::BrokenPipe(format!("open response stream failed: {e}")))?;
|
|
response_stream
|
|
.write_all(&[ACCEPT_BYTE])
|
|
.await
|
|
.map_err(|e| StreamError::BrokenPipe(format!("write accept byte failed: {e}")))?;
|
|
response_stream
|
|
.finish()
|
|
.map_err(|e| StreamError::BrokenPipe(format!("finish response stream failed: {e}")))?;
|
|
|
|
let stripe_count = config.stripe_count as usize;
|
|
let (handle, endpoints) = create_stream_handle(stream_id, config, 32, 16);
|
|
|
|
// Spawn send stripe tasks — we open uni-streams to write
|
|
{
|
|
let pool = endpoints.pool.clone();
|
|
let mut cmd_rx = endpoints.send_cmd_rx;
|
|
let evt_tx = endpoints.send_evt_tx;
|
|
let conn_clone = conn.clone();
|
|
let sc = stripe_count;
|
|
|
|
tokio::spawn(async move {
|
|
let mut writers = Vec::with_capacity(sc);
|
|
for _ in 0..sc {
|
|
match conn_clone.open_uni().await {
|
|
Ok(send_stream) => writers.push(send_stream),
|
|
Err(e) => {
|
|
let _ = evt_tx
|
|
.send(crate::streams::channel::SendEvent::Error(StreamError::BrokenPipe(
|
|
format!("open_uni failed: {e}"),
|
|
)))
|
|
.await;
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut stripe_idx = 0;
|
|
while let Some(cmd) = cmd_rx.recv().await {
|
|
match cmd {
|
|
crate::streams::channel::SendCommand::Data(buf) => {
|
|
let frame = wire::encode_data_frame(buf.written());
|
|
let writer = &mut writers[stripe_idx];
|
|
if let Err(e) = writer.write_all(&frame).await {
|
|
pool.checkin(buf);
|
|
let _ = evt_tx
|
|
.send(crate::streams::channel::SendEvent::Error(
|
|
StreamError::BrokenPipe(e.to_string()),
|
|
))
|
|
.await;
|
|
return;
|
|
}
|
|
pool.checkin(buf);
|
|
stripe_idx = (stripe_idx + 1) % sc;
|
|
}
|
|
crate::streams::channel::SendCommand::Flush => {
|
|
for writer in &mut writers {
|
|
let _ = writer.flush().await;
|
|
}
|
|
}
|
|
crate::streams::channel::SendCommand::Close => {
|
|
let sentinel = wire::encode_end_of_stripe();
|
|
for writer in &mut writers {
|
|
let _ = writer.write_all(&sentinel).await;
|
|
let _ = writer.finish();
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
// Spawn recv stripe tasks — accept uni-streams from remote
|
|
{
|
|
let pool = endpoints.pool.clone();
|
|
let evt_tx = endpoints.recv_evt_tx;
|
|
let conn_clone = conn.clone();
|
|
|
|
tokio::spawn(async move {
|
|
loop {
|
|
match conn_clone.accept_uni().await {
|
|
Ok(recv_stream) => {
|
|
let pool = pool.clone();
|
|
let tx = evt_tx.clone();
|
|
tokio::spawn(async move {
|
|
let _ =
|
|
data_plane::recv_stripe_task(recv_stream, tx, pool, None).await;
|
|
});
|
|
}
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
Ok(handle)
|
|
}
|
|
|
|
/// Send reject on a connection (best-effort).
|
|
async fn reject_stream_async(
|
|
conn: &iroh::endpoint::Connection,
|
|
) -> Result<(), StreamError> {
|
|
let mut response_stream = conn
|
|
.open_uni()
|
|
.await
|
|
.map_err(|e| StreamError::BrokenPipe(format!("open response stream failed: {e}")))?;
|
|
response_stream
|
|
.write_all(&[REJECT_BYTE])
|
|
.await
|
|
.map_err(|e| StreamError::BrokenPipe(format!("write reject byte failed: {e}")))?;
|
|
response_stream
|
|
.finish()
|
|
.map_err(|e| StreamError::BrokenPipe(format!("finish response stream failed: {e}")))?;
|
|
Ok(())
|
|
}
|