From eea507e389efcfc3db854b90552f9c423b31cf93 Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Mon, 23 Feb 2026 11:46:35 +0700 Subject: [PATCH] feat: data streams primitive Allows streaming blobs without interference from the actor runtime. --- Cargo.lock | 21 + Cargo.toml | 1 + crates/dashboard/src/bin/swactor-node.rs | 1 + crates/datastore/Cargo.toml | 7 + crates/datastore/src/actors/datastore_node.rs | 184 +++- crates/datastore/src/actors/mod.rs | 3 + .../datastore/src/actors/stream_downloader.rs | 171 ++++ .../datastore/src/actors/stream_listener.rs | 75 ++ crates/datastore/src/actors/stream_server.rs | 151 ++++ crates/datastore/src/blob_transfer.rs | 787 ++++++++++++++++++ crates/datastore/src/bridge.rs | 26 + crates/datastore/src/lib.rs | 1 + crates/datastore/src/messages.rs | 63 +- crates/datastore/tests/stream_integration.rs | 415 +++++++++ crates/distribution/src/iroh_driver.rs | 47 +- crates/distribution/tests/common/iroh.rs | 3 + crates/streams/Cargo.toml | 20 + crates/streams/src/accept.rs | 74 ++ crates/streams/src/buffer.rs | 252 ++++++ crates/streams/src/channel.rs | 41 + crates/streams/src/connection.rs | 65 ++ crates/streams/src/ctx_ext.rs | 201 +++++ crates/streams/src/data_plane.rs | 499 +++++++++++ crates/streams/src/handle.rs | 248 ++++++ crates/streams/src/lib.rs | 23 + crates/streams/src/manager.rs | 622 ++++++++++++++ crates/streams/src/messages.rs | 176 ++++ crates/streams/src/notify.rs | 172 ++++ crates/streams/src/types.rs | 148 ++++ crates/streams/src/wire.rs | 306 +++++++ crates/swactor-node/Cargo.toml | 1 + crates/swactor-node/src/main.rs | 41 + docs/development_history/STREAMS.md | 475 +++++++++++ .../STREAMS_IMPLEMENTATION.md | 338 ++++++++ 34 files changed, 5647 insertions(+), 11 deletions(-) create mode 100644 crates/datastore/src/actors/stream_downloader.rs create mode 100644 crates/datastore/src/actors/stream_listener.rs create mode 100644 crates/datastore/src/actors/stream_server.rs create mode 100644 crates/datastore/src/blob_transfer.rs create mode 100644 crates/datastore/tests/stream_integration.rs create mode 100644 crates/streams/Cargo.toml create mode 100644 crates/streams/src/accept.rs create mode 100644 crates/streams/src/buffer.rs create mode 100644 crates/streams/src/channel.rs create mode 100644 crates/streams/src/connection.rs create mode 100644 crates/streams/src/ctx_ext.rs create mode 100644 crates/streams/src/data_plane.rs create mode 100644 crates/streams/src/handle.rs create mode 100644 crates/streams/src/lib.rs create mode 100644 crates/streams/src/manager.rs create mode 100644 crates/streams/src/messages.rs create mode 100644 crates/streams/src/notify.rs create mode 100644 crates/streams/src/types.rs create mode 100644 crates/streams/src/wire.rs create mode 100644 docs/development_history/STREAMS.md create mode 100644 docs/development_history/STREAMS_IMPLEMENTATION.md diff --git a/Cargo.lock b/Cargo.lock index 8bf8325..b61e2de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5693,14 +5693,17 @@ dependencies = [ "dashboard", "distribution", "getrandom 0.2.17", + "iroh", "proptest", "serde", "serde_json", "shared-types", "swactor", "swactor-std", + "swactor-streams", "tempfile", "tiny_http", + "tokio", "ureq", ] @@ -5718,6 +5721,7 @@ dependencies = [ "swactor", "swactor-datastore", "swactor-std", + "swactor-streams", "toml 0.8.23", ] @@ -5744,6 +5748,23 @@ dependencies = [ "swactor", ] +[[package]] +name = "swactor-streams" +version = "0.1.0" +dependencies = [ + "blake3", + "crossbeam-queue", + "distribution", + "getrandom 0.2.17", + "iroh", + "proptest", + "serde", + "shared-types", + "swactor", + "swactor-std", + "tokio", +] + [[package]] name = "syn" version = "2.0.116" diff --git a/Cargo.toml b/Cargo.toml index 391d6f7..1f5565d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ members = [ "crates/datastore", "crates/shared-types", "crates/swactor-node", + "crates/streams", "tests/docker", "crates/ci", "xtask", diff --git a/crates/dashboard/src/bin/swactor-node.rs b/crates/dashboard/src/bin/swactor-node.rs index d2cd4f0..e06df89 100644 --- a/crates/dashboard/src/bin/swactor-node.rs +++ b/crates/dashboard/src/bin/swactor-node.rs @@ -221,6 +221,7 @@ fn run_iroh( relay_mode: RelayMode::Default, node: node_config, peer_auth: None, + additional_alpns: vec![], }; let mut driver = IrohDriver::new(iroh_config).expect("failed to create iroh driver"); diff --git a/crates/datastore/Cargo.toml b/crates/datastore/Cargo.toml index fcd9323..bda53e1 100644 --- a/crates/datastore/Cargo.toml +++ b/crates/datastore/Cargo.toml @@ -16,12 +16,19 @@ ureq = { version = "2", features = ["json"], optional = true } getrandom = { version = "0.2", optional = true } dashboard = { path = "../dashboard" } swactor-std = { path = "../std" } +swactor-streams = { path = "../streams" } +tokio = { version = "1", features = ["sync", "rt", "time"] } [dev-dependencies] serde_json = "1" proptest = "1" tempfile = "3" swactor = { path = "../.." } +swactor-streams = { path = "../streams" } +swactor-std = { path = "../std" } +distribution = { path = "../distribution", features = ["iroh"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "io-util", "time"] } +iroh = "0.96" ureq = { version = "2", features = ["json"] } [features] diff --git a/crates/datastore/src/actors/datastore_node.rs b/crates/datastore/src/actors/datastore_node.rs index 3970e0e..a358884 100644 --- a/crates/datastore/src/actors/datastore_node.rs +++ b/crates/datastore/src/actors/datastore_node.rs @@ -5,17 +5,27 @@ //! Delete, List, Status) and receive responses. Also routes incoming network //! protocol messages to the appropriate internal actors. -use std::collections::BTreeMap; +use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; use swactor::actor::{ActorAddress, ActorInterface, Ctx}; +use swactor::runtime::Runtime; use distribution::types::NodeId; +use crate::actors::stream_downloader::StreamDownloader; +use crate::actors::stream_server::StreamServer; use crate::chunking::chunk_blob; use crate::messages::{ BlobStoreMsg, DatastoreNodeMsg, DatastoreResponse, MetadataMsg, }; -use crate::types::{ContentHash, DatastoreConfig, ObjectEntry}; +use crate::types::{ContentHash, DatastoreConfig, ObjectEntry, ObjectManifest}; + +/// Progress from a partially-completed stream download, used for resume. +struct PartialDownload { + _source_node: [u8; 32], + chunks_completed: u64, +} /// Top-level coordinator actor for the datastore. /// @@ -26,6 +36,12 @@ pub struct DatastoreNode { blob_store: ActorAddress, metadata: ActorAddress, config: DatastoreConfig, + // Stream support (configured lazily via ConfigureStreams) + runtime: Option>, + tokio_handle: Option, + stream_manager: Option, + // Resume state for interrupted stream downloads + partial_downloads: HashMap, } impl DatastoreNode { @@ -40,6 +56,10 @@ impl DatastoreNode { blob_store, metadata, config, + runtime: None, + tokio_handle: None, + stream_manager: None, + partial_downloads: HashMap::new(), } } @@ -231,6 +251,138 @@ impl DatastoreNode { }, ); } + + // ── Stream-based transfer handlers ─────────────────────────────────── + + fn handle_configure_streams( + &mut self, + stream_manager: ActorAddress, + tokio_handle: tokio::runtime::Handle, + runtime: Arc, + ) { + self.stream_manager = Some(stream_manager); + self.tokio_handle = Some(tokio_handle); + self.runtime = Some(runtime); + } + + fn handle_download_via_stream( + &self, + ctx: &Ctx, + content_hash: ContentHash, + source_node: [u8; 32], + reply_to: ActorAddress, + ) { + let (stream_manager, tokio_handle, runtime) = + match (&self.stream_manager, &self.tokio_handle, &self.runtime) { + (Some(sm), Some(th), Some(rt)) => (*sm, th.clone(), Arc::clone(rt)), + _ => { + let _ = ctx.send( + reply_to, + DatastoreResponse::TransferFailed { + reason: "stream support not configured".into(), + }, + ); + return; + } + }; + + // Check for partial progress from a previous attempt + let skip_chunks = self + .partial_downloads + .get(&content_hash) + .map(|p| p.chunks_completed) + .unwrap_or(0); + + let downloader = StreamDownloader::new( + content_hash, + source_node, + ctx.self_addr(), + self.blob_store, + reply_to, + stream_manager, + tokio_handle, + runtime, + skip_chunks, + ); + let _ = ctx.spawn(downloader); + } + + fn handle_stream_offer( + &self, + ctx: &Ctx, + stream_id: swactor_streams::types::StreamId, + content_hash: ContentHash, + _from_node: [u8; 32], + stream_manager: ActorAddress, + resume_from_chunk: u64, + ) { + let (tokio_handle, runtime) = match (&self.tokio_handle, &self.runtime) { + (Some(th), Some(rt)) => (th.clone(), Arc::clone(rt)), + _ => return, + }; + + let server = StreamServer::new( + stream_id, + content_hash, + self.blob_store, + stream_manager, + tokio_handle, + runtime, + resume_from_chunk, + ); + let _ = ctx.spawn(server); + } + + fn handle_stream_download_complete( + &mut self, + ctx: &Ctx, + content_hash: ContentHash, + manifest: ObjectManifest, + reply_to: ActorAddress, + ) { + // Clear any partial progress now that download is complete + self.partial_downloads.remove(&content_hash); + + let entry = ObjectEntry { + content_hash, + name: None, + node_id: self.node_id, + tags: BTreeMap::new(), + size_bytes: manifest.total_size, + created_at: 0, + }; + + let _ = ctx.send( + self.metadata, + MetadataMsg::PutObject { + entry, + manifest, + reply_to, + }, + ); + } + + fn handle_stream_download_failed( + &mut self, + ctx: &Ctx, + content_hash: ContentHash, + reason: String, + chunks_completed: u64, + source_node: [u8; 32], + reply_to: ActorAddress, + ) { + // Store partial progress so next attempt can resume + if chunks_completed > 0 { + self.partial_downloads.insert( + content_hash, + PartialDownload { + _source_node: source_node, + chunks_completed, + }, + ); + } + let _ = ctx.send(reply_to, DatastoreResponse::TransferFailed { reason }); + } } impl ActorInterface for DatastoreNode { @@ -277,6 +429,34 @@ impl ActorInterface for DatastoreNode { DatastoreNodeMsg::IncomingListObjects { request, reply_to } => { self.handle_incoming_list_objects(ctx, request, reply_to) } + DatastoreNodeMsg::DownloadViaStream { + content_hash, + source_node, + reply_to, + } => self.handle_download_via_stream(ctx, content_hash, source_node, reply_to), + DatastoreNodeMsg::HandleStreamOffer { + stream_id, + content_hash, + from_node, + stream_manager, + resume_from_chunk, + } => self.handle_stream_offer(ctx, stream_id, content_hash, from_node, stream_manager, resume_from_chunk), + DatastoreNodeMsg::StreamDownloadComplete { + content_hash, + manifest, + reply_to, + } => self.handle_stream_download_complete(ctx, content_hash, manifest, reply_to), + DatastoreNodeMsg::StreamDownloadFailed { + content_hash, + reason, + chunks_completed, + reply_to, + } => self.handle_stream_download_failed(ctx, content_hash, reason, chunks_completed, [0; 32], reply_to), + DatastoreNodeMsg::ConfigureStreams { + stream_manager, + tokio_handle, + runtime, + } => self.handle_configure_streams(stream_manager, tokio_handle, runtime), } } } diff --git a/crates/datastore/src/actors/mod.rs b/crates/datastore/src/actors/mod.rs index c6e185d..6f874cd 100644 --- a/crates/datastore/src/actors/mod.rs +++ b/crates/datastore/src/actors/mod.rs @@ -2,6 +2,9 @@ pub mod blob_store; pub mod datastore_node; pub mod gateway; pub mod metadata; +pub mod stream_downloader; +pub mod stream_listener; +pub mod stream_server; pub mod transfer; pub use blob_store::BlobStoreActor; diff --git a/crates/datastore/src/actors/stream_downloader.rs b/crates/datastore/src/actors/stream_downloader.rs new file mode 100644 index 0000000..2fc2459 --- /dev/null +++ b/crates/datastore/src/actors/stream_downloader.rs @@ -0,0 +1,171 @@ +//! StreamDownloader — opens a stream to a remote node and downloads a blob. +//! +//! Lifecycle: +//! 1. on_start: sends Open to StreamManager +//! 2. StreamReady: spawns a tokio task for I/O, then stops self +//! 3. tokio task: recv_blob, write chunks to BlobStore, notify DatastoreNode + +use std::sync::Arc; + +use swactor::actor::{ActorAddress, ActorInterface, Ctx}; +use swactor::runtime::Runtime; + +use swactor_streams::messages::{StreamManagerMsg, StreamNotification}; +use swactor_streams::types::{StreamConfig, StreamMode}; + +use crate::blob_transfer::{encode_metadata, recv_blob, BlobTransferMetadata}; +use crate::messages::{BlobStoreMsg, DatastoreNodeMsg}; +use crate::types::ContentHash; + +pub struct StreamDownloader { + content_hash: ContentHash, + source_node: [u8; 32], + datastore_node: ActorAddress, + blob_store: ActorAddress, + reply_to: ActorAddress, + stream_manager: ActorAddress, + tokio_handle: tokio::runtime::Handle, + runtime: Arc, + skip_chunks: u64, +} + +impl StreamDownloader { + #[allow(clippy::too_many_arguments)] + pub fn new( + content_hash: ContentHash, + source_node: [u8; 32], + datastore_node: ActorAddress, + blob_store: ActorAddress, + reply_to: ActorAddress, + stream_manager: ActorAddress, + tokio_handle: tokio::runtime::Handle, + runtime: Arc, + skip_chunks: u64, + ) -> Self { + Self { + content_hash, + source_node, + datastore_node, + blob_store, + reply_to, + stream_manager, + tokio_handle, + runtime, + skip_chunks, + } + } +} + +impl ActorInterface for StreamDownloader { + type Incoming = StreamNotification; + type Response = (); + + fn on_start(&mut self, ctx: &Ctx) { + let resume = if self.skip_chunks > 0 { + Some(self.skip_chunks) + } else { + None + }; + let meta = BlobTransferMetadata { + content_hash: self.content_hash, + resume_from_chunk: resume, + }; + let config = StreamConfig { + metadata: encode_metadata(&meta), + stripe_count: 1, // blob transfer is sequential — one stripe avoids empty-stripe Closed races + ..Default::default() + }; + let _ = ctx.send( + self.stream_manager, + StreamManagerMsg::Open { + target_node: self.source_node, + mode: StreamMode::BlobTransfer, + config, + reply_to: ctx.self_addr(), + }, + ); + } + + fn handle(&mut self, ctx: &Ctx, msg: StreamNotification) { + match msg { + StreamNotification::StreamReady { handle, .. } => { + let stream_handle = match handle.take() { + Some(h) => h, + None => return, + }; + + let runtime = Arc::clone(&self.runtime); + let blob_store = self.blob_store; + let datastore_node = self.datastore_node; + let reply_to = self.reply_to; + let content_hash = self.content_hash; + let skip_chunks = self.skip_chunks; + + self.tokio_handle.spawn(async move { + let (_send, mut recv) = (stream_handle.send, stream_handle.recv); + + match recv_blob(&mut recv, skip_chunks).await { + Ok(received) => { + // Write chunks to BlobStore (fire-and-forget) + for (hash, data) in &received.chunks { + let _ = runtime.send_to( + blob_store, + BlobStoreMsg::WriteChunk { + hash: *hash, + data: data.clone(), + reply_to: datastore_node, // response ignored + }, + ); + } + + // Write manifest to BlobStore (fire-and-forget) + let _ = runtime.send_to( + blob_store, + BlobStoreMsg::WriteManifest { + manifest: received.manifest.clone(), + reply_to: datastore_node, // response ignored + }, + ); + + // Notify DatastoreNode of completion + let _ = runtime.send_to( + datastore_node, + DatastoreNodeMsg::StreamDownloadComplete { + content_hash, + manifest: received.manifest, + reply_to, + }, + ); + } + Err(e) => { + let _ = runtime.send_to( + datastore_node, + DatastoreNodeMsg::StreamDownloadFailed { + content_hash, + reason: e.to_string(), + chunks_completed: skip_chunks, + reply_to, + }, + ); + } + } + }); + + ctx.stop_self(); + } + StreamNotification::StreamFailed { error, .. } => { + let _ = ctx.send( + self.datastore_node, + DatastoreNodeMsg::StreamDownloadFailed { + content_hash: self.content_hash, + reason: error.to_string(), + chunks_completed: self.skip_chunks, + reply_to: self.reply_to, + }, + ); + ctx.stop_self(); + } + _ => {} + } + } +} diff --git a/crates/datastore/src/actors/stream_listener.rs b/crates/datastore/src/actors/stream_listener.rs new file mode 100644 index 0000000..326ade6 --- /dev/null +++ b/crates/datastore/src/actors/stream_listener.rs @@ -0,0 +1,75 @@ +//! StreamListener — listens for incoming BlobTransfer stream offers and +//! forwards them to DatastoreNode for handling. + +use swactor::actor::{ActorAddress, ActorInterface, Ctx}; + +use swactor_streams::messages::{StreamManagerMsg, StreamNotification}; +use swactor_streams::types::StreamMode; + +use crate::blob_transfer::parse_metadata; +use crate::messages::DatastoreNodeMsg; + +pub struct StreamListener { + datastore_node: ActorAddress, + stream_manager: ActorAddress, +} + +impl StreamListener { + pub fn new(datastore_node: ActorAddress, stream_manager: ActorAddress) -> Self { + Self { + datastore_node, + stream_manager, + } + } +} + +impl ActorInterface for StreamListener { + type Incoming = StreamNotification; + type Response = (); + + fn on_start(&mut self, ctx: &Ctx) { + let _ = ctx.send( + self.stream_manager, + StreamManagerMsg::Listen { + mode: StreamMode::BlobTransfer, + listener: ctx.self_addr(), + }, + ); + } + + fn handle(&mut self, ctx: &Ctx, msg: StreamNotification) { + match msg { + StreamNotification::StreamOffer { + stream_id, + metadata, + from_node, + .. + } => { + // Parse metadata (supports both legacy 32-byte and new versioned format) + let meta = match parse_metadata(&metadata) { + Some(m) => m, + None => { + // Reject malformed offer + let _ = ctx.send(self.stream_manager, StreamManagerMsg::Reject { stream_id }); + return; + } + }; + + let stream_manager = self.stream_manager; + + let _ = ctx.send( + self.datastore_node, + DatastoreNodeMsg::HandleStreamOffer { + stream_id, + content_hash: meta.content_hash, + from_node, + stream_manager, + resume_from_chunk: meta.resume_from_chunk.unwrap_or(0), + }, + ); + } + // Ignore other notifications + _ => {} + } + } +} diff --git a/crates/datastore/src/actors/stream_server.rs b/crates/datastore/src/actors/stream_server.rs new file mode 100644 index 0000000..5575413 --- /dev/null +++ b/crates/datastore/src/actors/stream_server.rs @@ -0,0 +1,151 @@ +//! StreamServer — accepts an incoming stream and serves blob data. +//! +//! Lifecycle: +//! 1. on_start: sends Accept to StreamManager +//! 2. StreamReady: spawns a tokio task for I/O, then stops self +//! 3. tokio task: reads manifest from BlobStore, streams each chunk on-demand + +use std::sync::Arc; +use std::time::Duration; + +use swactor::actor::{ActorAddress, ActorInterface, Ctx}; +use swactor::runtime::Runtime; + +use swactor_streams::messages::{StreamManagerMsg, StreamNotification}; +use swactor_streams::types::StreamId; + +use crate::blob_transfer::{poll_inbox, send_blob, BlobTransferError}; +use crate::messages::{BlobStoreMsg, DatastoreResponse}; +use crate::types::ContentHash; + +const INBOX_TIMEOUT: Duration = Duration::from_secs(10); + +pub struct StreamServer { + stream_id: StreamId, + content_hash: ContentHash, + blob_store: ActorAddress, + stream_manager: ActorAddress, + tokio_handle: tokio::runtime::Handle, + runtime: Arc, + skip_chunks: u64, +} + +impl StreamServer { + pub fn new( + stream_id: StreamId, + content_hash: ContentHash, + blob_store: ActorAddress, + stream_manager: ActorAddress, + tokio_handle: tokio::runtime::Handle, + runtime: Arc, + skip_chunks: u64, + ) -> Self { + Self { + stream_id, + content_hash, + blob_store, + stream_manager, + tokio_handle, + runtime, + skip_chunks, + } + } +} + +impl ActorInterface for StreamServer { + type Incoming = StreamNotification; + type Response = (); + + fn on_start(&mut self, ctx: &Ctx) { + let _ = ctx.send( + self.stream_manager, + StreamManagerMsg::Accept { + stream_id: self.stream_id, + reply_to: ctx.self_addr(), + }, + ); + } + + fn handle(&mut self, ctx: &Ctx, msg: StreamNotification) { + match msg { + StreamNotification::StreamReady { handle, .. } => { + let stream_handle = match handle.take() { + Some(h) => h, + None => return, + }; + + let runtime = Arc::clone(&self.runtime); + let blob_store = self.blob_store; + let content_hash = self.content_hash; + let skip_chunks = self.skip_chunks; + + self.tokio_handle.spawn(async move { + let (mut send, _recv) = (stream_handle.send, stream_handle.recv); + + // Read manifest from BlobStore + let manifest_inbox = match runtime.new_inbox::() { + Ok(inbox) => inbox, + Err(_) => return, + }; + let _ = runtime.send_to( + blob_store, + BlobStoreMsg::ReadManifest { + hash: content_hash, + reply_to: *manifest_inbox.addr(), + }, + ); + + let manifest = match poll_inbox(&manifest_inbox, INBOX_TIMEOUT).await { + Some(DatastoreResponse::ManifestOk { manifest }) => manifest, + other => { + // Manifest not found or timeout — close stream + eprintln!("StreamServer: manifest read failed: {other:?}"); + let _ = send.close(); + return; + } + }; + + // Stream each chunk on-demand (one at a time) + let rt = Arc::clone(&runtime); + let bs = blob_store; + let result = send_blob(&mut send, &manifest, |chunk_hash| { + let rt = Arc::clone(&rt); + async move { + let chunk_inbox = rt + .new_inbox::() + .map_err(|e| { + BlobTransferError::Storage(format!( + "failed to create inbox: {e}" + )) + })?; + let _ = rt.send_to( + bs, + BlobStoreMsg::ReadChunk { + hash: chunk_hash, + reply_to: *chunk_inbox.addr(), + }, + ); + match poll_inbox(&chunk_inbox, INBOX_TIMEOUT).await { + Some(DatastoreResponse::ChunkOk { data, .. }) => Ok(data), + _ => Err(BlobTransferError::Storage( + "chunk not found or timeout".into(), + )), + } + } + }, skip_chunks) + .await; + + if let Err(e) = result { + eprintln!("StreamServer: send_blob failed: {e}"); + } + }); + + ctx.stop_self(); + } + StreamNotification::StreamFailed { .. } => { + ctx.stop_self(); + } + _ => {} + } + } +} diff --git a/crates/datastore/src/blob_transfer.rs b/crates/datastore/src/blob_transfer.rs new file mode 100644 index 0000000..4cf8e59 --- /dev/null +++ b/crates/datastore/src/blob_transfer.rs @@ -0,0 +1,787 @@ +//! BlobTransfer wire protocol — async functions for sending and receiving +//! content-addressed blobs over a QUIC stream (SendHalf / RecvHalf). +//! +//! Wire format: +//! ```text +//! [4B manifest_json_length (u32 BE)] +//! [N bytes manifest JSON] +//! [chunk_0 raw bytes] ← size from manifest.chunks[0].size +//! [chunk_1 raw bytes] +//! ... +//! ``` +//! +//! These run inside tokio tasks (NOT actor handlers). + +use std::future::Future; +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use swactor::runtime::Inbox; +use swactor::actor::Message; +use swactor_streams::handle::{RecvHalf, SendHalf}; + +use crate::types::{ContentHash, ObjectManifest}; + +// ─── Metadata encoding ────────────────────────────────────────────────── + +/// Version tag for new-format metadata (byte 0). +const METADATA_VERSION_1: u8 = 0x01; + +/// Structured metadata sent in `StreamConfig.metadata` for blob transfers. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct BlobTransferMetadata { + pub content_hash: ContentHash, + #[serde(default)] + pub resume_from_chunk: Option, +} + +/// Encode metadata into the wire format: `[0x01][JSON bytes]`. +pub fn encode_metadata(meta: &BlobTransferMetadata) -> Vec { + let json = serde_json::to_vec(meta).expect("BlobTransferMetadata serialization cannot fail"); + let mut buf = Vec::with_capacity(1 + json.len()); + buf.push(METADATA_VERSION_1); + buf.extend_from_slice(&json); + buf +} + +/// Parse metadata from either the legacy 32-byte format or the new versioned format. +pub fn parse_metadata(data: &[u8]) -> Option { + if data.len() == 32 { + // Legacy format: raw 32-byte ContentHash + let mut hash_bytes = [0u8; 32]; + hash_bytes.copy_from_slice(data); + return Some(BlobTransferMetadata { + content_hash: ContentHash(hash_bytes), + resume_from_chunk: None, + }); + } + if data.len() > 1 && data[0] == METADATA_VERSION_1 { + return serde_json::from_slice(&data[1..]).ok(); + } + None +} + +// ─── Error type ────────────────────────────────────────────────────────── + +#[derive(Debug, Clone)] +pub enum BlobTransferError { + /// The stream was closed or disconnected before the transfer completed. + IncompleteTransfer(String), + /// A chunk failed blake3 verification. + ChunkVerificationFailed { + index: usize, + expected: ContentHash, + actual: ContentHash, + }, + /// Manifest JSON could not be parsed. + InvalidManifest(String), + /// An error from the underlying storage layer. + Storage(String), +} + +impl std::fmt::Display for BlobTransferError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + BlobTransferError::IncompleteTransfer(msg) => { + write!(f, "incomplete transfer: {msg}") + } + BlobTransferError::ChunkVerificationFailed { + index, + expected, + actual, + } => write!( + f, + "chunk {index} verification failed: expected {expected}, got {actual}" + ), + BlobTransferError::InvalidManifest(msg) => { + write!(f, "invalid manifest: {msg}") + } + BlobTransferError::Storage(msg) => write!(f, "storage error: {msg}"), + } + } +} + +impl std::error::Error for BlobTransferError {} + +// ─── ReceivedBlob ──────────────────────────────────────────────────────── + +/// Result of a successful `recv_blob` call. +#[derive(Debug)] +pub struct ReceivedBlob { + pub manifest: ObjectManifest, + pub chunks: Vec<(ContentHash, Vec)>, +} + +// ─── Core protocol ─────────────────────────────────────────────────────── + +/// Send a blob over a stream. Chunks are read on-demand via `read_chunk`. +/// +/// `read_chunk` is called once per chunk — at most one chunk is in memory +/// at a time on the sender side. +/// +/// `skip_chunks` allows resuming a previous transfer: the first `skip_chunks` +/// chunks are not read or written. The manifest preamble is always sent so the +/// receiver can verify integrity. +pub async fn send_blob( + send: &mut SendHalf, + manifest: &ObjectManifest, + read_chunk: F, + skip_chunks: u64, +) -> Result<(), BlobTransferError> +where + F: Fn(ContentHash) -> Fut, + Fut: Future, BlobTransferError>>, +{ + // Serialize manifest + let manifest_json = serde_json::to_vec(manifest) + .map_err(|e| BlobTransferError::InvalidManifest(e.to_string()))?; + + // Write manifest preamble: [4B length BE] [manifest JSON] + let len_bytes = (manifest_json.len() as u32).to_be_bytes(); + write_all(send, &len_bytes).await?; + write_all(send, &manifest_json).await?; + + // Write chunks, skipping already-transferred ones + for (i, chunk_ref) in manifest.chunks.iter().enumerate() { + if (i as u64) < skip_chunks { + continue; + } + let data = read_chunk(chunk_ref.hash).await?; + write_all(send, &data).await?; + } + + // Flush and close + send.flush() + .map_err(|e| BlobTransferError::IncompleteTransfer(e.to_string()))?; + send.close() + .map_err(|e| BlobTransferError::IncompleteTransfer(e.to_string()))?; + + Ok(()) +} + +/// Receive a blob from a stream. Reads manifest, then reads and verifies +/// each chunk via blake3. +/// +/// `skip_chunks` allows resuming: the sender skipped the first `skip_chunks` +/// chunks, so the receiver only reads chunks from `skip_chunks` onward. +/// The manifest preamble is always read. +pub async fn recv_blob( + recv: &mut RecvHalf, + skip_chunks: u64, +) -> Result { + // Read manifest preamble + let mut len_buf = [0u8; 4]; + read_exact(recv, &mut len_buf).await?; + let manifest_len = u32::from_be_bytes(len_buf) as usize; + + // Read manifest JSON + let mut manifest_buf = vec![0u8; manifest_len]; + read_exact(recv, &mut manifest_buf).await?; + let manifest: ObjectManifest = serde_json::from_slice(&manifest_buf) + .map_err(|e| BlobTransferError::InvalidManifest(e.to_string()))?; + + // Read and verify each chunk (only those the sender actually sent) + let total = manifest.chunks.len(); + let start = (skip_chunks as usize).min(total); + let mut chunks = Vec::with_capacity(total - start); + for (i, chunk_ref) in manifest.chunks.iter().enumerate().skip(start) { + let mut chunk_data = vec![0u8; chunk_ref.size as usize]; + read_exact(recv, &mut chunk_data).await?; + + // Verify blake3 + let actual_hash = ContentHash::of(&chunk_data); + if actual_hash != chunk_ref.hash { + return Err(BlobTransferError::ChunkVerificationFailed { + index: i, + expected: chunk_ref.hash, + actual: actual_hash, + }); + } + + chunks.push((chunk_ref.hash, chunk_data)); + } + + Ok(ReceivedBlob { manifest, chunks }) +} + +// ─── Helpers ───────────────────────────────────────────────────────────── + +/// Write all bytes to a SendHalf, yielding when backpressured. +async fn write_all(send: &mut SendHalf, data: &[u8]) -> Result<(), BlobTransferError> { + let mut offset = 0; + while offset < data.len() { + match send.try_write(&data[offset..]) { + Ok(0) => { + // Backpressure — yield and retry + tokio::task::yield_now().await; + } + Ok(n) => { + offset += n; + } + Err(e) => { + return Err(BlobTransferError::IncompleteTransfer(e.to_string())); + } + } + } + Ok(()) +} + +/// Read exactly `buf.len()` bytes from a RecvHalf, yielding when no data. +async fn read_exact(recv: &mut RecvHalf, buf: &mut [u8]) -> Result<(), BlobTransferError> { + let mut offset = 0; + while offset < buf.len() { + match recv.try_read(&mut buf[offset..]) { + Ok(0) => { + // No data available — yield and retry + tokio::task::yield_now().await; + } + Ok(n) => { + offset += n; + } + Err(e) => { + return Err(BlobTransferError::IncompleteTransfer(e.to_string())); + } + } + } + Ok(()) +} + +/// Async version of `bridge.rs:poll_response` — yields instead of thread::sleep. +pub async fn poll_inbox(inbox: &Inbox, timeout: Duration) -> Option { + let start = tokio::time::Instant::now(); + loop { + if let Some(msg) = inbox.try_recv() { + return Some(msg); + } + if start.elapsed() > timeout { + return None; + } + tokio::task::yield_now().await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::chunking::chunk_blob; + use swactor_streams::handle::create_stream_handle; + use swactor_streams::types::{StreamConfig, StreamId}; + + /// Helper: create a pair of (SendHalf, RecvHalf) connected via tokio tasks + /// that relay data through a DuplexStream. + fn create_test_pair() -> (SendHalf, RecvHalf) { + let stream_id = StreamId::new_random(); + let config = StreamConfig { + stripe_count: 1, + frame_size: 256 * 1024, + metadata: Vec::new(), + }; + let (handle_a, endpoints_a) = create_stream_handle(stream_id, &config, 128, 64); + let (handle_b, endpoints_b) = create_stream_handle(stream_id, &config, 128, 64); + + // Wire a's send → b's recv via a DuplexStream + let (client, server) = tokio::io::duplex(1024 * 1024); + let (client_read, client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + + // a's send data-plane task: read from cmd_rx, write to client_write + spawn_send_task(endpoints_a.send_cmd_rx, endpoints_a.send_evt_tx, endpoints_a.pool.clone(), client_write); + // b's recv data-plane task: read from server_read, push to evt_tx + spawn_recv_task(server_read, endpoints_b.recv_evt_tx, endpoints_b.pool.clone()); + + // b's send data-plane task: for the other direction (not used in basic tests) + spawn_send_task(endpoints_b.send_cmd_rx, endpoints_b.send_evt_tx, endpoints_b.pool.clone(), server_write); + // a's recv data-plane task + spawn_recv_task(client_read, endpoints_a.recv_evt_tx, endpoints_a.pool.clone()); + + // Return a's send half and b's recv half for unidirectional testing + (handle_a.send, handle_b.recv) + } + + fn spawn_send_task( + mut cmd_rx: tokio::sync::mpsc::Receiver, + evt_tx: tokio::sync::mpsc::Sender, + pool: swactor_streams::BufferPool, + mut writer: tokio::io::WriteHalf, + ) { + use tokio::io::AsyncWriteExt; + tokio::spawn(async move { + while let Some(cmd) = cmd_rx.recv().await { + match cmd { + swactor_streams::channel::SendCommand::Data(buf) => { + let data = buf.written(); + // Write length-prefixed frame + let len = (data.len() as u32).to_be_bytes(); + if writer.write_all(&len).await.is_err() { + pool.checkin(buf); + let _ = evt_tx.send(swactor_streams::channel::SendEvent::Error( + swactor_streams::StreamError::Disconnected, + )).await; + return; + } + if writer.write_all(data).await.is_err() { + pool.checkin(buf); + let _ = evt_tx.send(swactor_streams::channel::SendEvent::Error( + swactor_streams::StreamError::Disconnected, + )).await; + return; + } + pool.checkin(buf); + } + swactor_streams::channel::SendCommand::Flush => { + let _ = writer.flush().await; + } + swactor_streams::channel::SendCommand::Close => { + let _ = writer.shutdown().await; + break; + } + } + } + }); + } + + fn spawn_recv_task( + mut reader: tokio::io::ReadHalf, + evt_tx: tokio::sync::mpsc::Sender, + pool: swactor_streams::BufferPool, + ) { + use tokio::io::AsyncReadExt; + tokio::spawn(async move { + loop { + // Read length-prefixed frame + let mut len_buf = [0u8; 4]; + match reader.read_exact(&mut len_buf).await { + Ok(_) => {} + Err(_) => { + let _ = evt_tx.send(swactor_streams::channel::RecvEvent::Closed).await; + return; + } + } + let len = u32::from_be_bytes(len_buf) as usize; + let mut data = vec![0u8; len]; + match reader.read_exact(&mut data).await { + Ok(_) => {} + Err(_) => { + let _ = evt_tx.send(swactor_streams::channel::RecvEvent::Closed).await; + return; + } + } + + // Write data into FrameBufs and send + let mut offset = 0; + while offset < data.len() { + let mut buf = match pool.checkout() { + Some(b) => b, + None => { + let _ = evt_tx.send(swactor_streams::channel::RecvEvent::Error( + swactor_streams::StreamError::BufferExhausted, + )).await; + return; + } + }; + let written = buf.write(&data[offset..]); + offset += written; + let _ = evt_tx.send(swactor_streams::channel::RecvEvent::Data(buf)).await; + } + } + }); + } + + #[tokio::test] + async fn small_blob_round_trips() { + let data = b"hello, world!"; + let (_, manifest, chunks) = chunk_blob(data, 1024); + + let (mut send, mut recv) = create_test_pair(); + + let manifest_clone = manifest.clone(); + let chunks_clone = chunks.clone(); + let send_task = tokio::spawn(async move { + send_blob(&mut send, &manifest_clone, |hash| { + let chunks = chunks_clone.clone(); + async move { + chunks + .iter() + .find(|(h, _)| *h == hash) + .map(|(_, d)| d.clone()) + .ok_or_else(|| BlobTransferError::Storage("chunk not found".into())) + } + }, 0) + .await + .unwrap(); + }); + + let recv_task = tokio::spawn(async move { + recv_blob(&mut recv, 0).await.unwrap() + }); + + send_task.await.unwrap(); + let received = recv_task.await.unwrap(); + + assert_eq!(received.manifest, manifest); + assert_eq!(received.chunks.len(), chunks.len()); + for (i, (hash, data)) in received.chunks.iter().enumerate() { + assert_eq!(*hash, chunks[i].0); + assert_eq!(*data, chunks[i].1); + } + } + + #[tokio::test] + async fn multi_chunk_round_trips() { + // 4MB with 256KB chunks = 16 chunks + let data: Vec = (0..4 * 1024 * 1024).map(|i| (i % 251) as u8).collect(); + let (_, manifest, chunks) = chunk_blob(&data, 256 * 1024); + assert_eq!(chunks.len(), 16); + + let (mut send, mut recv) = create_test_pair(); + + let manifest_clone = manifest.clone(); + let chunks_clone = chunks.clone(); + let send_task = tokio::spawn(async move { + send_blob(&mut send, &manifest_clone, |hash| { + let chunks = chunks_clone.clone(); + async move { + chunks + .iter() + .find(|(h, _)| *h == hash) + .map(|(_, d)| d.clone()) + .ok_or_else(|| BlobTransferError::Storage("chunk not found".into())) + } + }, 0) + .await + .unwrap(); + }); + + let recv_task = tokio::spawn(async move { + recv_blob(&mut recv, 0).await.unwrap() + }); + + send_task.await.unwrap(); + let received = recv_task.await.unwrap(); + + assert_eq!(received.manifest, manifest); + assert_eq!(received.chunks.len(), 16); + // Verify all chunk data matches + for (i, (hash, cdata)) in received.chunks.iter().enumerate() { + assert_eq!(*hash, chunks[i].0); + assert_eq!(cdata.len(), chunks[i].1.len()); + } + } + + #[tokio::test] + async fn corrupted_chunk_detected() { + let data = b"integrity test data here"; + let (_, manifest, mut chunks) = chunk_blob(data, 1024); + + // Flip a byte in the chunk data + chunks[0].1[0] ^= 0xFF; + + let (mut send, mut recv) = create_test_pair(); + + let manifest_clone = manifest.clone(); + let chunks_clone = chunks.clone(); + let send_task = tokio::spawn(async move { + send_blob(&mut send, &manifest_clone, |hash| { + let chunks = chunks_clone.clone(); + async move { + // Note: we send the corrupted data (hash won't match) + chunks + .iter() + .find(|(h, _)| *h == hash) + .map(|(_, d)| d.clone()) + .ok_or_else(|| BlobTransferError::Storage("chunk not found".into())) + } + }, 0) + .await + .unwrap(); + }); + + let recv_task = tokio::spawn(async move { + recv_blob(&mut recv, 0).await + }); + + send_task.await.unwrap(); + let result = recv_task.await.unwrap(); + + match result { + Err(BlobTransferError::ChunkVerificationFailed { index: 0, .. }) => {} + other => panic!("expected ChunkVerificationFailed, got: {other:?}"), + } + } + + #[tokio::test] + async fn truncated_stream_detected() { + // Create a blob with multiple chunks + let data: Vec = vec![42u8; 4096]; + let (_, manifest, chunks) = chunk_blob(&data, 1024); + + let (mut send, mut recv) = create_test_pair(); + + let manifest_clone = manifest.clone(); + // Only send the manifest + first chunk, then close + let send_task = tokio::spawn(async move { + let manifest_json = serde_json::to_vec(&manifest_clone).unwrap(); + let len_bytes = (manifest_json.len() as u32).to_be_bytes(); + write_all(&mut send, &len_bytes).await.unwrap(); + write_all(&mut send, &manifest_json).await.unwrap(); + // Write first chunk + write_all(&mut send, &chunks[0].1).await.unwrap(); + // Close without writing remaining chunks + send.flush().unwrap(); + send.close().unwrap(); + }); + + let recv_task = tokio::spawn(async move { + recv_blob(&mut recv, 0).await + }); + + send_task.await.unwrap(); + let result = recv_task.await.unwrap(); + + match result { + Err(BlobTransferError::IncompleteTransfer(_)) => {} + other => panic!("expected IncompleteTransfer, got: {other:?}"), + } + } + + // ── Resume token tests ────────────────────────────────────────────── + + #[tokio::test] + async fn resume_skips_first_n_chunks() { + // 16 chunks, resume from chunk 8 → only chunks 8-15 transferred + let data: Vec = (0..4 * 1024 * 1024).map(|i| (i % 251) as u8).collect(); + let (_, manifest, chunks) = chunk_blob(&data, 256 * 1024); + assert_eq!(chunks.len(), 16); + + let skip = 8u64; + let (mut send, mut recv) = create_test_pair(); + + let manifest_clone = manifest.clone(); + let chunks_clone = chunks.clone(); + let send_task = tokio::spawn(async move { + send_blob(&mut send, &manifest_clone, |hash| { + let chunks = chunks_clone.clone(); + async move { + chunks + .iter() + .find(|(h, _)| *h == hash) + .map(|(_, d)| d.clone()) + .ok_or_else(|| BlobTransferError::Storage("chunk not found".into())) + } + }, skip) + .await + .unwrap(); + }); + + let recv_task = tokio::spawn(async move { + recv_blob(&mut recv, skip).await.unwrap() + }); + + send_task.await.unwrap(); + let received = recv_task.await.unwrap(); + + assert_eq!(received.manifest, manifest); + assert_eq!(received.chunks.len(), 8); // only chunks 8-15 + for (j, (hash, cdata)) in received.chunks.iter().enumerate() { + let orig_idx = skip as usize + j; + assert_eq!(*hash, chunks[orig_idx].0); + assert_eq!(*cdata, chunks[orig_idx].1); + } + } + + #[tokio::test] + async fn resume_from_zero_is_full_transfer() { + let data: Vec = (0..4 * 1024 * 1024).map(|i| (i % 251) as u8).collect(); + let (_, manifest, chunks) = chunk_blob(&data, 256 * 1024); + + let (mut send, mut recv) = create_test_pair(); + + let manifest_clone = manifest.clone(); + let chunks_clone = chunks.clone(); + let send_task = tokio::spawn(async move { + send_blob(&mut send, &manifest_clone, |hash| { + let chunks = chunks_clone.clone(); + async move { + chunks + .iter() + .find(|(h, _)| *h == hash) + .map(|(_, d)| d.clone()) + .ok_or_else(|| BlobTransferError::Storage("chunk not found".into())) + } + }, 0) + .await + .unwrap(); + }); + + let recv_task = tokio::spawn(async move { + recv_blob(&mut recv, 0).await.unwrap() + }); + + send_task.await.unwrap(); + let received = recv_task.await.unwrap(); + + assert_eq!(received.chunks.len(), chunks.len()); + } + + #[tokio::test] + async fn resume_from_last_chunk() { + // 16 chunks, skip=15 → only the final chunk transfers + let data: Vec = (0..4 * 1024 * 1024).map(|i| (i % 251) as u8).collect(); + let (_, manifest, chunks) = chunk_blob(&data, 256 * 1024); + assert_eq!(chunks.len(), 16); + + let skip = 15u64; + let (mut send, mut recv) = create_test_pair(); + + let manifest_clone = manifest.clone(); + let chunks_clone = chunks.clone(); + let send_task = tokio::spawn(async move { + send_blob(&mut send, &manifest_clone, |hash| { + let chunks = chunks_clone.clone(); + async move { + chunks + .iter() + .find(|(h, _)| *h == hash) + .map(|(_, d)| d.clone()) + .ok_or_else(|| BlobTransferError::Storage("chunk not found".into())) + } + }, skip) + .await + .unwrap(); + }); + + let recv_task = tokio::spawn(async move { + recv_blob(&mut recv, skip).await.unwrap() + }); + + send_task.await.unwrap(); + let received = recv_task.await.unwrap(); + + assert_eq!(received.chunks.len(), 1); + assert_eq!(received.chunks[0].0, chunks[15].0); + assert_eq!(received.chunks[0].1, chunks[15].1); + } + + #[test] + fn metadata_encoding_round_trip() { + let meta = BlobTransferMetadata { + content_hash: ContentHash([42u8; 32]), + resume_from_chunk: Some(50), + }; + let encoded = encode_metadata(&meta); + assert_eq!(encoded[0], 0x01); + let decoded = parse_metadata(&encoded).unwrap(); + assert_eq!(decoded, meta); + } + + #[test] + fn metadata_backward_compat() { + // Legacy 32-byte format → parsed as no resume offset + let raw_hash = [0xABu8; 32]; + let decoded = parse_metadata(&raw_hash).unwrap(); + assert_eq!(decoded.content_hash, ContentHash(raw_hash)); + assert_eq!(decoded.resume_from_chunk, None); + } + + mod proptests { + use super::*; + use proptest::prelude::*; + + proptest! { + #[test] + fn arbitrary_blob_round_trips( + data in proptest::collection::vec(any::(), 1..=128 * 1024), + chunk_size in 256u32..=64 * 1024, + ) { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let (_, manifest, chunks) = chunk_blob(&data, chunk_size); + let (mut send, mut recv) = create_test_pair(); + + let manifest_clone = manifest.clone(); + let chunks_clone = chunks.clone(); + let send_task = tokio::spawn(async move { + send_blob(&mut send, &manifest_clone, |hash| { + let chunks = chunks_clone.clone(); + async move { + chunks + .iter() + .find(|(h, _)| *h == hash) + .map(|(_, d)| d.clone()) + .ok_or_else(|| BlobTransferError::Storage("not found".into())) + } + }, 0) + .await + .unwrap(); + }); + + let recv_task = tokio::spawn(async move { + recv_blob(&mut recv, 0).await.unwrap() + }); + + send_task.await.unwrap(); + let received = recv_task.await.unwrap(); + + prop_assert_eq!(received.manifest, manifest); + prop_assert_eq!(received.chunks.len(), chunks.len()); + for (i, (hash, cdata)) in received.chunks.iter().enumerate() { + prop_assert_eq!(*hash, chunks[i].0); + prop_assert_eq!(cdata, &chunks[i].1); + } + + Ok(()) + })?; + } + + #[test] + fn arbitrary_resume_offset_round_trips( + data in proptest::collection::vec(any::(), 1..=128 * 1024), + chunk_size in 256u32..=64 * 1024, + skip_frac in 0.0f64..1.0, + ) { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let (_, manifest, chunks) = chunk_blob(&data, chunk_size); + let total = chunks.len() as u64; + let skip = (skip_frac * total as f64).floor() as u64; + + let (mut send, mut recv) = create_test_pair(); + + let manifest_clone = manifest.clone(); + let chunks_clone = chunks.clone(); + let send_task = tokio::spawn(async move { + send_blob(&mut send, &manifest_clone, |hash| { + let chunks = chunks_clone.clone(); + async move { + chunks + .iter() + .find(|(h, _)| *h == hash) + .map(|(_, d)| d.clone()) + .ok_or_else(|| BlobTransferError::Storage("not found".into())) + } + }, skip) + .await + .unwrap(); + }); + + let recv_task = tokio::spawn(async move { + recv_blob(&mut recv, skip).await.unwrap() + }); + + send_task.await.unwrap(); + let received = recv_task.await.unwrap(); + + let expected_count = total - skip; + prop_assert_eq!(received.chunks.len() as u64, expected_count); + for (j, (hash, cdata)) in received.chunks.iter().enumerate() { + let orig_idx = skip as usize + j; + prop_assert_eq!(*hash, chunks[orig_idx].0); + prop_assert_eq!(cdata, &chunks[orig_idx].1); + } + + Ok(()) + })?; + } + } + } +} diff --git a/crates/datastore/src/bridge.rs b/crates/datastore/src/bridge.rs index 5e77b38..0b23127 100644 --- a/crates/datastore/src/bridge.rs +++ b/crates/datastore/src/bridge.rs @@ -381,6 +381,7 @@ pub struct DatastoreAuthConfig { /// Owns the full lifecycle of a datastore actor group: BlobStore, Metadata, /// DatastoreNode, and optional GatewayActor. pub struct DatastoreGroup { + datastore_addr: ActorAddress, metadata_addr: ActorAddress, gateway_addr: Option, bridge: Arc, @@ -472,6 +473,7 @@ impl DatastoreGroup { } Ok(Self { + datastore_addr, metadata_addr, gateway_addr, bridge, @@ -500,6 +502,30 @@ impl DatastoreGroup { pub fn bridge(&self) -> &Arc { &self.bridge } + + /// Configure stream support: sends ConfigureStreams to DatastoreNode and + /// spawns a StreamListener actor. + pub fn configure_streams( + &self, + stream_manager: ActorAddress, + tokio_handle: tokio::runtime::Handle, + ) { + let _ = self.runtime.send_to( + self.datastore_addr, + DatastoreNodeMsg::ConfigureStreams { + stream_manager, + tokio_handle, + runtime: Arc::clone(&self.runtime), + }, + ); + + // Spawn StreamListener + use crate::actors::stream_listener::StreamListener; + use swactor_std::RuntimeNaming; + if let Ok(addr) = self.runtime.spawn(StreamListener::new(self.datastore_addr, stream_manager)) { + let _ = self.runtime.register_name("StreamListener", addr); + } + } } fn generate_node_id() -> NodeId { diff --git a/crates/datastore/src/lib.rs b/crates/datastore/src/lib.rs index 1578c07..1170505 100644 --- a/crates/datastore/src/lib.rs +++ b/crates/datastore/src/lib.rs @@ -8,6 +8,7 @@ pub mod cli; pub mod metrics; pub mod api; pub mod ui_html; +pub mod blob_transfer; pub mod bridge; pub use types::{ChunkRef, ContentHash, DatastoreConfig, ObjectEntry, ObjectManifest}; diff --git a/crates/datastore/src/messages.rs b/crates/datastore/src/messages.rs index 92d0d4b..10f7d9b 100644 --- a/crates/datastore/src/messages.rs +++ b/crates/datastore/src/messages.rs @@ -9,12 +9,15 @@ use std::collections::BTreeMap; use std::collections::HashSet; use std::net::SocketAddr; +use std::sync::Arc; use serde::{Deserialize, Serialize}; use swactor::actor::ActorAddress; +use swactor::runtime::Runtime; use swactor::transport::NetworkMessage; use distribution::types::NodeId; +use swactor_streams::types::StreamId; use crate::auth::{AccessRequestInfo, AuthorizedKeyInfo, DeniedReason, SignedRequest}; use crate::types::{ContentHash, ObjectEntry, ObjectManifest}; @@ -276,7 +279,7 @@ pub enum TransferMsg { // ─── DatastoreNodeMsg ─────────────────────────────────────────────────────── /// Messages handled by the `DatastoreNode` coordinator actor. -#[derive(Debug, Clone)] +#[derive(Clone)] pub enum DatastoreNodeMsg { // ── User-facing commands ──────────────────────────────────────────── /// Store a blob with optional name and tags. @@ -333,6 +336,64 @@ pub enum DatastoreNodeMsg { request: ListObjectsRequest, reply_to: ActorAddress, }, + + // ── Stream-based blob transfer ───────────────────────────────────── + /// Download a blob via QUIC stream from a remote node. + DownloadViaStream { + content_hash: ContentHash, + source_node: [u8; 32], + reply_to: ActorAddress, + }, + /// Handle an incoming stream offer (from StreamListener). + HandleStreamOffer { + stream_id: StreamId, + content_hash: ContentHash, + from_node: [u8; 32], + stream_manager: ActorAddress, + resume_from_chunk: u64, + }, + /// A stream download completed successfully. + StreamDownloadComplete { + content_hash: ContentHash, + manifest: ObjectManifest, + reply_to: ActorAddress, + }, + /// A stream download failed. + StreamDownloadFailed { + content_hash: ContentHash, + reason: String, + chunks_completed: u64, + reply_to: ActorAddress, + }, + /// Configure stream support (StreamManager address + tokio handle). + ConfigureStreams { + stream_manager: ActorAddress, + tokio_handle: tokio::runtime::Handle, + runtime: Arc, + }, +} + +impl std::fmt::Debug for DatastoreNodeMsg { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Put { name, .. } => f.debug_struct("Put").field("name", name).finish_non_exhaustive(), + Self::Get { content_hash, .. } => f.debug_struct("Get").field("content_hash", content_hash).finish_non_exhaustive(), + Self::Delete { content_hash, .. } => f.debug_struct("Delete").field("content_hash", content_hash).finish_non_exhaustive(), + Self::List { name_filter, all, .. } => f.debug_struct("List").field("name_filter", name_filter).field("all", all).finish_non_exhaustive(), + Self::Status { .. } => write!(f, "Status"), + Self::ReadChunk { hash, .. } => f.debug_struct("ReadChunk").field("hash", hash).finish_non_exhaustive(), + Self::IncomingGetChunk { .. } => write!(f, "IncomingGetChunk"), + Self::IncomingGetManifest { .. } => write!(f, "IncomingGetManifest"), + Self::IncomingStoreObject { .. } => write!(f, "IncomingStoreObject"), + Self::IncomingFindObject { .. } => write!(f, "IncomingFindObject"), + Self::IncomingListObjects { .. } => write!(f, "IncomingListObjects"), + Self::DownloadViaStream { content_hash, .. } => f.debug_struct("DownloadViaStream").field("content_hash", content_hash).finish_non_exhaustive(), + Self::HandleStreamOffer { stream_id, content_hash, resume_from_chunk, .. } => f.debug_struct("HandleStreamOffer").field("stream_id", stream_id).field("content_hash", content_hash).field("resume_from_chunk", resume_from_chunk).finish_non_exhaustive(), + Self::StreamDownloadComplete { content_hash, .. } => f.debug_struct("StreamDownloadComplete").field("content_hash", content_hash).finish_non_exhaustive(), + Self::StreamDownloadFailed { content_hash, reason, chunks_completed, .. } => f.debug_struct("StreamDownloadFailed").field("content_hash", content_hash).field("reason", reason).field("chunks_completed", chunks_completed).finish_non_exhaustive(), + Self::ConfigureStreams { stream_manager, .. } => f.debug_struct("ConfigureStreams").field("stream_manager", stream_manager).finish_non_exhaustive(), + } + } } // ─── DatastoreResponse ────────────────────────────────────────────────────── diff --git a/crates/datastore/tests/stream_integration.rs b/crates/datastore/tests/stream_integration.rs new file mode 100644 index 0000000..110ff46 --- /dev/null +++ b/crates/datastore/tests/stream_integration.rs @@ -0,0 +1,415 @@ +//! End-to-end test: store a blob on node A, download via QUIC stream on node B. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use distribution::iroh_driver::{IrohDriver, IrohDriverConfig}; +use distribution::node::DistributedNodeConfig; +use iroh::RelayMode; + +use swactor::config::RuntimeConfig; +use swactor::runtime::Runtime; +use swactor_datastore::bridge::{DatastoreGroup, DatastoreGroupConfig}; +use swactor_datastore::messages::{DatastoreNodeMsg, DatastoreResponse}; +use swactor_std::RuntimeNaming; + +fn make_driver_with_streams() -> IrohDriver { + IrohDriver::new(IrohDriverConfig { + secret_key: None, + relay_mode: RelayMode::Disabled, + node: DistributedNodeConfig::default(), + peer_auth: None, + additional_alpns: vec![swactor_streams::ALPN.to_vec()], + }) + .expect("create iroh driver") +} + +fn make_runtime() -> Arc { + let rt = Runtime::new(RuntimeConfig { + num_threads: 2, + max_actors: 256, + channel_buffer_size: 2000, + ..Default::default() + }) + .with_extension(Arc::new(swactor_std::StdExtension::new())); + + let handle = rt.run().expect("start runtime"); + handle.runtime +} + +fn spawn_stream_manager( + runtime: &Arc, + driver: &IrohDriver, +) -> swactor::actor::ActorAddress { + let mgr = swactor_streams::StreamManager::new( + driver.endpoint().clone(), + driver.tokio_handle(), + Arc::clone(runtime), + ); + let addr = runtime.spawn(mgr).expect("spawn StreamManager"); + runtime + .register_name(swactor_streams::STREAM_MANAGER_NAME, addr) + .expect("register StreamManager"); + addr +} + +fn spawn_datastore( + runtime: &Arc, + driver: &IrohDriver, + mgr_addr: swactor::actor::ActorAddress, +) -> DatastoreGroup { + let node_id = driver.node_id(); + let group = DatastoreGroup::spawn( + Arc::clone(runtime), + DatastoreGroupConfig { + node_id, + node_id_hex: format!("{:?}", node_id), + chunk_size: 256, + storage_path: None, // in-memory + auth: None, + gc_interval: u64::MAX, + disseminate_interval: u64::MAX, + }, + ) + .expect("spawn datastore"); + group.configure_streams(mgr_addr, driver.tokio_handle()); + group +} + +/// Poll for a response from the inbox, routing stream connections between +/// the two nodes. Actor message processing is handled by worker threads. +fn pump_until_response( + rt_a: &Arc, + rt_b: &Arc, + driver_a: &mut IrohDriver, + driver_b: &mut IrohDriver, + mgr_a: swactor::actor::ActorAddress, + mgr_b: swactor::actor::ActorAddress, + inbox: &swactor::runtime::Inbox, + timeout: Duration, +) -> Option { + let deadline = Instant::now() + timeout; + let tokio_handle = driver_a.tokio_handle(); + + loop { + // SWIM protocol ticks + driver_a.recv(); + driver_a.tick(); + driver_b.recv(); + driver_b.tick(); + + // Route incoming stream connections on node A + for (node_id, conn) in driver_a.drain_other_connections() { + let rt = Arc::clone(rt_a); + let mgr = mgr_a; + let node_bytes = node_id.0; + tokio_handle.spawn(async move { + let _ = swactor_streams::accept::handle_incoming(node_bytes, conn, &rt, mgr).await; + }); + } + + // Route incoming stream connections on node B + for (node_id, conn) in driver_b.drain_other_connections() { + let rt = Arc::clone(rt_b); + let mgr = mgr_b; + let node_bytes = node_id.0; + tokio_handle.spawn(async move { + let _ = swactor_streams::accept::handle_incoming(node_bytes, conn, &rt, mgr).await; + }); + } + + // Check for completion + if let Some(resp) = inbox.try_recv() { + return Some(resp); + } + + if Instant::now() >= deadline { + return None; + } + + std::thread::sleep(Duration::from_millis(10)); + } +} + +#[test] +fn small_blob_transfers_between_two_nodes_via_stream() { + let mut driver_a = make_driver_with_streams(); + let mut driver_b = make_driver_with_streams(); + + let rt_a = make_runtime(); + let rt_b = make_runtime(); + + let mgr_a = spawn_stream_manager(&rt_a, &driver_a); + let mgr_b = spawn_stream_manager(&rt_b, &driver_b); + + let _ds_a = spawn_datastore(&rt_a, &driver_a, mgr_a); + let ds_b = spawn_datastore(&rt_b, &driver_b, mgr_b); + let _ = &ds_b; // keep alive + + // Have both nodes discover each other via SWIM + let addr_a = driver_a.endpoint_addr(); + let addr_b = driver_b.endpoint_addr(); + driver_a.join(&[addr_b]); + driver_b.join(&[addr_a]); + + // Pump until SWIM membership converges + let swim_deadline = Instant::now() + Duration::from_secs(10); + loop { + driver_a.recv(); + driver_a.tick(); + driver_b.recv(); + driver_b.tick(); + + let snap_a = driver_a.snapshot(); + let snap_b = driver_b.snapshot(); + if snap_a.alive_count >= 1 && snap_b.alive_count >= 1 { + break; + } + if Instant::now() >= swim_deadline { + panic!("SWIM convergence timed out"); + } + std::thread::sleep(Duration::from_millis(50)); + } + + // Give worker threads time to process spawned actors + std::thread::sleep(Duration::from_millis(100)); + + // Store a small blob on Node A + let test_data = b"Hello from node A! This is a stream integration test."; + let put_inbox = rt_a + .new_inbox::() + .expect("create inbox"); + let ds_a_addr = rt_a.where_is("Datastore").expect("Datastore registered on A"); + let _ = rt_a.send_to( + ds_a_addr, + DatastoreNodeMsg::Put { + data: test_data.to_vec(), + name: Some("test-blob".into()), + tags: Default::default(), + reply_to: *put_inbox.addr(), + }, + ); + + // Wait for PutOk (worker threads process messages) + let content_hash = loop { + if let Some(resp) = put_inbox.try_recv() { + match resp { + DatastoreResponse::PutOk { content_hash } => break content_hash, + other => panic!("expected PutOk, got: {other:?}"), + } + } + std::thread::sleep(Duration::from_millis(10)); + }; + + // PutOk comes from MetadataActor; BlobStore writes are fire-and-forget. + // Give BlobStore time to finish writing chunks + manifest. + std::thread::sleep(Duration::from_millis(200)); + + // Download the blob on Node B via stream + let download_inbox = rt_b + .new_inbox::() + .expect("create inbox"); + let ds_b_addr = rt_b.where_is("Datastore").expect("Datastore registered on B"); + let _ = rt_b.send_to( + ds_b_addr, + DatastoreNodeMsg::DownloadViaStream { + content_hash, + source_node: driver_a.node_id().0, + reply_to: *download_inbox.addr(), + }, + ); + + // Pump loop until we get a response + let resp = pump_until_response( + &rt_a, + &rt_b, + &mut driver_a, + &mut driver_b, + mgr_a, + mgr_b, + &download_inbox, + Duration::from_secs(15), + ); + + match resp { + Some(DatastoreResponse::PutOk { content_hash: h }) => { + assert_eq!(h, content_hash, "downloaded blob hash should match"); + } + other => panic!("expected PutOk from download, got: {other:?}"), + } + + // Verify: read the blob back from Node B's datastore + let verify_inbox = rt_b + .new_inbox::() + .expect("create inbox"); + let _ = rt_b.send_to( + ds_b_addr, + DatastoreNodeMsg::Get { + content_hash, + reply_to: *verify_inbox.addr(), + }, + ); + + let verify_deadline = Instant::now() + Duration::from_secs(5); + loop { + if let Some(resp) = verify_inbox.try_recv() { + match resp { + DatastoreResponse::GetOk { entry, manifest } => { + assert_eq!(entry.content_hash, content_hash); + assert_eq!(manifest.total_size, test_data.len() as u64); + break; + } + other => panic!("expected GetOk, got: {other:?}"), + } + } + if Instant::now() >= verify_deadline { + panic!("verify timed out — blob not found on Node B"); + } + std::thread::sleep(Duration::from_millis(10)); + } + + driver_a.shutdown(); + driver_b.shutdown(); + rt_a.shutdown(); + rt_b.shutdown(); +} + +#[test] +fn multi_chunk_blob_transfers_between_two_nodes_via_stream() { + let mut driver_a = make_driver_with_streams(); + let mut driver_b = make_driver_with_streams(); + + let rt_a = make_runtime(); + let rt_b = make_runtime(); + + let mgr_a = spawn_stream_manager(&rt_a, &driver_a); + let mgr_b = spawn_stream_manager(&rt_b, &driver_b); + + let _ds_a = spawn_datastore(&rt_a, &driver_a, mgr_a); + let ds_b = spawn_datastore(&rt_b, &driver_b, mgr_b); + let _ = &ds_b; + + // Discover each other via SWIM + let addr_a = driver_a.endpoint_addr(); + let addr_b = driver_b.endpoint_addr(); + driver_a.join(&[addr_b]); + driver_b.join(&[addr_a]); + + let swim_deadline = Instant::now() + Duration::from_secs(10); + loop { + driver_a.recv(); + driver_a.tick(); + driver_b.recv(); + driver_b.tick(); + + let snap_a = driver_a.snapshot(); + let snap_b = driver_b.snapshot(); + if snap_a.alive_count >= 1 && snap_b.alive_count >= 1 { + break; + } + if Instant::now() >= swim_deadline { + panic!("SWIM convergence timed out"); + } + std::thread::sleep(Duration::from_millis(50)); + } + + // Give worker threads time to process spawned actors + std::thread::sleep(Duration::from_millis(100)); + + // 4 chunks at 256 bytes each = 1024 bytes + let test_data: Vec = (0..1024).map(|i| (i % 251) as u8).collect(); + let put_inbox = rt_a + .new_inbox::() + .expect("create inbox"); + let ds_a_addr = rt_a.where_is("Datastore").expect("Datastore registered on A"); + let _ = rt_a.send_to( + ds_a_addr, + DatastoreNodeMsg::Put { + data: test_data.clone(), + name: Some("multi-chunk".into()), + tags: Default::default(), + reply_to: *put_inbox.addr(), + }, + ); + + let content_hash = loop { + if let Some(resp) = put_inbox.try_recv() { + match resp { + DatastoreResponse::PutOk { content_hash } => break content_hash, + other => panic!("expected PutOk, got: {other:?}"), + } + } + std::thread::sleep(Duration::from_millis(10)); + }; + + // PutOk comes from MetadataActor; BlobStore writes are fire-and-forget. + std::thread::sleep(Duration::from_millis(200)); + + // Download on Node B + let download_inbox = rt_b + .new_inbox::() + .expect("create inbox"); + let ds_b_addr = rt_b.where_is("Datastore").expect("Datastore registered on B"); + let _ = rt_b.send_to( + ds_b_addr, + DatastoreNodeMsg::DownloadViaStream { + content_hash, + source_node: driver_a.node_id().0, + reply_to: *download_inbox.addr(), + }, + ); + + let resp = pump_until_response( + &rt_a, + &rt_b, + &mut driver_a, + &mut driver_b, + mgr_a, + mgr_b, + &download_inbox, + Duration::from_secs(15), + ); + + match resp { + Some(DatastoreResponse::PutOk { content_hash: h }) => { + assert_eq!(h, content_hash); + } + other => panic!("expected PutOk from download, got: {other:?}"), + } + + // Verify the data on Node B by reading each chunk + let verify_inbox = rt_b + .new_inbox::() + .expect("create inbox"); + let _ = rt_b.send_to( + ds_b_addr, + DatastoreNodeMsg::Get { + content_hash, + reply_to: *verify_inbox.addr(), + }, + ); + + let verify_deadline = Instant::now() + Duration::from_secs(5); + loop { + if let Some(resp) = verify_inbox.try_recv() { + match resp { + DatastoreResponse::GetOk { entry, manifest } => { + assert_eq!(entry.content_hash, content_hash); + assert_eq!(manifest.total_size, test_data.len() as u64); + assert!(manifest.chunks.len() > 1, "should be multi-chunk"); + break; + } + other => panic!("expected GetOk, got: {other:?}"), + } + } + if Instant::now() >= verify_deadline { + panic!("verify timed out — blob not found on Node B"); + } + std::thread::sleep(Duration::from_millis(10)); + } + + driver_a.shutdown(); + driver_b.shutdown(); + rt_a.shutdown(); + rt_b.shutdown(); +} diff --git a/crates/distribution/src/iroh_driver.rs b/crates/distribution/src/iroh_driver.rs index 66fc1e1..8bd749b 100644 --- a/crates/distribution/src/iroh_driver.rs +++ b/crates/distribution/src/iroh_driver.rs @@ -40,6 +40,8 @@ pub struct IrohDriverConfig { pub node: DistributedNodeConfig, /// Optional peer allow-list. If provided, only allowed peers can connect. pub peer_auth: Option>>, + /// Additional ALPNs to register beyond SWIM. Opaque to the driver. + pub additional_alpns: Vec>, /// If set, start an embedded relay server on this address. /// Requires the `relay` feature. On success, the driver uses the embedded /// relay for `RelayMode::Custom`; on failure, falls back to `relay_mode`. @@ -74,8 +76,10 @@ pub struct IrohDriver { peer_auth: Option>>, /// Collects connections from background join tasks. pending_joins: Arc>>, - /// Connections accepted by the background accept loop. + /// Connections accepted by the background accept loop (SWIM ALPN). accepted_conns: Arc>>, + /// Connections accepted on non-SWIM ALPNs (streams, etc.). + other_accepted_conns: Arc>>, /// Relay URLs learned from join seeds, used for reconnection. peer_relay_urls: HashMap, /// Embedded relay server (if started). @@ -122,8 +126,10 @@ impl IrohDriver { let (relay_url, effective_relay_mode) = (None::, config.relay_mode); let endpoint = rt.block_on(async { + let mut alpns = vec![ALPN.to_vec()]; + alpns.extend(config.additional_alpns.iter().cloned()); let mut builder = Endpoint::empty_builder(effective_relay_mode) - .alpns(vec![ALPN.to_vec()]); + .alpns(alpns); if let Some(key) = config.secret_key { builder = builder.secret_key(key); @@ -141,10 +147,13 @@ impl IrohDriver { // Spawn background accept loop so incoming connections are never missed let accepted_conns: Arc>> = Arc::new(Mutex::new(Vec::new())); + let other_accepted_conns: Arc>> = + Arc::new(Mutex::new(Vec::new())); { let ep = endpoint.clone(); let peer_auth = config.peer_auth.clone(); - let buf = Arc::clone(&accepted_conns); + let swim_buf = Arc::clone(&accepted_conns); + let other_buf = Arc::clone(&other_accepted_conns); rt.spawn(async move { loop { match ep.accept().await { @@ -165,11 +174,22 @@ impl IrohDriver { conn.close(0u32.into(), b"unauthorized"); continue; } - eprintln!( - "iroh driver: accepted connection from {}", - crate::identity::hex_encode(&node_id.0[..4]) - ); - buf.lock().unwrap().push((node_id, conn)); + // Route by negotiated ALPN + let negotiated_alpn = conn.alpn(); + if negotiated_alpn == ALPN { + eprintln!( + "iroh driver: accepted SWIM connection from {}", + crate::identity::hex_encode(&node_id.0[..4]) + ); + swim_buf.lock().unwrap().push((node_id, conn)); + } else { + eprintln!( + "iroh driver: accepted non-SWIM connection from {} (ALPN: {})", + crate::identity::hex_encode(&node_id.0[..4]), + String::from_utf8_lossy(&negotiated_alpn), + ); + other_buf.lock().unwrap().push((node_id, conn)); + } } Err(e) => { eprintln!("iroh driver: incoming connection error: {e}"); @@ -189,6 +209,7 @@ impl IrohDriver { peer_auth: config.peer_auth, pending_joins: Arc::new(Mutex::new(Vec::new())), accepted_conns, + other_accepted_conns, peer_relay_urls: HashMap::new(), #[cfg(feature = "relay")] relay_server, @@ -201,6 +222,16 @@ impl IrohDriver { self.rt.handle().clone() } + /// Get a reference to the iroh endpoint (for creating outbound connections). + pub fn endpoint(&self) -> &Endpoint { + &self.endpoint + } + + /// Drain connections accepted on non-SWIM ALPNs. + pub fn drain_other_connections(&self) -> Vec<(NodeId, Connection)> { + self.other_accepted_conns.lock().unwrap().drain(..).collect() + } + /// The node's identity. pub fn node_id(&self) -> NodeId { self.node.node_id() diff --git a/crates/distribution/tests/common/iroh.rs b/crates/distribution/tests/common/iroh.rs index f4e18f5..e651e76 100644 --- a/crates/distribution/tests/common/iroh.rs +++ b/crates/distribution/tests/common/iroh.rs @@ -22,6 +22,7 @@ pub fn make_driver() -> IrohDriver { relay_mode: RelayMode::Disabled, node: test_config(), peer_auth: None, + additional_alpns: vec![], #[cfg(feature = "relay")] embedded_relay_bind: None, #[cfg(feature = "relay")] @@ -36,6 +37,7 @@ pub fn make_driver_with_auth(auth: Arc>) -> IrohDriver { relay_mode: RelayMode::Disabled, node: test_config(), peer_auth: Some(auth), + additional_alpns: vec![], #[cfg(feature = "relay")] embedded_relay_bind: None, #[cfg(feature = "relay")] @@ -50,6 +52,7 @@ pub fn make_driver_with_relay(relay_url: iroh::RelayUrl) -> IrohDriver { relay_mode: RelayMode::Custom(relay_url.into()), node: test_config(), peer_auth: None, + additional_alpns: vec![], #[cfg(feature = "relay")] embedded_relay_bind: None, #[cfg(feature = "relay")] diff --git a/crates/streams/Cargo.toml b/crates/streams/Cargo.toml new file mode 100644 index 0000000..3f81fc0 --- /dev/null +++ b/crates/streams/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "swactor-streams" +version = "0.1.0" +edition = "2024" + +[dependencies] +swactor = { path = "../..", features = ["serde"] } +swactor-std = { path = "../std" } +shared-types = { path = "../shared-types" } +distribution = { path = "../distribution" } +crossbeam-queue = "0.3.12" +tokio = { version = "1", features = ["sync", "io-util"] } +iroh = { version = "0.96" } +blake3 = "1" +serde = { version = "1", features = ["derive"] } +getrandom = "0.2" + +[dev-dependencies] +proptest = "1" +tokio = { version = "1", features = ["rt-multi-thread", "macros", "test-util", "io-util"] } diff --git a/crates/streams/src/accept.rs b/crates/streams/src/accept.rs new file mode 100644 index 0000000..93ab644 --- /dev/null +++ b/crates/streams/src/accept.rs @@ -0,0 +1,74 @@ +use std::sync::Arc; + +use iroh::endpoint::Connection; + +use swactor::actor::ActorAddress; +use swactor::runtime::Runtime; + +use crate::messages::{OneShot, StreamManagerMsg}; +use crate::wire; + +/// Spawn a bridge task that processes incoming stream connections and +/// forwards them to the StreamManager actor. +/// +/// For each `(node_id, conn)` received: +/// 1. Accept the control bi-stream +/// 2. Read the stream header +/// 3. Send `StreamManagerMsg::IncomingConnection` to the StreamManager +pub fn spawn_accept_bridge( + accepted_rx: tokio::sync::mpsc::Receiver<([u8; 32], Connection)>, + runtime: Arc, + manager_addr: ActorAddress, +) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + accept_bridge_loop(accepted_rx, runtime, manager_addr).await; + }) +} + +async fn accept_bridge_loop( + mut accepted_rx: tokio::sync::mpsc::Receiver<([u8; 32], Connection)>, + runtime: Arc, + manager_addr: ActorAddress, +) { + while let Some((node_id, conn)) = accepted_rx.recv().await { + let runtime = Arc::clone(&runtime); + let manager_addr = manager_addr; + tokio::spawn(async move { + if let Err(e) = handle_incoming(node_id, conn, &runtime, manager_addr).await + { + eprintln!("accept bridge: failed to handle incoming connection: {e}"); + } + }); + } +} + +/// Handle a single incoming stream connection: accept the control bi-stream, +/// read the header, and forward to the StreamManager. +pub async fn handle_incoming( + node_id: [u8; 32], + conn: Connection, + runtime: &Runtime, + manager_addr: ActorAddress, +) -> Result<(), Box> { + // Accept the control bi-stream (opener sends header here) + let (_, mut recv_ctrl) = conn.accept_bi().await?; + + // Read the full header into a buffer. The opener finishes the send side + // after writing the header, so read_to_end collects all header bytes. + let header_bytes = recv_ctrl.read_to_end(4096).await?; + + let header = wire::decode_header(&header_bytes) + .map_err(|e| format!("invalid stream header: {e}"))?; + + let msg = StreamManagerMsg::IncomingConnection { + node_id, + stream_id: header.stream_id, + mode: header.mode, + config: header.config, + conn: OneShot::new(conn), + }; + runtime + .send_to(manager_addr, msg) + .map_err(|e| format!("send to StreamManager failed: {e}"))?; + Ok(()) +} diff --git a/crates/streams/src/buffer.rs b/crates/streams/src/buffer.rs new file mode 100644 index 0000000..d22b29e --- /dev/null +++ b/crates/streams/src/buffer.rs @@ -0,0 +1,252 @@ +use std::sync::Arc; + +use crossbeam_queue::ArrayQueue; + +/// A pre-allocated frame buffer with read/write cursors for zero-alloc recycling. +/// +/// Data is written starting at `write_pos` and read starting at `read_pos`. +/// When recycled via `reset()`, only the cursors are zeroed -- no memset. +pub struct FrameBuf { + data: Box<[u8]>, + read_pos: usize, + write_pos: usize, +} + +impl FrameBuf { + /// Create a new buffer with the given capacity. + pub fn new(capacity: usize) -> Self { + FrameBuf { + data: vec![0u8; capacity].into_boxed_slice(), + read_pos: 0, + write_pos: 0, + } + } + + /// Write data into the buffer. Returns the number of bytes written. + pub fn write(&mut self, src: &[u8]) -> usize { + let available = self.data.len() - self.write_pos; + let n = src.len().min(available); + self.data[self.write_pos..self.write_pos + n].copy_from_slice(&src[..n]); + self.write_pos += n; + n + } + + /// Read data from the buffer. Returns the number of bytes read. + pub fn read(&mut self, dst: &mut [u8]) -> usize { + let available = self.write_pos - self.read_pos; + let n = dst.len().min(available); + dst[..n].copy_from_slice(&self.data[self.read_pos..self.read_pos + n]); + self.read_pos += n; + n + } + + /// Reset cursors for reuse. Does NOT zero the data. + pub fn reset(&mut self) { + self.read_pos = 0; + self.write_pos = 0; + } + + /// Number of unread bytes in the buffer. + pub fn remaining(&self) -> usize { + self.write_pos - self.read_pos + } + + /// Available space for writing. + pub fn available(&self) -> usize { + self.data.len() - self.write_pos + } + + /// Whether the buffer is full (no more write space). + pub fn is_full(&self) -> bool { + self.write_pos == self.data.len() + } + + /// Whether all written data has been read. + pub fn is_empty(&self) -> bool { + self.read_pos == self.write_pos + } + + /// Total capacity of the buffer. + pub fn capacity(&self) -> usize { + self.data.len() + } + + /// The written portion of the buffer as a slice. + pub fn written(&self) -> &[u8] { + &self.data[..self.write_pos] + } + + /// Load data directly into the buffer, replacing any existing content. + pub fn load(&mut self, src: &[u8]) { + assert!( + src.len() <= self.data.len(), + "source data exceeds buffer capacity" + ); + self.data[..src.len()].copy_from_slice(src); + self.read_pos = 0; + self.write_pos = src.len(); + } +} + +/// A fixed-size lock-free pool of `FrameBuf`s backed by `crossbeam::ArrayQueue`. +/// +/// Supports concurrent checkout/checkin between actor threads and tokio tasks. +#[derive(Clone)] +pub struct BufferPool { + inner: Arc>, + buf_capacity: usize, +} + +impl BufferPool { + /// Create a new pool with `count` buffers, each of `buf_capacity` bytes. + pub fn new(count: usize, buf_capacity: usize) -> Self { + let queue = ArrayQueue::new(count); + for _ in 0..count { + let _ = queue.push(FrameBuf::new(buf_capacity)); + } + BufferPool { + inner: Arc::new(queue), + buf_capacity, + } + } + + /// Check out a buffer from the pool. Returns `None` if exhausted. + pub fn checkout(&self) -> Option { + self.inner.pop() + } + + /// Return a buffer to the pool. The buffer is reset before being made available. + pub fn checkin(&self, mut buf: FrameBuf) { + buf.reset(); + // If push fails (pool full), the buffer is dropped -- this is fine. + let _ = self.inner.push(buf); + } + + /// Number of buffers currently available in the pool. + pub fn available(&self) -> usize { + self.inner.len() + } + + /// The capacity of each buffer in the pool. + pub fn buf_capacity(&self) -> usize { + self.buf_capacity + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn frame_buf_write_then_read_returns_same_data() { + let mut buf = FrameBuf::new(64); + let data = b"hello, streams!"; + + let written = buf.write(data); + assert_eq!(written, data.len()); + assert_eq!(buf.remaining(), data.len()); + + let mut out = vec![0u8; data.len()]; + let read = buf.read(&mut out); + assert_eq!(read, data.len()); + assert_eq!(&out, data); + assert!(buf.is_empty()); + } + + #[test] + fn frame_buf_partial_write_when_full() { + let mut buf = FrameBuf::new(8); + let written = buf.write(b"twelve chars"); + assert_eq!(written, 8); + assert!(buf.is_full()); + assert_eq!(buf.available(), 0); + } + + #[test] + fn frame_buf_reset_allows_reuse() { + let mut buf = FrameBuf::new(16); + buf.write(b"first"); + buf.reset(); + + assert!(buf.is_empty()); + assert_eq!(buf.remaining(), 0); + assert_eq!(buf.available(), 16); + + let written = buf.write(b"second"); + assert_eq!(written, 6); + + let mut out = vec![0u8; 6]; + buf.read(&mut out); + assert_eq!(&out, b"second"); + } + + #[test] + fn frame_buf_load_replaces_content() { + let mut buf = FrameBuf::new(32); + buf.write(b"old data"); + buf.load(b"new data here"); + assert_eq!(buf.remaining(), 13); + let mut out = vec![0u8; 13]; + buf.read(&mut out); + assert_eq!(&out, b"new data here"); + } + + #[test] + fn pool_checkout_checkin_cycle() { + let pool = BufferPool::new(4, 1024); + assert_eq!(pool.available(), 4); + + let b1 = pool.checkout().unwrap(); + let b2 = pool.checkout().unwrap(); + assert_eq!(pool.available(), 2); + + pool.checkin(b1); + assert_eq!(pool.available(), 3); + + pool.checkin(b2); + assert_eq!(pool.available(), 4); + } + + #[test] + fn pool_exhausted_returns_none() { + let pool = BufferPool::new(2, 64); + + let _b1 = pool.checkout().unwrap(); + let _b2 = pool.checkout().unwrap(); + assert!(pool.checkout().is_none()); + } + + #[test] + fn pool_checkin_after_exhaustion_restores_availability() { + let pool = BufferPool::new(1, 64); + + let buf = pool.checkout().unwrap(); + assert!(pool.checkout().is_none()); + + pool.checkin(buf); + assert!(pool.checkout().is_some()); + } + + #[test] + fn pool_checkin_resets_buffer() { + let pool = BufferPool::new(1, 64); + let mut buf = pool.checkout().unwrap(); + buf.write(b"dirty data"); + assert_eq!(buf.remaining(), 10); + + pool.checkin(buf); + + let recycled = pool.checkout().unwrap(); + assert!(recycled.is_empty()); + assert_eq!(recycled.available(), 64); + } + + #[test] + fn pool_clone_shares_same_backing() { + let pool = BufferPool::new(3, 128); + let pool2 = pool.clone(); + + let _b = pool.checkout().unwrap(); + assert_eq!(pool2.available(), 2); + } +} diff --git a/crates/streams/src/channel.rs b/crates/streams/src/channel.rs new file mode 100644 index 0000000..fcdee38 --- /dev/null +++ b/crates/streams/src/channel.rs @@ -0,0 +1,41 @@ +use crate::buffer::FrameBuf; +use crate::types::StreamError; + +/// Commands sent from the actor to the send-side data-plane task. +pub enum SendCommand { + /// A buffer of data to write to the wire. + Data(FrameBuf), + /// Flush any partially-filled buffers. + Flush, + /// Gracefully close the send side. + Close, +} + +/// Events sent from the send-side data-plane task back to the actor. +#[derive(Debug, Clone)] +pub enum SendEvent { + /// The data-plane task is ready to accept more data. + WriteReady, + /// An error occurred on the send side. + Error(StreamError), + /// The send side has been closed. + Closed, +} + +/// Commands sent from the actor to the recv-side data-plane task. +pub enum RecvCommand { + /// Return a consumed buffer to the pool. + Consumed(FrameBuf), + /// Close the receive side. + Close, +} + +/// Events sent from the recv-side data-plane task to the actor. +pub enum RecvEvent { + /// A buffer of received data. + Data(FrameBuf), + /// An error occurred on the recv side. + Error(StreamError), + /// The recv side has been closed (all stripes finished). + Closed, +} diff --git a/crates/streams/src/connection.rs b/crates/streams/src/connection.rs new file mode 100644 index 0000000..36ec1f6 --- /dev/null +++ b/crates/streams/src/connection.rs @@ -0,0 +1,65 @@ +use std::collections::HashMap; + +use iroh::endpoint::Connection; +use iroh::{Endpoint, PublicKey}; + +use crate::types::StreamError; +use crate::wire::ALPN; + +/// Cache of QUIC connections used for stream data transfer. +/// +/// Separate from the SWIM connection pool in IrohDriver. All connections +/// are established using the stream ALPN (`swactor/stream/1`). +pub struct StreamConnectionCache { + connections: HashMap<[u8; 32], Connection>, +} + +impl StreamConnectionCache { + pub fn new() -> Self { + StreamConnectionCache { + connections: HashMap::new(), + } + } + + /// Get an existing healthy connection or establish a new one. + pub async fn get_or_connect( + &mut self, + endpoint: &Endpoint, + node_id: [u8; 32], + ) -> Result { + // Check for cached connection that's still open + if let Some(conn) = self.connections.get(&node_id) { + if conn.close_reason().is_none() { + return Ok(conn.clone()); + } + // Connection closed, remove it + self.connections.remove(&node_id); + } + + let key = PublicKey::from_bytes(&node_id) + .map_err(|e| StreamError::BrokenPipe(format!("invalid public key: {e}")))?; + + let conn = endpoint + .connect(key, ALPN) + .await + .map_err(|e| StreamError::BrokenPipe(format!("connect failed: {e}")))?; + + self.connections.insert(node_id, conn.clone()); + Ok(conn) + } + + /// Remove dead connections from the cache. + pub fn prune_closed(&mut self) { + self.connections.retain(|_, conn| conn.close_reason().is_none()); + } + + /// Remove a specific connection. + pub fn remove(&mut self, node_id: &[u8; 32]) { + self.connections.remove(node_id); + } + + /// Insert a connection into the cache. + pub fn insert(&mut self, node_id: [u8; 32], conn: Connection) { + self.connections.insert(node_id, conn); + } +} diff --git a/crates/streams/src/ctx_ext.rs b/crates/streams/src/ctx_ext.rs new file mode 100644 index 0000000..f68f072 --- /dev/null +++ b/crates/streams/src/ctx_ext.rs @@ -0,0 +1,201 @@ +//! Convenience extension traits for actors that use streams. +//! +//! `CtxStreams` wraps StreamManager message construction for use inside actor +//! handlers (defaults `reply_to`/`listener` to `ctx.self_addr()`). +//! +//! `RuntimeStreams` provides the same operations from a `Runtime` handle, +//! requiring explicit addresses since there's no implicit "self". + +use swactor::actor::{ActorAddress, Ctx}; +use swactor::runtime::Runtime; +use swactor_std::CtxNaming; +use swactor_std::RuntimeNaming; + +use crate::messages::StreamManagerMsg; +use crate::types::{StreamConfig, StreamError, StreamId, StreamMode}; + +fn mgr_not_found() -> StreamError { + StreamError::BrokenPipe("StreamManager not found in name registry".into()) +} + +/// Stream operations available inside actor handlers via `Ctx`. +/// +/// All methods default `reply_to` / `listener` to `ctx.self_addr()`. +pub trait CtxStreams { + /// Open a new stream to a remote node. + fn stream_open( + &self, + target_node: [u8; 32], + mode: StreamMode, + config: StreamConfig, + ) -> Result<(), StreamError>; + + /// Register as a stream listener for the given mode. + fn stream_listen(&self, mode: StreamMode) -> Result<(), StreamError>; + + /// Accept an offered incoming stream. + fn stream_accept(&self, stream_id: StreamId) -> Result<(), StreamError>; + + /// Reject an offered incoming stream. + fn stream_reject(&self, stream_id: StreamId) -> Result<(), StreamError>; + + /// Close a stream. + fn stream_close(&self, stream_id: StreamId) -> Result<(), StreamError>; +} + +impl CtxStreams for Ctx<'_> { + fn stream_open( + &self, + target_node: [u8; 32], + mode: StreamMode, + config: StreamConfig, + ) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send( + mgr, + StreamManagerMsg::Open { + target_node, + mode, + config, + reply_to: self.self_addr(), + }, + ) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } + + fn stream_listen(&self, mode: StreamMode) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send( + mgr, + StreamManagerMsg::Listen { + mode, + listener: self.self_addr(), + }, + ) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } + + fn stream_accept(&self, stream_id: StreamId) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send( + mgr, + StreamManagerMsg::Accept { + stream_id, + reply_to: self.self_addr(), + }, + ) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } + + fn stream_reject(&self, stream_id: StreamId) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send(mgr, StreamManagerMsg::Reject { stream_id }) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } + + fn stream_close(&self, stream_id: StreamId) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send(mgr, StreamManagerMsg::Close { stream_id }) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } +} + +/// Stream operations available from a `Runtime` handle (outside actor handlers). +/// +/// Requires explicit `reply_to` / `listener` addresses. +pub trait RuntimeStreams { + /// Open a new stream to a remote node. + fn stream_open( + &self, + target_node: [u8; 32], + mode: StreamMode, + config: StreamConfig, + reply_to: ActorAddress, + ) -> Result<(), StreamError>; + + /// Register an address as a stream listener for the given mode. + fn stream_listen(&self, mode: StreamMode, listener: ActorAddress) -> Result<(), StreamError>; + + /// Accept an offered incoming stream. + fn stream_accept( + &self, + stream_id: StreamId, + reply_to: ActorAddress, + ) -> Result<(), StreamError>; + + /// Reject an offered incoming stream. + fn stream_reject(&self, stream_id: StreamId) -> Result<(), StreamError>; + + /// Close a stream. + fn stream_close(&self, stream_id: StreamId) -> Result<(), StreamError>; +} + +impl RuntimeStreams for Runtime { + fn stream_open( + &self, + target_node: [u8; 32], + mode: StreamMode, + config: StreamConfig, + reply_to: ActorAddress, + ) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send_to( + mgr, + StreamManagerMsg::Open { + target_node, + mode, + config, + reply_to, + }, + ) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } + + fn stream_listen(&self, mode: StreamMode, listener: ActorAddress) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send_to(mgr, StreamManagerMsg::Listen { mode, listener }) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } + + fn stream_accept( + &self, + stream_id: StreamId, + reply_to: ActorAddress, + ) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send_to(mgr, StreamManagerMsg::Accept { stream_id, reply_to }) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } + + fn stream_reject(&self, stream_id: StreamId) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send_to(mgr, StreamManagerMsg::Reject { stream_id }) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } + + fn stream_close(&self, stream_id: StreamId) -> Result<(), StreamError> { + let mgr = self + .where_is(crate::manager::STREAM_MANAGER_NAME) + .ok_or_else(mgr_not_found)?; + self.send_to(mgr, StreamManagerMsg::Close { stream_id }) + .map_err(|e| StreamError::BrokenPipe(e.to_string())) + } +} diff --git a/crates/streams/src/data_plane.rs b/crates/streams/src/data_plane.rs new file mode 100644 index 0000000..437669f --- /dev/null +++ b/crates/streams/src/data_plane.rs @@ -0,0 +1,499 @@ +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::sync::mpsc; + +use crate::buffer::BufferPool; +use crate::channel::{RecvEvent, SendCommand}; +use crate::notify::NotifySink; +use crate::types::StreamError; +use crate::wire; + +/// A send-side data-plane task for a single stripe. +/// +/// Reads `SendCommand`s from the actor's channel, encodes them as wire +/// frames, and writes them to the underlying transport. Returns consumed +/// buffers to the pool. +/// +/// Generic over `AsyncWrite` so it can be tested with `DuplexStream`. +pub async fn send_stripe_task( + mut writer: W, + mut cmd_rx: mpsc::Receiver, + pool: BufferPool, + notify: Option, +) -> Result<(), StreamError> +where + W: AsyncWriteExt + Unpin + Send + 'static, +{ + while let Some(cmd) = cmd_rx.recv().await { + match cmd { + SendCommand::Data(buf) => { + let frame = wire::encode_data_frame(buf.written()); + writer + .write_all(&frame) + .await + .map_err(|e| StreamError::BrokenPipe(e.to_string()))?; + + // Return buffer to pool + pool.checkin(buf); + + // Signal write ready + if let Some(ref sink) = notify { + sink.write_ready(); + } + } + SendCommand::Flush => { + writer + .flush() + .await + .map_err(|e| StreamError::BrokenPipe(e.to_string()))?; + } + SendCommand::Close => { + // Write end-of-stripe sentinel + let sentinel = wire::encode_end_of_stripe(); + writer + .write_all(&sentinel) + .await + .map_err(|e| StreamError::BrokenPipe(e.to_string()))?; + writer + .flush() + .await + .map_err(|e| StreamError::BrokenPipe(e.to_string()))?; + break; + } + } + } + + Ok(()) +} + +/// A recv-side data-plane task for a single stripe. +/// +/// Reads wire-encoded frames from the transport, fills `FrameBuf`s from +/// the pool, and sends them to the actor via the event channel. +/// +/// Generic over `AsyncRead` so it can be tested with `DuplexStream`. +pub async fn recv_stripe_task( + mut reader: R, + evt_tx: mpsc::Sender, + pool: BufferPool, + notify: Option, +) -> Result<(), StreamError> +where + R: AsyncReadExt + Unpin + Send + 'static, +{ + loop { + // Read the 4-byte length prefix + let mut len_buf = [0u8; 4]; + match reader.read_exact(&mut len_buf).await { + Ok(_) => {} + Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { + // Connection closed + let _ = evt_tx.send(RecvEvent::Closed).await; + if let Some(ref sink) = notify { + sink.closed(); + } + return Ok(()); + } + Err(e) => { + let err = StreamError::BrokenPipe(e.to_string()); + let _ = evt_tx.send(RecvEvent::Error(err.clone())).await; + if let Some(ref sink) = notify { + sink.error(); + } + return Err(err); + } + } + + let payload_len = u32::from_be_bytes(len_buf) as usize; + + // End-of-stripe sentinel + if payload_len == 0 { + let _ = evt_tx.send(RecvEvent::Closed).await; + if let Some(ref sink) = notify { + sink.closed(); + } + return Ok(()); + } + + // Read the payload into a buffer from the pool + let mut buf = match pool.checkout() { + Some(b) => b, + None => { + let err = StreamError::BufferExhausted; + let _ = evt_tx.send(RecvEvent::Error(err.clone())).await; + if let Some(ref sink) = notify { + sink.error(); + } + return Err(err); + } + }; + + let mut temp = vec![0u8; payload_len]; + match reader.read_exact(&mut temp).await { + Ok(_) => {} + Err(e) => { + pool.checkin(buf); + let err = StreamError::BrokenPipe(e.to_string()); + let _ = evt_tx.send(RecvEvent::Error(err.clone())).await; + if let Some(ref sink) = notify { + sink.error(); + } + return Err(err); + } + } + buf.load(&temp); + + // Send to actor + if evt_tx.send(RecvEvent::Data(buf)).await.is_err() { + return Err(StreamError::Disconnected); + } + + if let Some(ref sink) = notify { + sink.data_ready(); + } + } +} + +/// Spawn a complete set of send-side stripe tasks. +/// +/// Returns a Vec of `mpsc::Sender` -- one per stripe. +/// The caller assigns chunks round-robin: chunk `i` goes to stripe `i % stripe_count`. +pub fn spawn_send_stripes( + stripe_count: usize, + pool: BufferPool, + _notify: Option, + mut writer_factory: F, + channel_capacity: usize, +) -> Vec> +where + W: AsyncWriteExt + Unpin + Send + 'static, + F: FnMut(usize) -> W, +{ + let mut senders = Vec::with_capacity(stripe_count); + + for i in 0..stripe_count { + let (tx, rx) = mpsc::channel(channel_capacity); + let writer = writer_factory(i); + let pool = pool.clone(); + tokio::spawn(async move { + let _ = send_stripe_task(writer, rx, pool, None).await; + }); + senders.push(tx); + } + + senders +} + +/// Spawn a complete set of recv-side stripe tasks. +/// +/// Returns a single `mpsc::Receiver` that merges events from all stripes. +pub fn spawn_recv_stripes( + stripe_count: usize, + pool: BufferPool, + _notify: Option, + mut reader_factory: F, + channel_capacity: usize, +) -> mpsc::Receiver +where + R: AsyncReadExt + Unpin + Send + 'static, + F: FnMut(usize) -> R, +{ + // All stripes feed into a single merged channel + let (merged_tx, merged_rx) = mpsc::channel(channel_capacity * stripe_count); + + for i in 0..stripe_count { + let reader = reader_factory(i); + let pool = pool.clone(); + let tx = merged_tx.clone(); + tokio::spawn(async move { + let _ = recv_stripe_task(reader, tx, pool, None).await; + }); + } + + merged_rx +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::buffer::BufferPool; + use crate::types::StreamId; + + fn make_test_pool(count: usize, capacity: usize) -> BufferPool { + BufferPool::new(count, capacity) + } + + /// End-to-end: send data through a single stripe, receive it back. + #[tokio::test] + async fn single_stripe_transfer() { + let (client, server) = tokio::io::duplex(64 * 1024); + let pool = make_test_pool(16, 1024); + + let (cmd_tx, cmd_rx) = mpsc::channel(16); + let (evt_tx, mut evt_rx) = mpsc::channel(16); + + let send_pool = pool.clone(); + let recv_pool = pool.clone(); + + let send_handle = tokio::spawn(async move { + send_stripe_task(client, cmd_rx, send_pool, None).await + }); + let recv_handle = tokio::spawn(async move { + recv_stripe_task(server, evt_tx, recv_pool, None).await + }); + + // Send some data + let test_data = b"hello, streams!"; + let mut buf = pool.checkout().unwrap(); + buf.write(test_data); + cmd_tx.send(SendCommand::Data(buf)).await.unwrap(); + + // Send close + cmd_tx.send(SendCommand::Close).await.unwrap(); + + // Receive data + let evt = evt_rx.recv().await.unwrap(); + match evt { + RecvEvent::Data(mut buf) => { + let mut out = vec![0u8; test_data.len()]; + let n = buf.read(&mut out); + assert_eq!(n, test_data.len()); + assert_eq!(&out, test_data); + } + other => panic!("expected Data, got {:?}", std::mem::discriminant(&other)), + } + + // Receive close + let evt = evt_rx.recv().await.unwrap(); + assert!(matches!(evt, RecvEvent::Closed)); + + send_handle.await.unwrap().unwrap(); + recv_handle.await.unwrap().unwrap(); + } + + /// Multiple chunks through a single stripe. + #[tokio::test] + async fn multiple_chunks_single_stripe() { + let (client, server) = tokio::io::duplex(256 * 1024); + let pool = make_test_pool(32, 1024); + + let (cmd_tx, cmd_rx) = mpsc::channel(32); + let (evt_tx, mut evt_rx) = mpsc::channel(32); + + let sp = pool.clone(); + let rp = pool.clone(); + + tokio::spawn(async move { send_stripe_task(client, cmd_rx, sp, None).await }); + tokio::spawn(async move { recv_stripe_task(server, evt_tx, rp, None).await }); + + let chunk_count = 20; + for i in 0..chunk_count { + let mut buf = pool.checkout().unwrap(); + let data = format!("chunk-{i:04}"); + buf.write(data.as_bytes()); + cmd_tx.send(SendCommand::Data(buf)).await.unwrap(); + } + cmd_tx.send(SendCommand::Close).await.unwrap(); + + let mut received = Vec::new(); + loop { + match evt_rx.recv().await.unwrap() { + RecvEvent::Data(mut buf) => { + let mut out = vec![0u8; buf.remaining()]; + buf.read(&mut out); + received.push(String::from_utf8(out).unwrap()); + } + RecvEvent::Closed => break, + RecvEvent::Error(e) => panic!("unexpected error: {e}"), + } + } + + assert_eq!(received.len(), chunk_count); + for (i, chunk) in received.iter().enumerate() { + assert_eq!(chunk, &format!("chunk-{i:04}")); + } + } + + /// Multi-stripe transfer with round-robin assignment. + #[tokio::test] + async fn multi_stripe_round_robin() { + let stripe_count = 4; + let chunk_count = 100; + // Pool needs enough buffers for in-flight data on both sides + let pool = make_test_pool(256, 256); + + // Create duplex pairs for each stripe + let mut send_writers = Vec::new(); + let mut recv_readers = Vec::new(); + for _ in 0..stripe_count { + let (client, server) = tokio::io::duplex(64 * 1024); + send_writers.push(Some(client)); + recv_readers.push(Some(server)); + } + + // Spawn recv stripe tasks + let (merged_tx, mut merged_rx) = mpsc::channel(chunk_count * 2); + for i in 0..stripe_count { + let reader = recv_readers[i].take().unwrap(); + let p = pool.clone(); + let tx = merged_tx.clone(); + tokio::spawn(async move { + recv_stripe_task(reader, tx, p, None).await + }); + } + drop(merged_tx); // so merged_rx closes when all tasks finish + + // Spawn send stripe tasks + let mut stripe_txs = Vec::new(); + for i in 0..stripe_count { + let (tx, rx) = mpsc::channel(32); + let writer = send_writers[i].take().unwrap(); + let p = pool.clone(); + tokio::spawn(async move { + send_stripe_task(writer, rx, p, None).await + }); + stripe_txs.push(tx); + } + + // Send chunks round-robin + for i in 0..chunk_count { + let stripe_idx = i % stripe_count; + let mut buf = pool.checkout().unwrap(); + let data = format!("chunk-{i:04}"); + buf.write(data.as_bytes()); + stripe_txs[stripe_idx] + .send(SendCommand::Data(buf)) + .await + .unwrap(); + } + + // Close all stripes + for tx in &stripe_txs { + tx.send(SendCommand::Close).await.unwrap(); + } + + // Collect all received data (order may differ per stripe) + let mut received = Vec::new(); + let mut closed_count = 0; + while let Some(evt) = merged_rx.recv().await { + match evt { + RecvEvent::Data(mut buf) => { + let mut out = vec![0u8; buf.remaining()]; + buf.read(&mut out); + received.push(String::from_utf8(out).unwrap()); + } + RecvEvent::Closed => { + closed_count += 1; + if closed_count == stripe_count { + break; + } + } + RecvEvent::Error(e) => panic!("unexpected error: {e}"), + } + } + + // All chunks should have arrived (order may vary across stripes) + assert_eq!(received.len(), chunk_count); + received.sort(); + for (i, chunk) in received.iter().enumerate() { + assert_eq!(chunk, &format!("chunk-{i:04}")); + } + } + + /// Graceful close: writer closes, receiver sees end-of-stripe then Closed. + #[tokio::test] + async fn graceful_close() { + let (client, server) = tokio::io::duplex(64 * 1024); + let pool = make_test_pool(8, 256); + + let (cmd_tx, cmd_rx) = mpsc::channel(8); + let (evt_tx, mut evt_rx) = mpsc::channel(8); + + let sp = pool.clone(); + let rp = pool.clone(); + + tokio::spawn(async move { send_stripe_task(client, cmd_rx, sp, None).await }); + tokio::spawn(async move { recv_stripe_task(server, evt_tx, rp, None).await }); + + // Close immediately without sending data + cmd_tx.send(SendCommand::Close).await.unwrap(); + + // Should receive Closed + let evt = evt_rx.recv().await.unwrap(); + assert!(matches!(evt, RecvEvent::Closed)); + } + + /// Notification coalescing through NotifySink. + #[tokio::test] + async fn notification_coalescing() { + use crate::notify::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + let flag = Arc::new(NotifyFlag::new()); + let inject_count = Arc::new(AtomicUsize::new(0)); + let count_clone = inject_count.clone(); + + let stream_id = StreamId::new_random(); + let sink = NotifySink::new(flag.clone(), stream_id, move |_evt| { + count_clone.fetch_add(1, Ordering::SeqCst); + }); + + // First notification should inject + sink.data_ready(); + assert_eq!(inject_count.load(Ordering::SeqCst), 1); + + // Duplicate should coalesce (no inject) + sink.data_ready(); + assert_eq!(inject_count.load(Ordering::SeqCst), 1); + + // Clear and re-notify + flag.clear(DATA_READY); + sink.data_ready(); + assert_eq!(inject_count.load(Ordering::SeqCst), 2); + + // Different flag should still inject independently + sink.write_ready(); + assert_eq!(inject_count.load(Ordering::SeqCst), 3); + } + + /// Backpressure: when channel and active buffer are saturated, try_write returns 0. + #[tokio::test] + async fn send_backpressure() { + use crate::handle::create_stream_handle; + use crate::types::StreamConfig; + + let config = StreamConfig { + stripe_count: 1, + frame_size: 64, + metadata: vec![], + }; + + // Pool of 4, channel of 2 -- we can fill both quickly + let (mut handle, _endpoints) = create_stream_handle( + StreamId::new_random(), + &config, + 4, + 2, + ); + + let data = vec![0xAA; 64]; // exactly fills one buffer + + // Write 1: checks out buf, fills it (64 bytes), buf is full -> try_send succeeds + let n1 = handle.send.try_write(&data).unwrap(); + assert_eq!(n1, 64); + + // Write 2: checks out new buf, fills it, buf is full -> try_send succeeds + let n2 = handle.send.try_write(&data).unwrap(); + assert_eq!(n2, 64); + + // Write 3: checks out new buf, fills it, buf is full -> try_send fails (channel full) + // Data IS in the buffer (written=64), buffer kept locally + let n3 = handle.send.try_write(&data).unwrap(); + assert_eq!(n3, 64); + + // Write 4: active buf still full, buf.write() returns 0 (no space), + // try_send fails again -> returns 0 signaling backpressure + let n4 = handle.send.try_write(&data).unwrap(); + assert_eq!(n4, 0); // backpressure! + } +} diff --git a/crates/streams/src/handle.rs b/crates/streams/src/handle.rs new file mode 100644 index 0000000..17caef3 --- /dev/null +++ b/crates/streams/src/handle.rs @@ -0,0 +1,248 @@ +use tokio::sync::mpsc; + +use crate::buffer::{BufferPool, FrameBuf}; +use crate::channel::{RecvCommand, RecvEvent, SendCommand, SendEvent}; +use crate::types::{StreamConfig, StreamError, StreamId}; + +/// Sending half of a stream. Owned by the actor that sends data. +/// +/// Uses `try_send`/`try_recv` for non-blocking operation on the actor thread. +pub struct SendHalf { + stream_id: StreamId, + cmd_tx: mpsc::Sender, + evt_rx: mpsc::Receiver, + pool: BufferPool, + active_buf: Option, +} + +impl SendHalf { + pub fn stream_id(&self) -> StreamId { + self.stream_id + } + + /// Write data into the stream. Returns the number of bytes consumed. + /// + /// Fills the active buffer and sends full buffers to the data-plane task. + /// Returns 0 if the channel is full (backpressure) or the buffer pool + /// is exhausted. Does NOT block. + pub fn try_write(&mut self, data: &[u8]) -> Result { + if data.is_empty() { + return Ok(0); + } + + // Ensure we have an active buffer + if self.active_buf.is_none() { + self.active_buf = self.pool.checkout(); + if self.active_buf.is_none() { + return Err(StreamError::BufferExhausted); + } + } + + let buf = self.active_buf.as_mut().unwrap(); + let written = buf.write(data); + + // If the buffer is full, send it to the data-plane task + if buf.is_full() { + let full_buf = self.active_buf.take().unwrap(); + match self.cmd_tx.try_send(SendCommand::Data(full_buf)) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(cmd)) => { + // Put the buffer back -- channel is full (backpressure). + // Data is already written into the buffer. Next call to + // try_write will attempt try_send again. + if let SendCommand::Data(buf) = cmd { + self.active_buf = Some(buf); + } + return Ok(written); + } + Err(mpsc::error::TrySendError::Closed(_)) => { + return Err(StreamError::Disconnected); + } + } + } + + Ok(written) + } + + /// Flush any partially-filled buffer to the data-plane task. + pub fn flush(&mut self) -> Result<(), StreamError> { + if let Some(buf) = self.active_buf.take() { + if buf.remaining() > 0 || buf.written().len() > 0 { + self.cmd_tx + .try_send(SendCommand::Data(buf)) + .map_err(|_| StreamError::Disconnected)?; + } else { + self.pool.checkin(buf); + } + } + self.cmd_tx + .try_send(SendCommand::Flush) + .map_err(|_| StreamError::Disconnected)?; + Ok(()) + } + + /// Close the send side of the stream. + pub fn close(&mut self) -> Result<(), StreamError> { + if let Some(buf) = self.active_buf.take() { + if buf.written().len() > 0 { + let _ = self.cmd_tx.try_send(SendCommand::Data(buf)); + } else { + self.pool.checkin(buf); + } + } + self.cmd_tx + .try_send(SendCommand::Close) + .map_err(|_| StreamError::Disconnected)?; + Ok(()) + } + + /// Poll for events from the data-plane task (non-blocking). + pub fn try_recv_event(&mut self) -> Option { + self.evt_rx.try_recv().ok() + } +} + +/// Receiving half of a stream. Owned by the actor that receives data. +pub struct RecvHalf { + stream_id: StreamId, + evt_rx: mpsc::Receiver, + cmd_tx: mpsc::Sender, + pool: BufferPool, + active_buf: Option, +} + +impl RecvHalf { + pub fn stream_id(&self) -> StreamId { + self.stream_id + } + + /// Read data from the stream. Returns the number of bytes read. + /// + /// Drains the active buffer, then pulls new buffers from the channel. + /// Returns 0 if no data is currently available. Does NOT block. + pub fn try_read(&mut self, dst: &mut [u8]) -> Result { + if dst.is_empty() { + return Ok(0); + } + + // Drain any active buffer first + if let Some(buf) = &mut self.active_buf { + if buf.remaining() > 0 { + let read = buf.read(dst); + if buf.is_empty() { + let buf = self.active_buf.take().unwrap(); + self.pool.checkin(buf); + } + return Ok(read); + } else { + let buf = self.active_buf.take().unwrap(); + self.pool.checkin(buf); + } + } + + // Try to pull a new buffer from the channel + match self.evt_rx.try_recv() { + Ok(RecvEvent::Data(mut buf)) => { + let read = buf.read(dst); + if buf.is_empty() { + self.pool.checkin(buf); + } else { + self.active_buf = Some(buf); + } + Ok(read) + } + Ok(RecvEvent::Error(e)) => Err(e), + Ok(RecvEvent::Closed) => Err(StreamError::Closed), + Err(mpsc::error::TryRecvError::Empty) => Ok(0), + Err(mpsc::error::TryRecvError::Disconnected) => Err(StreamError::Disconnected), + } + } + + /// Check if data is available without consuming it. + pub fn has_data(&self) -> bool { + if let Some(buf) = &self.active_buf { + if buf.remaining() > 0 { + return true; + } + } + !self.evt_rx.is_empty() + } + + /// Close the receive side of the stream. + pub fn close(&mut self) -> Result<(), StreamError> { + if let Some(buf) = self.active_buf.take() { + self.pool.checkin(buf); + } + self.cmd_tx + .try_send(RecvCommand::Close) + .map_err(|_| StreamError::Disconnected)?; + Ok(()) + } +} + +/// Combined stream handle with both send and receive halves. +/// +/// `Send` but NOT `Clone` (mpsc::Receiver is not Clone). +pub struct StreamHandle { + pub send: SendHalf, + pub recv: RecvHalf, +} + +/// Channel endpoints for the data-plane tasks. +pub struct DataPlaneEndpoints { + /// Receive send commands from the actor. + pub send_cmd_rx: mpsc::Receiver, + /// Send events back to the actor. + pub send_evt_tx: mpsc::Sender, + /// Send received data to the actor. + pub recv_evt_tx: mpsc::Sender, + /// Receive consume/close commands from the actor. + pub recv_cmd_rx: mpsc::Receiver, + /// Shared buffer pool. + pub pool: BufferPool, +} + +/// Create a stream handle and its corresponding data-plane channel endpoints. +/// +/// `channel_capacity` controls how many FrameBufs can be in-flight between +/// the actor and the data-plane tasks. +pub fn create_stream_handle( + stream_id: StreamId, + config: &StreamConfig, + pool_size: usize, + channel_capacity: usize, +) -> (StreamHandle, DataPlaneEndpoints) { + let pool = BufferPool::new(pool_size, config.frame_size as usize); + + let (send_cmd_tx, send_cmd_rx) = mpsc::channel(channel_capacity); + let (send_evt_tx, send_evt_rx) = mpsc::channel(channel_capacity); + let (recv_evt_tx, recv_evt_rx) = mpsc::channel(channel_capacity); + let (recv_cmd_tx, recv_cmd_rx) = mpsc::channel(channel_capacity); + + let handle = StreamHandle { + send: SendHalf { + stream_id, + cmd_tx: send_cmd_tx, + evt_rx: send_evt_rx, + pool: pool.clone(), + active_buf: None, + }, + recv: RecvHalf { + stream_id, + evt_rx: recv_evt_rx, + cmd_tx: recv_cmd_tx, + pool: pool.clone(), + active_buf: None, + }, + }; + + let endpoints = DataPlaneEndpoints { + send_cmd_rx, + send_evt_tx, + recv_evt_tx, + recv_cmd_rx, + pool, + }; + + (handle, endpoints) +} diff --git a/crates/streams/src/lib.rs b/crates/streams/src/lib.rs new file mode 100644 index 0000000..258aa65 --- /dev/null +++ b/crates/streams/src/lib.rs @@ -0,0 +1,23 @@ +pub mod accept; +pub mod buffer; +pub mod channel; +pub mod connection; +pub mod ctx_ext; +pub mod data_plane; +pub mod handle; +pub mod manager; +pub mod messages; +pub mod notify; +pub mod types; +pub mod wire; + +pub use buffer::{BufferPool, FrameBuf}; +pub use channel::{RecvCommand, RecvEvent, SendCommand, SendEvent}; +pub use connection::StreamConnectionCache; +pub use ctx_ext::{CtxStreams, RuntimeStreams}; +pub use handle::{create_stream_handle, DataPlaneEndpoints, RecvHalf, SendHalf, StreamHandle}; +pub use manager::{StreamManager, STREAM_MANAGER_NAME}; +pub use messages::{OneShot, StreamManagerMsg, StreamNotification}; +pub use notify::{NotifyFlag, NotifySink, StreamEvent, StreamEventKind}; +pub use types::{ResumeToken, StreamConfig, StreamError, StreamId, StreamMode}; +pub use wire::{StreamHeader, ALPN, MAGIC, VERSION}; diff --git a/crates/streams/src/manager.rs b/crates/streams/src/manager.rs new file mode 100644 index 0000000..04b834d --- /dev/null +++ b/crates/streams/src/manager.rs @@ -0,0 +1,622 @@ +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::connection::StreamConnectionCache; +use crate::data_plane; +use crate::handle::{create_stream_handle, StreamHandle}; +use crate::messages::{OneShot, StreamManagerMsg, StreamNotification}; +use crate::types::{StreamConfig, StreamError, StreamId, StreamMode}; +use crate::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 { + _stream_id: StreamId, + owner: ActorAddress, + _mode: StreamMode, + _remote_node: [u8; 32], +} + +struct PendingIncoming { + _node_id: [u8; 32], + _stream_id: StreamId, + _mode: StreamMode, + config: StreamConfig, + conn: OneShot, +} + +pub struct StreamManager { + streams: HashMap, + pending_incoming: HashMap, + listeners: HashMap>, + #[allow(dead_code)] + conn_cache: StreamConnectionCache, + endpoint: Endpoint, + tokio_handle: tokio::runtime::Handle, + runtime: Arc, + self_addr: Option, +} + +impl StreamManager { + pub fn new( + endpoint: Endpoint, + tokio_handle: tokio::runtime::Handle, + runtime: Arc, + ) -> Self { + StreamManager { + streams: HashMap::new(), + pending_incoming: HashMap::new(), + listeners: HashMap::new(), + conn_cache: StreamConnectionCache::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>, + ) { + match result.take() { + Some(Ok(handle)) => { + self.streams.insert( + stream_id, + StreamState { + _stream_id: stream_id, + owner: reply_to, + _mode: StreamMode::BlobTransfer, + _remote_node: [0; 32], + }, + ); + 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, + ) { + // 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 { + _node_id: node_id, + _stream_id: stream_id, + _mode: mode, + 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>, + ) { + match result.take() { + Some(Ok(handle)) => { + self.streams.insert( + stream_id, + StreamState { + _stream_id: stream_id, + owner: reply_to, + _mode: StreamMode::BlobTransfer, + _remote_node: [0; 32], + }, + ); + 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_insert_with(Vec::new) + .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 { + 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::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::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::channel::SendEvent::Error( + StreamError::BrokenPipe(e.to_string()), + )) + .await; + return; + } + pool.checkin(buf); + stripe_idx = (stripe_idx + 1) % sc; + } + crate::channel::SendCommand::Flush => { + for writer in &mut writers { + let _ = writer.flush().await; + } + } + crate::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 { + // 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::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::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::channel::SendEvent::Error( + StreamError::BrokenPipe(e.to_string()), + )) + .await; + return; + } + pool.checkin(buf); + stripe_idx = (stripe_idx + 1) % sc; + } + crate::channel::SendCommand::Flush => { + for writer in &mut writers { + let _ = writer.flush().await; + } + } + crate::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(()) +} diff --git a/crates/streams/src/messages.rs b/crates/streams/src/messages.rs new file mode 100644 index 0000000..907ae21 --- /dev/null +++ b/crates/streams/src/messages.rs @@ -0,0 +1,176 @@ +use std::fmt; +use std::sync::{Arc, Mutex}; + +use iroh::endpoint::Connection; +use swactor::actor::ActorAddress; + +use crate::handle::StreamHandle; +use crate::types::{StreamConfig, StreamError, StreamId, StreamMode}; + +// ─── OneShot ──────────────────────────────────────────────────────────── + +/// Clone-friendly wrapper for non-Clone data (StreamHandle, Connection). +/// +/// The first `.take()` extracts the value; subsequent clones/takes get `None`. +/// This allows non-Clone payloads to live inside Clone message enums required +/// by the actor system's `Message` trait. +pub struct OneShot(Arc>>); + +impl OneShot { + pub fn new(val: T) -> Self { + OneShot(Arc::new(Mutex::new(Some(val)))) + } + + /// Extract the value. Returns `Some` exactly once; all subsequent calls + /// (including from clones) return `None`. + pub fn take(&self) -> Option { + self.0.lock().unwrap().take() + } +} + +impl Clone for OneShot { + fn clone(&self) -> Self { + OneShot(Arc::clone(&self.0)) + } +} + +impl fmt::Debug for OneShot { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let has_value = self.0.lock().unwrap().is_some(); + write!(f, "OneShot({})", if has_value { "Some" } else { "None" }) + } +} + +// SAFETY: OneShot is Send+Sync because access is guarded by Mutex, +// and Arc provides shared ownership. +unsafe impl Send for OneShot {} +unsafe impl Sync for OneShot {} + +// ─── StreamManagerMsg ─────────────────────────────────────────────────── + +/// Messages sent TO the StreamManager actor. +#[derive(Clone, Debug)] +pub enum StreamManagerMsg { + /// Open a new stream to a remote node. + Open { + target_node: [u8; 32], + mode: StreamMode, + config: StreamConfig, + /// The requesting actor's address; receives StreamNotification. + reply_to: ActorAddress, + }, + /// Accept an offered incoming stream. + Accept { + stream_id: StreamId, + /// Receives StreamNotification::StreamReady. + reply_to: ActorAddress, + }, + /// Reject an offered incoming stream. + Reject { + stream_id: StreamId, + }, + /// Register as a stream listener for a given mode. + Listen { + mode: StreamMode, + /// Receives StreamNotification::StreamOffer. + listener: ActorAddress, + }, + /// Close a stream. + Close { + stream_id: StreamId, + }, + + // -- Internal (from tokio tasks back to StreamManager) -- + /// Incoming connection from the accept bridge task. + IncomingConnection { + node_id: [u8; 32], + stream_id: StreamId, + mode: StreamMode, + config: StreamConfig, + /// QUIC connection for data stripes. + conn: OneShot, + }, + /// Async open task completed. + OpenCompleted { + stream_id: StreamId, + reply_to: ActorAddress, + result: OneShot>, + }, + /// Async accept task completed (data-plane tasks spawned). + AcceptCompleted { + stream_id: StreamId, + reply_to: ActorAddress, + result: OneShot>, + }, +} + +// ─── StreamNotification ───────────────────────────────────────────────── + +/// Notifications sent FROM StreamManager TO user actors. +#[derive(Clone, Debug)] +pub enum StreamNotification { + /// A stream is ready for use (open or accept completed successfully). + StreamReady { + stream_id: StreamId, + handle: OneShot, + }, + /// A remote node is offering a new stream. + StreamOffer { + stream_id: StreamId, + mode: StreamMode, + metadata: Vec, + from_node: [u8; 32], + }, + /// A stream was closed. + StreamClosed { + stream_id: StreamId, + reason: Option, + }, + /// A stream open/accept failed. + StreamFailed { + stream_id: StreamId, + error: StreamError, + }, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn oneshot_take_once_semantics() { + let os = OneShot::new(42u64); + assert_eq!(os.take(), Some(42)); + assert_eq!(os.take(), None); + } + + #[test] + fn oneshot_clone_shares_value() { + let os = OneShot::new("hello".to_string()); + let clone = os.clone(); + // First take from clone succeeds + assert_eq!(clone.take(), Some("hello".to_string())); + // Original now gets None + assert_eq!(os.take(), None); + } + + #[test] + fn oneshot_debug_format() { + let os = OneShot::new(1); + assert_eq!(format!("{os:?}"), "OneShot(Some)"); + os.take(); + assert_eq!(format!("{os:?}"), "OneShot(None)"); + } + + fn assert_message() {} + + #[test] + fn stream_manager_msg_is_message() { + assert_message::(); + } + + #[test] + fn stream_notification_is_message() { + assert_message::(); + } +} diff --git a/crates/streams/src/notify.rs b/crates/streams/src/notify.rs new file mode 100644 index 0000000..bd20a07 --- /dev/null +++ b/crates/streams/src/notify.rs @@ -0,0 +1,172 @@ +use std::sync::atomic::{AtomicU8, Ordering}; +use std::sync::Arc; + +/// Bit positions for notification flags. +pub const DATA_READY: u8 = 0b0000_0001; +pub const WRITE_READY: u8 = 0b0000_0010; +pub const CLOSED: u8 = 0b0000_0100; +pub const ERROR: u8 = 0b0000_1000; + +/// Atomic bitflags for coalescing notifications to an actor. +/// +/// Multiple data-plane tasks may set flags concurrently. The actor clears +/// flags after handling them. If a flag is already set when a task tries +/// to set it, the notification is coalesced (deduplicated). +pub struct NotifyFlag { + flags: AtomicU8, +} + +impl NotifyFlag { + pub fn new() -> Self { + NotifyFlag { + flags: AtomicU8::new(0), + } + } + + /// Set a flag bit. Returns `true` if the bit was previously clear + /// (i.e., this is a new notification that should trigger an inject). + /// Returns `false` if already set (coalesced, no inject needed). + pub fn set(&self, kind: u8) -> bool { + let prev = self.flags.fetch_or(kind, Ordering::AcqRel); + (prev & kind) == 0 + } + + /// Clear a flag bit. Called by the actor after handling. + pub fn clear(&self, kind: u8) { + self.flags.fetch_and(!kind, Ordering::AcqRel); + } + + /// Read all currently-set flags. + pub fn read(&self) -> u8 { + self.flags.load(Ordering::Acquire) + } + + /// Check if a specific flag is set. + pub fn is_set(&self, kind: u8) -> bool { + (self.read() & kind) != 0 + } +} + +impl Default for NotifyFlag { + fn default() -> Self { + Self::new() + } +} + +/// The kind of stream event delivered to an actor. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StreamEventKind { + DataReady, + WriteReady, + Closed, + Error, +} + +/// A lightweight notification message injected into an actor's mailbox. +#[derive(Debug, Clone)] +pub struct StreamEvent { + pub stream_id: crate::types::StreamId, + pub kind: StreamEventKind, +} + +/// Sink that data-plane tasks use to inject notifications into the actor system. +/// +/// Holds the shared `NotifyFlag` for coalescing, and an inject closure +/// that sends a `StreamEvent` into the actor's mailbox when a truly new +/// notification needs to fire. +pub struct NotifySink { + flag: Arc, + inject: Box, + stream_id: crate::types::StreamId, +} + +impl NotifySink { + pub fn new( + flag: Arc, + stream_id: crate::types::StreamId, + inject: impl Fn(StreamEvent) + Send + Sync + 'static, + ) -> Self { + NotifySink { + flag, + inject: Box::new(inject), + stream_id, + } + } + + /// Notify the actor of a stream event. Coalesces duplicate notifications. + pub fn notify(&self, kind_flag: u8, kind: StreamEventKind) { + if self.flag.set(kind_flag) { + (self.inject)(StreamEvent { + stream_id: self.stream_id, + kind, + }); + } + } + + /// Convenience: notify data ready. + pub fn data_ready(&self) { + self.notify(DATA_READY, StreamEventKind::DataReady); + } + + /// Convenience: notify write ready. + pub fn write_ready(&self) { + self.notify(WRITE_READY, StreamEventKind::WriteReady); + } + + /// Convenience: notify closed. + pub fn closed(&self) { + self.notify(CLOSED, StreamEventKind::Closed); + } + + /// Convenience: notify error. + pub fn error(&self) { + self.notify(ERROR, StreamEventKind::Error); + } + + /// Access the shared flag for the actor side to clear bits. + pub fn flag(&self) -> &Arc { + &self.flag + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn set_returns_true_first_time_false_on_duplicate() { + let flag = NotifyFlag::new(); + assert!(flag.set(DATA_READY)); + assert!(!flag.set(DATA_READY)); + } + + #[test] + fn clear_allows_re_notification() { + let flag = NotifyFlag::new(); + assert!(flag.set(DATA_READY)); + flag.clear(DATA_READY); + assert!(flag.set(DATA_READY)); + } + + #[test] + fn independent_flags_do_not_interfere() { + let flag = NotifyFlag::new(); + assert!(flag.set(DATA_READY)); + assert!(flag.set(WRITE_READY)); + assert!(!flag.set(DATA_READY)); // still set + flag.clear(DATA_READY); + assert!(flag.is_set(WRITE_READY)); // unaffected + assert!(!flag.is_set(DATA_READY)); + } + + #[test] + fn read_shows_all_set_flags() { + let flag = NotifyFlag::new(); + flag.set(DATA_READY); + flag.set(ERROR); + let bits = flag.read(); + assert_eq!(bits & DATA_READY, DATA_READY); + assert_eq!(bits & ERROR, ERROR); + assert_eq!(bits & WRITE_READY, 0); + } +} diff --git a/crates/streams/src/types.rs b/crates/streams/src/types.rs new file mode 100644 index 0000000..707b143 --- /dev/null +++ b/crates/streams/src/types.rs @@ -0,0 +1,148 @@ +use std::fmt; + +use serde::{Deserialize, Serialize}; + +/// Unique identifier for a stream, generated randomly. +#[derive(Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct StreamId(pub [u8; 16]); + +impl StreamId { + pub fn new_random() -> Self { + let mut bytes = [0u8; 16]; + getrandom::getrandom(&mut bytes).expect("getrandom failed"); + StreamId(bytes) + } +} + +impl fmt::Debug for StreamId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "StreamId(")?; + for b in &self.0[..4] { + write!(f, "{b:02x}")?; + } + write!(f, "\u{2026})") + } +} + +impl fmt::Display for StreamId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + for b in &self.0[..8] { + write!(f, "{b:02x}")?; + } + write!(f, "\u{2026}") + } +} + +/// Mode of a stream -- what kind of data flows through it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum StreamMode { + BlobTransfer, +} + +/// Configuration for a stream. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct StreamConfig { + /// Number of parallel QUIC stripes for data transfer. + pub stripe_count: u8, + /// Maximum frame payload size in bytes. + pub frame_size: u32, + /// Opaque metadata attached to the stream negotiation. + pub metadata: Vec, +} + +impl Default for StreamConfig { + fn default() -> Self { + StreamConfig { + stripe_count: 4, + frame_size: 256 * 1024, // 256 KB + metadata: Vec::new(), + } + } +} + +/// Errors that can occur during stream operations. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StreamError { + Closed, + BrokenPipe(String), + Disconnected, + BufferExhausted, + InvalidHeader(String), + ChunkVerificationFailed { + chunk_index: usize, + expected: [u8; 32], + actual: [u8; 32], + }, +} + +impl fmt::Display for StreamError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + StreamError::Closed => write!(f, "stream closed"), + StreamError::BrokenPipe(msg) => write!(f, "broken pipe: {msg}"), + StreamError::Disconnected => write!(f, "disconnected"), + StreamError::BufferExhausted => write!(f, "buffer pool exhausted"), + StreamError::InvalidHeader(msg) => write!(f, "invalid header: {msg}"), + StreamError::ChunkVerificationFailed { + chunk_index, + expected, + actual, + } => { + write!(f, "chunk {chunk_index} verification failed: expected ")?; + for b in &expected[..4] { + write!(f, "{b:02x}")?; + } + write!(f, "\u{2026}, got ")?; + for b in &actual[..4] { + write!(f, "{b:02x}")?; + } + write!(f, "\u{2026}") + } + } + } +} + +impl std::error::Error for StreamError {} + +/// Token that allows resuming an interrupted stream transfer. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ResumeToken { + pub stream_id: StreamId, + pub mode: StreamMode, + pub chunks_completed: u64, + pub bytes_transferred: u64, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn stream_id_random_is_unique() { + let a = StreamId::new_random(); + let b = StreamId::new_random(); + assert_ne!(a, b); + } + + #[test] + fn stream_id_debug_shows_prefix() { + let id = StreamId([0xab, 0xcd, 0xef, 0x01, 0x23, 0x45, 0x67, 0x89, 0, 0, 0, 0, 0, 0, 0, 0]); + let dbg = format!("{id:?}"); + assert_eq!(dbg, "StreamId(abcdef01\u{2026})"); + } + + #[test] + fn stream_id_display_shows_8_bytes() { + let id = StreamId([0xab, 0xcd, 0xef, 0x01, 0x23, 0x45, 0x67, 0x89, 0xaa, 0xbb, 0, 0, 0, 0, 0, 0]); + let disp = format!("{id}"); + assert_eq!(disp, "abcdef0123456789\u{2026}"); + } + + #[test] + fn stream_config_defaults() { + let cfg = StreamConfig::default(); + assert_eq!(cfg.stripe_count, 4); + assert_eq!(cfg.frame_size, 256 * 1024); + assert!(cfg.metadata.is_empty()); + } +} diff --git a/crates/streams/src/wire.rs b/crates/streams/src/wire.rs new file mode 100644 index 0000000..ff84c7b --- /dev/null +++ b/crates/streams/src/wire.rs @@ -0,0 +1,306 @@ +use crate::types::{StreamConfig, StreamError, StreamId, StreamMode}; + +/// Magic bytes identifying the swactor stream protocol. +pub const MAGIC: [u8; 2] = [0x53, 0x57]; + +/// Wire protocol version. +pub const VERSION: u8 = 0x01; + +/// ALPN protocol identifier for QUIC negotiation. +pub const ALPN: &[u8] = b"swactor/stream/1"; + +/// Header sent at the beginning of a stream connection. +/// +/// Wire layout: +/// ```text +/// [2B magic] [1B version] [16B stream_id] [1B mode] [1B stripe_count] +/// [4B frame_size] [4B metadata_len] [metadata_len B metadata] +/// ``` +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StreamHeader { + pub stream_id: StreamId, + pub mode: StreamMode, + pub config: StreamConfig, +} + +/// Fixed portion of the header (before variable-length metadata). +const HEADER_FIXED_SIZE: usize = 2 + 1 + 16 + 1 + 1 + 4 + 4; // 29 bytes + +/// Encode a stream header into bytes. +pub fn encode_header(header: &StreamHeader) -> Vec { + let meta_len = header.config.metadata.len() as u32; + let total = HEADER_FIXED_SIZE + header.config.metadata.len(); + let mut buf = Vec::with_capacity(total); + + // Magic + version + buf.extend_from_slice(&MAGIC); + buf.push(VERSION); + + // Stream ID + buf.extend_from_slice(&header.stream_id.0); + + // Mode + let mode_byte = match header.mode { + StreamMode::BlobTransfer => 0x01, + }; + buf.push(mode_byte); + + // Config: stripe_count, frame_size, metadata + buf.push(header.config.stripe_count); + buf.extend_from_slice(&header.config.frame_size.to_be_bytes()); + buf.extend_from_slice(&meta_len.to_be_bytes()); + buf.extend_from_slice(&header.config.metadata); + + buf +} + +/// Decode a stream header from bytes. +pub fn decode_header(data: &[u8]) -> Result { + if data.len() < HEADER_FIXED_SIZE { + return Err(StreamError::InvalidHeader(format!( + "too short: {} bytes, need at least {HEADER_FIXED_SIZE}", + data.len() + ))); + } + + // Magic + if data[0..2] != MAGIC { + return Err(StreamError::InvalidHeader(format!( + "bad magic: [{:#04x}, {:#04x}]", + data[0], data[1] + ))); + } + + // Version + if data[2] != VERSION { + return Err(StreamError::InvalidHeader(format!( + "unsupported version: {}", + data[2] + ))); + } + + // Stream ID + let mut id_bytes = [0u8; 16]; + id_bytes.copy_from_slice(&data[3..19]); + let stream_id = StreamId(id_bytes); + + // Mode + let mode = match data[19] { + 0x01 => StreamMode::BlobTransfer, + other => { + return Err(StreamError::InvalidHeader(format!( + "unknown mode: {other:#04x}" + ))); + } + }; + + // Config + let stripe_count = data[20]; + let frame_size = u32::from_be_bytes([data[21], data[22], data[23], data[24]]); + let meta_len = u32::from_be_bytes([data[25], data[26], data[27], data[28]]) as usize; + + if data.len() < HEADER_FIXED_SIZE + meta_len { + return Err(StreamError::InvalidHeader(format!( + "metadata truncated: have {} bytes after fixed header, need {meta_len}", + data.len() - HEADER_FIXED_SIZE + ))); + } + + let metadata = data[HEADER_FIXED_SIZE..HEADER_FIXED_SIZE + meta_len].to_vec(); + + Ok(StreamHeader { + stream_id, + mode, + config: StreamConfig { + stripe_count, + frame_size, + metadata, + }, + }) +} + +/// Encode a data frame: `[4B payload_len (big-endian)] [payload]`. +/// A payload length of 0 signals end-of-stripe. +pub fn encode_data_frame(payload: &[u8]) -> Vec { + let len = payload.len() as u32; + let mut buf = Vec::with_capacity(4 + payload.len()); + buf.extend_from_slice(&len.to_be_bytes()); + buf.extend_from_slice(payload); + buf +} + +/// End-of-stripe sentinel: a frame with zero-length payload. +pub fn encode_end_of_stripe() -> [u8; 4] { + [0, 0, 0, 0] +} + +/// Result of decoding a data frame from a byte slice. +#[derive(Debug, PartialEq, Eq)] +pub enum DataFrameDecoded<'a> { + /// A data frame with payload. + Data(&'a [u8]), + /// End-of-stripe sentinel. + EndOfStripe, +} + +/// Decode a data frame from a byte slice. +/// Returns the decoded frame and the number of bytes consumed. +pub fn decode_data_frame(data: &[u8]) -> Result<(DataFrameDecoded<'_>, usize), StreamError> { + if data.len() < 4 { + return Err(StreamError::InvalidHeader( + "data frame too short for length prefix".into(), + )); + } + + let len = u32::from_be_bytes([data[0], data[1], data[2], data[3]]) as usize; + + if len == 0 { + return Ok((DataFrameDecoded::EndOfStripe, 4)); + } + + if data.len() < 4 + len { + return Err(StreamError::InvalidHeader(format!( + "data frame truncated: need {len} bytes, have {}", + data.len() - 4 + ))); + } + + Ok((DataFrameDecoded::Data(&data[4..4 + len]), 4 + len)) +} + +#[cfg(test)] +mod tests { + use super::*; + use proptest::prelude::*; + + #[test] + fn header_round_trip_basic() { + let header = StreamHeader { + stream_id: StreamId([1; 16]), + mode: StreamMode::BlobTransfer, + config: StreamConfig { + stripe_count: 4, + frame_size: 262144, + metadata: vec![10, 20, 30], + }, + }; + let encoded = encode_header(&header); + let decoded = decode_header(&encoded).unwrap(); + assert_eq!(header, decoded); + } + + #[test] + fn header_rejects_bad_magic() { + let mut encoded = encode_header(&StreamHeader { + stream_id: StreamId([0; 16]), + mode: StreamMode::BlobTransfer, + config: StreamConfig::default(), + }); + encoded[0] = 0xFF; + assert!(matches!( + decode_header(&encoded), + Err(StreamError::InvalidHeader(_)) + )); + } + + #[test] + fn header_rejects_bad_version() { + let mut encoded = encode_header(&StreamHeader { + stream_id: StreamId([0; 16]), + mode: StreamMode::BlobTransfer, + config: StreamConfig::default(), + }); + encoded[2] = 0xFF; + assert!(matches!( + decode_header(&encoded), + Err(StreamError::InvalidHeader(_)) + )); + } + + #[test] + fn header_rejects_truncated() { + let encoded = encode_header(&StreamHeader { + stream_id: StreamId([0; 16]), + mode: StreamMode::BlobTransfer, + config: StreamConfig { + metadata: vec![1, 2, 3], + ..StreamConfig::default() + }, + }); + // Chop off the metadata + let truncated = &encoded[..HEADER_FIXED_SIZE]; + assert!(matches!( + decode_header(truncated), + Err(StreamError::InvalidHeader(_)) + )); + } + + #[test] + fn data_frame_round_trip() { + let payload = b"hello world"; + let encoded = encode_data_frame(payload); + let (decoded, consumed) = decode_data_frame(&encoded).unwrap(); + assert_eq!(decoded, DataFrameDecoded::Data(b"hello world")); + assert_eq!(consumed, encoded.len()); + } + + #[test] + fn end_of_stripe_sentinel() { + let sentinel = encode_end_of_stripe(); + assert_eq!(sentinel, [0, 0, 0, 0]); + let (decoded, consumed) = decode_data_frame(&sentinel).unwrap(); + assert_eq!(decoded, DataFrameDecoded::EndOfStripe); + assert_eq!(consumed, 4); + } + + #[test] + fn data_frame_rejects_truncated() { + let encoded = encode_data_frame(b"hello"); + // Only give the length prefix + partial payload + let truncated = &encoded[..6]; + assert!(matches!( + decode_data_frame(truncated), + Err(StreamError::InvalidHeader(_)) + )); + } + + proptest! { + #[test] + fn header_round_trip_arbitrary( + id_bytes in prop::array::uniform16(any::()), + stripe_count in 1u8..=16, + frame_size in 1024u32..=1_048_576, + metadata in prop::collection::vec(any::(), 0..256), + ) { + let header = StreamHeader { + stream_id: StreamId(id_bytes), + mode: StreamMode::BlobTransfer, + config: StreamConfig { + stripe_count, + frame_size, + metadata, + }, + }; + let encoded = encode_header(&header); + let decoded = decode_header(&encoded).unwrap(); + prop_assert_eq!(header, decoded); + } + + #[test] + fn data_frame_round_trip_arbitrary( + payload in prop::collection::vec(any::(), 0..262144), + ) { + if payload.is_empty() { + // Empty payload encodes as end-of-stripe + let encoded = encode_data_frame(&payload); + let (decoded, _) = decode_data_frame(&encoded).unwrap(); + prop_assert_eq!(decoded, DataFrameDecoded::EndOfStripe); + } else { + let encoded = encode_data_frame(&payload); + let (decoded, consumed) = decode_data_frame(&encoded).unwrap(); + prop_assert_eq!(decoded, DataFrameDecoded::Data(&payload)); + prop_assert_eq!(consumed, 4 + payload.len()); + } + } + } +} diff --git a/crates/swactor-node/Cargo.toml b/crates/swactor-node/Cargo.toml index 1ebd0a7..7499647 100644 --- a/crates/swactor-node/Cargo.toml +++ b/crates/swactor-node/Cargo.toml @@ -9,6 +9,7 @@ swactor-std = { path = "../std" } dashboard = { path = "../dashboard", features = ["distribution"] } swactor-datastore = { path = "../datastore" } distribution = { path = "../distribution" } +swactor-streams = { path = "../streams" } clap = { version = "4", features = ["derive"] } ctrlc = "3" iroh = { version = "0.96", optional = true } diff --git a/crates/swactor-node/src/main.rs b/crates/swactor-node/src/main.rs index af46d82..e43520c 100644 --- a/crates/swactor-node/src/main.rs +++ b/crates/swactor-node/src/main.rs @@ -713,6 +713,7 @@ fn run_iroh( ) { use distribution::iroh_driver::{IrohDriver, IrohDriverConfig}; use iroh::{RelayMode, SecretKey}; + use swactor_std::RuntimeNaming; // Evaluate relay candidacy and determine embedded relay bind address #[cfg(feature = "relay")] @@ -775,6 +776,7 @@ fn run_iroh( relay_mode, node: node_config, peer_auth: Some(peer_auth.clone()), + additional_alpns: vec![swactor_streams::ALPN.to_vec()], #[cfg(feature = "relay")] embedded_relay_bind, #[cfg(feature = "relay")] @@ -782,6 +784,26 @@ fn run_iroh( }; let mut driver = IrohDriver::new(iroh_config).expect("failed to create iroh driver"); + // Spawn StreamManager actor + let stream_mgr = swactor_streams::StreamManager::new( + driver.endpoint().clone(), + driver.tokio_handle(), + Arc::clone(&handle.runtime), + ); + let stream_mgr_addr = handle + .runtime + .spawn(stream_mgr) + .expect("spawn StreamManager"); + handle + .runtime + .register_name(swactor_streams::STREAM_MANAGER_NAME, stream_mgr_addr) + .expect("register StreamManager"); + + // Wire streams into the datastore + if let Some(group) = ds_group { + group.configure_streams(stream_mgr_addr, driver.tokio_handle()); + } + eprintln!("Node {} started (iroh)", hex(&driver.node_id().0[..4])); // Join seed if provided (accepts hex or base58) @@ -835,6 +857,25 @@ fn run_iroh( driver.recv(); driver.tick(); + // Forward incoming stream connections to StreamManager + for (node_id, conn) in driver.drain_other_connections() { + let rt_clone = Arc::clone(&handle.runtime); + let mgr_addr = stream_mgr_addr; + let node_bytes = node_id.0; + driver.tokio_handle().spawn(async move { + match swactor_streams::accept::handle_incoming( + node_bytes, conn, &rt_clone, mgr_addr, + ) + .await + { + Ok(()) => {} + Err(e) => { + eprintln!("stream accept: failed to handle incoming: {e}"); + } + } + }); + } + // Drain discovered peers (dashboard "Add Peer") and auto-join them { let mut new_peers = Vec::new(); diff --git a/docs/development_history/STREAMS.md b/docs/development_history/STREAMS.md new file mode 100644 index 0000000..1876c0a --- /dev/null +++ b/docs/development_history/STREAMS.md @@ -0,0 +1,475 @@ + Swactor Stream Primitive -- Architectural Design + + Context + + Swactor has a distributed actor runtime with SWIM membership, Kademlia routing, and a content-addressed datastore. The current datastore + transfers blobs one chunk at a time via actor message round-trips -- extremely slow for large objects. Beyond the datastore, the system + needs a general-purpose bulk data transfer primitive for ML workloads (training data, weight checkpoints, gradient exchange), real-time + media (video/voice), and future game state replication. + + The stream primitive is a high-performance data channel between nodes that actors negotiate and manage but do not sit on the data path of. + It should achieve top-class throughput by leveraging QUIC's multiplexed streams directly, bypassing the actor mailbox system for data + transfer. + + Decisions made: + - Data path: StreamHandle with try_read/try_write; actors receive lightweight notification messages but data bypasses mailboxes + - Reliability: Reliable-only MVP; abstraction designed so unreliable (QUIC datagrams) can be added later + - Locality: Cross-node only; same-node actors use regular messages + - Crate: New crates/streams/ crate + + --- + 1. Core Concept: Control Plane vs Data Plane + + The fundamental architecture separates stream management (control plane) from data transfer (data plane). + + Control plane -- actor messages through normal mailboxes: + - Stream negotiation (open, accept, reject) + - Parameter configuration (buffer sizes, chunk sizes, parallelism) + - Lifecycle events (established, closed, error) + - Progress/health notifications + + Data plane -- bypasses actors entirely: + - Raw bytes flow through QUIC streams on the iroh transport + - Managed by async tasks on the IrohDriver's tokio runtime + - Actors interact via StreamHandle objects (try_read/try_write), not mailbox messages + - QUIC's built-in flow control handles backpressure + + CONTROL PLANE (actor messages, mailboxes, worker ticks) + +--------+ StreamOpen +-----------+ StreamAccept +--------+ + | Actor | -----------> | Stream | <------------- | Actor | + | (nodeA)| | Manager | |(nodeB) | + +--------+ +-----------+ +--------+ + | | | + | StreamReady(handle) | | StreamReady(handle) + v v v + DATA PLANE (tokio tasks, QUIC streams, pre-allocated buffers) + +----------+ bytes +----------+ bytes +----------+ + | SendHalf | =========> | QUIC | =========> | RecvHalf | + | (writer) | N parallel| streams | N parallel | (reader) | + +----------+ stripes +----------+ stripes +----------+ + + --- + 2. Stream Identity and Addressing + + StreamId: A 16-byte random identifier, generated by the initiator during negotiation. Deliberately not an ActorAddress -- streams are not + actors, are not placed on workers, and are not discoverable via Kademlia. Keeping them out of the AddressMap avoids polluting the actor + routing hot path. + + Full stream address: The tuple (NodeId, StreamId) is globally unique. A node can host many concurrent streams to many peers. + + ALPN separation: Streams use a new protocol identifier swactor/stream/1, separate from the existing swactor/swim/1 used for membership. + This means: + - The iroh accept loop can distinguish stream connections from protocol messages immediately + - Stream data never blocks or interferes with cluster heartbeats + - Stream connections can have different tuning in the future + + --- + 3. QUIC Stream Utilization + + Parallel Stripes for Blob Transfers + + For a single large transfer, multiple QUIC streams are opened in parallel on the same QUIC connection. Each stream carries a disjoint + range of the data. This is the stripe count, negotiated during handshake (default: 4). + + Why multiple streams? A single QUIC stream can be limited by per-stream receive-window backpressure. Multiple streams allow the sender to + push data into QUIC's send buffer more aggressively, keeping the congestion window filled. Measurements from quinn/s2n-quic show 2-8 + parallel streams can improve throughput 2-4x on high-bandwidth-delay-product links. + + Stream layout per transfer: + - Stream 0 (control stream): Bidirectional QUIC stream. Carries the handshake header and out-of-band signals (completion, cancel, errors, + health). Stays open for the transfer's lifetime. + - Streams 1..N (data stripes): Unidirectional QUIC streams, each carrying sequential chunks. Stripe assignment is round-robin by chunk + index. + + Connection Reuse + + Multiple concurrent streams between the same two nodes share one QUIC connection (on the stream ALPN). QUIC multiplexing handles this + natively. The streams crate maintains a connection cache separate from the SWIM connection cache. + + --- + 4. Wire Format + + Two layers of wire format: the stream-level protocol (negotiation + data framing) and the blob transfer application protocol that rides on + top of it. + + Control Stream Header (stream-level) + + [2B magic: 0x53 0x57] -- "SW" + [1B version: 0x01] + [16B StreamId] + [1B mode] -- 0x01=BlobTransfer, 0x02=ContinuousStream (future) + [1B stripe_count] -- parallel data stripes (1-255) + [4B frame_size (BE u32)] -- maximum frame payload size in bytes + [4B metadata_len (BE u32)] + [N bytes metadata] -- negotiation payload (e.g., ContentHash for blob transfer) + + Data Stripe Frame Format (stream-level) + + [4B frame_len (BE u32)] -- 0 = end-of-stripe + [N bytes payload] -- raw data bytes + + Deliberately minimal. No per-frame type tags (QUIC provides ordered reliable delivery), no per-frame checksums on the wire (QUIC provides + TLS integrity for transport), no per-frame metadata. Every byte of overhead on the hot path costs throughput. + + BlobTransfer Application Protocol + + For blob transfers, the `StreamConfig.metadata` carries the 32-byte `ContentHash` of the requested blob (so the serve side knows what to + send). The actual blob data flows over the StreamHandle with this application-level framing: + + [4B manifest_json_length (u32 BE)] + [N bytes manifest JSON] -- serialized ObjectManifest + [chunk_0 raw bytes] -- size from manifest.chunks[0].size + [chunk_1 raw bytes] -- size from manifest.chunks[1].size + ... + + The receiver knows each chunk's expected size and blake3 hash from the manifest. Each chunk is verified individually on arrival: + blake3(chunk_data) == chunk_ref.hash. Corrupted chunks cause immediate transfer failure. This is implemented by the `send_blob` and + `recv_blob` async functions in `crates/datastore/src/blob_transfer.rs`. + + Note: the blob transfer protocol sends chunks sequentially through the StreamHandle, which distributes data frames across stripes via the + data-plane layer's round-robin. Individual chunks are not split across stripes -- the stripe layer is transparent to the application + protocol. + + --- + 5. Buffering Strategy + + Pre-allocated Sliding Window (Zero Allocation on Hot Path) + + The buffer pool is a sliding window, not a store. It never holds the entire blob in memory -- data flows through it like water through a + pipe. A 1TB transfer uses the same ~4MB of buffer memory as a 1MB transfer; only the duration changes. + + All buffers are allocated during stream setup, not per-frame. + + Sender pipeline (per stripe, double-buffered): + Source (disk/memory/computation) + → [Buffer A: being filled from source] + → [Buffer B: being written to QUIC] + → Buffer B recycled → becomes the next Buffer A + → repeat until source exhausted + One buffer is being filled while the other is being sent. When QUIC accepts Buffer B's bytes, it's recycled and refilled from the source. + The source can be disk I/O, a computation producing data, or anything that yields bytes. + + Receiver pipeline (per stripe, double-buffered): + QUIC recv stream + → [Buffer A: being filled from QUIC] + → [Buffer B: being written to disk/consumed] + → Buffer B recycled → becomes the next Buffer A + → repeat until stream ends + The receiver reads from QUIC into one buffer while the previous buffer is being written to disk (for blob transfer) or consumed by the + application. Buffers are recycled, never allocated mid-transfer. + + Backpressure chain (end-to-end): + Source read speed + → fills sender buffer pool (2 per stripe) + → QUIC congestion window + → network bandwidth + → QUIC receive window + → fills receiver buffer pool (2 per stripe) + → sink write speed (disk I/O, consumer processing) + + If ANY link is slow, pressure propagates backward automatically. + No custom flow control needed -- QUIC handles it. + + Sizing: + - Pool: stripe_count * 2 buffers per side = 8 buffers (at 4 stripes) + - Frame size: 256KB per frame (separate from the datastore's 1MB storage chunk size) + - Total memory per stream direction: 8 x 256KB = 2MB + - Total for a bidirectional transfer: ~4MB, regardless of blob size + - At ~1200 bytes per QUIC packet, 256KB = ~213 packets. Smaller blast radius on packet loss than 1MB, better interleaving across stripes, + aligns with OS page sizes. + + TB-Scale Considerations + + For very large transfers (100GB+ ML weights, TB-scale training data), additional design considerations apply: + + Manifest streaming: At 1MB datastore chunks, a 1TB blob has ~1M chunks. Each ChunkRef is ~40 bytes, so the manifest is ~40MB. This is too + large for a single negotiation payload. The current implementation sends the manifest as a JSON preamble on the data stream itself (not in + the negotiation metadata). For TB-scale, the manifest could be streamed progressively instead of loaded all at once. + + Per-chunk verification on arrival: The receiver verifies each chunk individually as it arrives: blake3(chunk_data) == chunk_ref.hash. + This is implemented in `recv_blob`. A corrupted chunk causes immediate transfer failure. This catches problems early rather than waiting + for full reassembly. + + Progressive resume tokens: Resume tokens are emitted periodically (e.g., every 1000 chunks or every 256MB, whichever comes first), not + just on failure. The sender acknowledges receipt of resume tokens. On connection loss, the receiver persists the latest resume token, and + a new stream can resume from that point. For a 1TB transfer, a resume token with a 1M-bit BitVec is ~125KB -- cheap to exchange. + (Not yet implemented -- the ResumeToken type exists but nothing emits or consumes it.) + + Disk I/O as the bottleneck: For TB-scale over fast networks (10Gbps+), disk I/O often becomes the bottleneck rather than the network. The + buffering strategy handles this naturally: when disk writes slow down, the receiver's buffer pool fills, QUIC backpressure kicks in, and + the sender slows to match. No special handling needed -- the pipeline self-regulates. For maximum disk throughput, the receiver can use + O_DIRECT or memory-mapped writes, but this is an implementation optimization, not an architectural decision. + + Stripe count scaling: For very high bandwidth links, the default 4 stripes may not be enough to saturate the connection. The stripe count + should be configurable up to 16, negotiated during handshake based on the expected transfer size and link characteristics. A heuristic: + min(16, max(4, total_chunks / 1000)) -- more stripes for larger transfers. + + --- + 6. StreamHandle -- The Actor-Facing API + + The StreamHandle is a lightweight, Send (but not Clone) object that actors store in their state. It communicates with the data-plane tokio + tasks via channels internally. + + Writer interface: + - try_write(data: &[u8]) -> Result -- Non-blocking. Returns bytes accepted. + - flush() -- Signal that buffered data should be sent. + - close() -- Graceful close. + + Reader interface: + - try_read(buf: &mut [u8]) -> Result -- Non-blocking. Returns bytes read, 0 if none available. + - has_data() -> bool -- Check if data is available without consuming it. + + BlobTransfer Async Functions + + Rather than a wrapper object, blob transfer uses standalone async functions that run inside tokio tasks (spawned after StreamReady). These + functions loop over try_write/try_read with tokio::task::yield_now() for cooperative scheduling: + + - send_blob(send, manifest, read_chunk) -- Writes the manifest preamble, then calls read_chunk(hash) for each chunk on-demand and writes + it. At most one chunk is in memory at a time on the sender side. The read_chunk callback allows any data source (BlobStore via Inbox, + in-memory, etc). + - recv_blob(recv) -- Reads the manifest, then reads and blake3-verifies each chunk. Returns ReceivedBlob { manifest, chunks }. + - poll_inbox(inbox, timeout) -- Async version of the bridge.rs poll_response pattern. Yields instead of thread::sleep. + + These live in crates/datastore/src/blob_transfer.rs. The key insight: since actors can't await futures, the pattern is for the actor to + receive StreamReady, extract the StreamHandle via OneShot::take(), spawn a tokio task for the I/O loop, then stop itself. The tokio task + sends results back to other actors via runtime.send_to(). + + Why non-blocking? Actor handlers are synchronous (fn handle(&mut self, ctx: &Ctx, msg)). They cannot await futures. The try_read/try_write + pattern fits naturally. The tokio task bridge is the mechanism for async I/O. + + --- + 7. Actor Integration: Negotiation Protocol + + Opening a Stream (Initiator) + + 1. Actor sends a StreamOpen control message (through normal actor mailbox routing) to a StreamManager system actor. Contains: target_node: + NodeId, mode, metadata (e.g., ContentHash + manifest for blob transfer), reply_to: ActorAddress. + 2. StreamManager validates the request, allocates a StreamId, and posts an async task to the tokio runtime that: + - Opens a QUIC connection to the target (stream ALPN) + - Opens the control bidirectional stream + - Sends the stream header + - Waits for accept/reject + 3. On accept: StreamManager sends StreamReady { stream_id, handle: StreamHandle } back to the requesting actor. + + Accepting a Stream (Receiver) + + 1. IrohDriver's accept loop receives connection on stream ALPN. + 2. Reads control stream header, extracts StreamId + mode + metadata. + 3. Sends StreamIncoming actor message to local StreamManager. + 4. StreamManager routes to registered stream acceptors (actors that called StreamListen). + 5. Matching actor receives StreamOffer { stream_id, mode, metadata } in its mailbox. + 6. Actor sends StreamAccept or StreamReject back to StreamManager. + 7. On accept: StreamManager allocates buffers, spawns data-plane tasks, sends StreamReady { handle } to the accepting actor. + + Notification Model (Hybrid) + + Stream data bypasses mailboxes, but actors need to know when data is available: + + - The data-plane tasks inject lightweight StreamEvent sentinel messages into the owning actor's mailbox when state changes: DataReady, + WriteReady, Closed, Error. + - Coalescing: An atomic flag prevents duplicate notifications. Set when notification posted, cleared when actor handles it. A + high-throughput stream generates at most one DataReady per actor tick, not one per frame. + - The actor's handle_any dispatches StreamEvent via downcast (same mechanism as Down and ActorExited today -- no core trait changes + needed). + - Actors can also proactively call handle.try_read() from any handler, not just in response to DataReady. + + --- + 8. The StreamManager Actor + + A system actor spawned alongside the IrohDriver, registered under a well-known name. It is the bridge between the actor world and the + stream data plane. + + Responsibilities: + - Registry of active streams: StreamId -> StreamState + - Handle StreamOpen / StreamAccept / StreamReject / StreamListen / StreamClose messages + - Spawn and supervise data-plane tokio tasks + - Monitor stream-holding actors; clean up streams when actors die + - Expose stream metrics (active streams, throughput, errors) for the dashboard + + Communication with tokio runtime: Uses tokio::sync::mpsc and tokio::sync::oneshot channels. Posts commands to async tasks, receives + results as actor messages (via the Inbox pattern already used by DatastoreBridge). + + --- + 9. Flow Control and Backpressure + + Three layers, all leveraging what QUIC already provides: + + 1. QUIC-level: Per-stream and per-connection flow control (receive window, congestion window). This is the primary mechanism. Not + duplicated. + 2. Buffer pool saturation: When receiver's pre-allocated buffer pool is full, the recv-side tokio task stops reading from QUIC. QUIC's + receive window closes, sender stops transmitting. Natural backpressure without custom protocol. + 3. StreamHandle backpressure: try_write() returns 0 bytes accepted when the send buffer is full. The actor knows to back off or buffer + internally. + + No custom flow control protocol. QUIC's congestion control (Cubic/BBR) is battle-tested. Adding application-level flow control would add + complexity and latency without benefit. + + Cancellation + + - Cooperative: StreamCancel signal on the control stream + - Abrupt: reset() on the QUIC streams + - Nuclear: close the QUIC connection (node shutdown only) + + --- + 10. Error Handling and Recovery + + Failure Modes + + ┌──────────────────┬──────────────────────┬────────────────────────────────────────────────┐ + │ Failure │ Detection │ Behavior │ + ├──────────────────┼──────────────────────┼────────────────────────────────────────────────┤ + │ Frame corruption │ QUIC TLS + checksums │ Automatic retransmit │ + ├──────────────────┼──────────────────────┼────────────────────────────────────────────────┤ + │ Stream reset │ QUIC RST_STREAM │ StreamEvent::Error to owning actor │ + ├──────────────────┼──────────────────────┼────────────────────────────────────────────────┤ + │ Connection loss │ QUIC timeout │ StreamEvent::Error on all streams to that node │ + ├──────────────────┼──────────────────────┼────────────────────────────────────────────────┤ + │ Node death │ SWIM declares Dead │ StreamEvent::Error on all streams to that node │ + ├──────────────────┼──────────────────────┼────────────────────────────────────────────────┤ + │ Owner actor dies │ Worker cleanup phase │ Stream closed, remote side notified │ + └──────────────────┴──────────────────────┴────────────────────────────────────────────────┘ + + Resume Tokens for Blob Transfers (Not Yet Implemented) + + For large transfers, the receiver periodically emits a ResumeToken on the control channel: + + ResumeToken { + stream_id: StreamId, + manifest_hash: ContentHash, + chunks_received: BitVec, -- which chunks confirmed stored + } + + On failure, the initiator can open a new stream with the ResumeToken. The sender skips confirmed chunks. This avoids retransmitting + terabytes when a checkpoint transfer fails near completion. Leverages the existing ObjectManifest/ChunkRef model. + + The ResumeToken type is defined in crates/streams/src/types.rs but emission/consumption logic is deferred to a future stage. + + --- + 11. Integration with Existing Datastore + + The stream primitive adds a parallel transfer path to the datastore. The existing chunk-at-a-time TransferActor is preserved for + compatibility; the new stream path is used when stream support is configured. + + Architecture: + + DOWNLOAD SIDE: SERVE SIDE: + + DatastoreNode StreamListener + │ DownloadViaStream (Incoming = StreamNotification) + │ ctx.spawn(StreamDownloader) │ on StreamOffer → HandleStreamOffer + ▼ ▼ + StreamDownloader DatastoreNode + (Incoming = StreamNotification) │ HandleStreamOffer + │ on_start: Open → StreamManager │ ctx.spawn(StreamServer) + │ StreamReady → tokio task: ▼ + │ recv_blob → verify → write chunks StreamServer + │ send completion to DatastoreNode (Incoming = StreamNotification) + ▼ │ on_start: Accept → StreamManager + DatastoreNode │ StreamReady → tokio task: + │ StreamDownloadComplete │ read manifest from BlobStore (Inbox) + │ persist metadata, reply to caller │ for each chunk: read from BlobStore, + │ write to stream (one at a time) + │ close stream + + Design principles: + - No bridge/shim actors. Stream-facing actors use Incoming = StreamNotification directly. + - No preloading all chunks into memory. Chunks flow on-demand: storage → network. + - DatastoreNode stays simple (fire-and-forget coordination). The stream actors own the full I/O lifecycle. + - After receiving StreamReady, actors spawn tokio tasks for I/O. Tokio tasks communicate results back via runtime.send_to(). + - StreamServer reads chunks on-demand — at most one chunk in memory at a time. + + Flow: + 1. DatastoreNode receives DownloadViaStream { content_hash, source_node, reply_to }. + 2. Spawns a StreamDownloader actor, which sends Open to StreamManager with metadata = content_hash.0 (32 bytes). + 3. Remote StreamListener receives StreamOffer, extracts ContentHash from metadata, sends HandleStreamOffer to DatastoreNode. + 4. Remote DatastoreNode spawns a StreamServer actor, which sends Accept to StreamManager. + 5. StreamServer receives StreamReady, spawns tokio task: reads manifest from BlobStore, then streams each chunk on-demand via send_blob. + 6. StreamDownloader receives StreamReady, spawns tokio task: calls recv_blob, writes chunks to BlobStore (fire-and-forget), notifies + DatastoreNode of completion. + 7. DatastoreNode creates ObjectEntry and persists via MetadataActor, which sends PutOk to the original caller. + + This eliminates the round-trip-per-chunk bottleneck. A 1GB object with 1MB chunks currently requires 1,024 sequential round-trips. With + streams and 4 parallel stripes, the entire blob flows in a single burst limited only by network bandwidth. + + --- + 12. Growth Path + + Phase 1 (MVP): Reliable Ordered Blob Transfer — IMPLEMENTED + + - StreamConfig with BlobTransfer mode only + - New ALPN swactor/stream/1 handler + - StreamOpen/StreamAccept handshake + - Parallel striped data transfer + - StreamHandle with try_read/try_write + - StreamManager actor + - Datastore integration (StreamListener, StreamDownloader, StreamServer actors) + - BlobTransfer application protocol (send_blob/recv_blob with per-chunk blake3 verification) + + Remaining MVP work: + - Two-node integration test (real QUIC, full download flow) + - Resume tokens (ResumeToken type exists, emission/consumption not yet wired) + + Phase 2: Continuous Streams + + - ContinuousStream mode (no total size known) + - Single bidirectional QUIC stream (no striping) + - Variable-sized message frames + - Bounded ring buffer backpressure + - Enables: federated learning gradient streams, data pipelines + + Phase 3: Unreliable Datagrams + + - UnreliableSequenced reliability mode using QUIC datagrams + - Sequence-based frame dropping (latest-wins) + - Receiver-side jitter buffer + - Advisory StreamThrottle on control channel + - Enables: voice/video, game entity state replication + + Phase 4: Priority and QoS + + - priority: u8 in StreamConfig + - Priority-aware write scheduler across concurrent streams + - QUIC stream priority hints + - Per-stream health reporting and dashboard integration + - Enables: simultaneous video + checkpoint without starvation + + Phase 5: Parallel Unordered Transfer + + - ReliableUnordered mode: parallel QUIC streams per chunk, independent delivery + - Configurable parallelism + - Enables: gradient exchange for distributed ML (any chunk consumable independently) + + --- + 13. Key Design Decisions Summary + + ┌───────────────────┬───────────────────────────────────────────────┬────────────────────────────────────────────────────────────┐ + │ Decision │ Choice │ Rationale │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Data path │ StreamHandle bypass, tokio task bridge │ Max throughput; actors manage, don't bottleneck │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Stream identity │ 16-byte StreamId, not ActorAddress │ Streams are not actors; avoid polluting address space │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ ALPN │ Separate swactor/stream/1 │ Isolate from SWIM; no interference with heartbeats │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Parallel stripes │ 4 QUIC streams per blob transfer │ Saturate congestion window on high-BDP links │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Flow control │ QUIC built-in only │ Don't duplicate what the transport does well │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Wire format │ 4-byte length prefix, no type tags │ Minimal per-frame overhead │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Stream chunk size │ 256KB │ Better packet-loss resilience, page-aligned │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Buffering │ Pre-allocated slab per stream │ Zero allocation on hot path │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Blob protocol │ Async functions, not wrapper object │ Simpler; tokio tasks own the I/O loop after StreamReady │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Actor pattern │ Spawn actor → StreamReady → tokio task → stop │ Clean separation; actor negotiates, task does I/O │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Chunk I/O │ On-demand via Inbox polling (poll_inbox) │ At most 1 chunk in memory; no preloading │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Crate │ New crates/streams/ │ Optional, clean dependency graph │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ MVP scope │ Reliable ordered only │ Covers ML + datastore; unreliable added later │ + ├───────────────────┼───────────────────────────────────────────────┼────────────────────────────────────────────────────────────┤ + │ Locality │ Cross-node only │ Focused scope; same-node uses regular messages │ + └───────────────────┴───────────────────────────────────────────────┴────────────────────────────────────────────────────────────┘ diff --git a/docs/development_history/STREAMS_IMPLEMENTATION.md b/docs/development_history/STREAMS_IMPLEMENTATION.md new file mode 100644 index 0000000..14406b3 --- /dev/null +++ b/docs/development_history/STREAMS_IMPLEMENTATION.md @@ -0,0 +1,338 @@ +# Swactor Streams -- Implementation Status + +## What Was Built + +Stages 1-4 are implemented. Stages 1-3 built the stream primitive in `crates/streams/` (the `swactor-streams` crate). Stage 4 connected streams to the datastore so blob transfers use QUIC streams instead of sequential actor-message round-trips. All 42 tests pass (37 streams + 5 blob_transfer). + +### Stage 1: Types, Wire Format, and Buffer Pool + +Pure Rust -- no tokio, no iroh, no network. Compiles and tests in isolation. + +#### `src/types.rs` + +Core domain types for the stream system. + +- **`StreamId([u8; 16])`** -- 16-byte random identifier. `Copy`, `Hash`, `Eq`, `Serialize`/`Deserialize`. Custom `Debug` (4-byte hex prefix) and `Display` (8-byte hex prefix) following the codebase's ID conventions. Not an `ActorAddress` -- streams are not actors and don't pollute the address space. +- **`StreamMode`** -- enum with `BlobTransfer` variant. Extensible for future modes (continuous streams, datagrams). +- **`StreamConfig`** -- negotiation parameters: `stripe_count` (default 4), `frame_size` (default 256KB), `metadata` (opaque bytes for application-level negotiation payloads like ContentHash). +- **`StreamError`** -- error enum covering `Closed`, `BrokenPipe`, `Disconnected`, `BufferExhausted`, `InvalidHeader`, and `ChunkVerificationFailed` (with expected/actual hashes for diagnostics). +- **`ResumeToken`** -- checkpoint for resuming interrupted transfers, carrying stream identity and progress counters. + +#### `src/wire.rs` + +Binary wire format for stream headers and data frames. Pure functions, no I/O. + +- **Constants**: `MAGIC: [0x53, 0x57]` ("SW"), `VERSION: 0x01`, `ALPN: b"swactor/stream/1"`. +- **`StreamHeader`** -- the negotiation header sent at connection establishment. Wire layout: `[2B magic][1B version][16B stream_id][1B mode][1B stripe_count][4B frame_size][4B metadata_len][N metadata]`. +- **`encode_header` / `decode_header`** -- round-trippable serialization with validation (magic, version, mode, truncation checks). +- **Data frame format**: `[4B payload_len (big-endian)][payload]`. Deliberately minimal -- no per-frame type tags or checksums (QUIC provides TLS integrity). A zero-length payload signals end-of-stripe. +- **`encode_data_frame` / `decode_data_frame` / `encode_end_of_stripe`** -- frame-level codec. + +#### `src/buffer.rs` + +Pre-allocated buffer pool for zero-allocation data transfer. + +- **`FrameBuf`** -- a `Box<[u8]>` with read/write cursors. `write(&[u8]) -> usize` fills from the write cursor, `read(&mut [u8]) -> usize` drains from the read cursor. `reset()` zeroes only the cursors (not the data) for fast recycling. `load(&[u8])` replaces content directly. +- **`BufferPool`** -- a fixed-size pool backed by `crossbeam::ArrayQueue` (lock-free MPMC). `checkout() -> Option` and `checkin(buf)` enable concurrent use between actor threads and tokio tasks without locks. `Clone` shares the underlying `Arc` so send/recv sides reference the same pool. + +### Stage 2: StreamHandle, Channels, and Data Plane + +Introduces tokio channels and async tasks but NOT iroh. Data-plane tasks are generic over `AsyncRead`/`AsyncWrite`, fully testable with `tokio::io::DuplexStream`. + +#### `src/channel.rs` + +Typed channel messages that move `FrameBuf`s by ownership (zero-copy handoff). + +- **`SendCommand`** -- `Data(FrameBuf)`, `Flush`, `Close`. Actor -> send task. +- **`SendEvent`** -- `WriteReady`, `Error(StreamError)`, `Closed`. Send task -> actor. +- **`RecvCommand`** -- `Consumed(FrameBuf)`, `Close`. Actor -> recv task. +- **`RecvEvent`** -- `Data(FrameBuf)`, `Error(StreamError)`, `Closed`. Recv task -> actor. + +#### `src/notify.rs` + +Notification coalescing to prevent flooding actor mailboxes. + +- **`NotifyFlag`** -- `AtomicU8` bitflags (`DATA_READY`, `WRITE_READY`, `CLOSED`, `ERROR`). `set(kind) -> bool` returns true only if the bit was previously clear, signaling a new notification should be injected. `clear(kind)` is called by the actor after handling. +- **`StreamEvent`** / **`StreamEventKind`** -- the lightweight sentinel message injected into actor mailboxes. Carries `stream_id` and `kind` (DataReady, WriteReady, Closed, Error). +- **`NotifySink`** -- held by data-plane tasks. Combines the shared `NotifyFlag` with an inject closure. Convenience methods: `data_ready()`, `write_ready()`, `closed()`, `error()`. + +#### `src/handle.rs` + +The actor-facing API for reading and writing stream data. + +- **`SendHalf`** -- owns `mpsc::Sender`, `mpsc::Receiver`, a `BufferPool` clone, and an active `FrameBuf`. `try_write(&[u8]) -> Result` fills the active buffer and sends full buffers via `try_send` (non-blocking). Returns 0 on backpressure. `flush()` sends partial buffers. `close()` flushes remaining data and sends the Close command. +- **`RecvHalf`** -- owns `mpsc::Receiver`, `mpsc::Sender`, a `BufferPool` clone, and an active `FrameBuf`. `try_read(&mut [u8]) -> Result` drains the active buffer then pulls new buffers from the channel. Returns 0 when no data is available. `has_data()` peeks without consuming. +- **`StreamHandle`** -- combines `SendHalf` and `RecvHalf`. `Send` but not `Clone` (the mpsc receivers are not cloneable). +- **`create_stream_handle(stream_id, config, pool_size, channel_capacity)`** -- factory that returns `(StreamHandle, DataPlaneEndpoints)`. The handle goes to the actor; the endpoints go to the data-plane tasks. + +#### `src/data_plane.rs` + +Async tasks that bridge `StreamHandle` channels to actual byte streams. + +- **`send_stripe_task`** -- reads `SendCommand`s from the channel, wire-encodes them as data frames, writes to the transport, returns consumed buffers to the pool, and optionally notifies the actor via `NotifySink`. +- **`recv_stripe_task`** -- reads wire-encoded frames from the transport, loads payloads into `FrameBuf`s from the pool, sends `RecvEvent::Data` to the actor channel. Handles end-of-stripe sentinel and connection closure. +- **`spawn_send_stripes` / `spawn_recv_stripes`** -- spawn a set of stripe tasks from a writer/reader factory. The recv spawner merges all stripe outputs into a single `mpsc::Receiver`. + +Generic over `AsyncRead + AsyncWrite + Send + Unpin + 'static`, so tests use `tokio::io::DuplexStream` with no network stack. + +### Stage 3: QUIC Integration and StreamManager Actor + +Connects the data-plane tasks to real QUIC streams via iroh. Introduces the `StreamManager` system actor with full open/accept/reject lifecycle. Modifies `IrohDriver` for generic ALPN routing and bootstraps the StreamManager in `swactor-node`. + +#### `src/messages.rs` + +Protocol types for the stream control plane. + +- **`OneShot`** -- Clone-friendly wrapper for non-Clone data (`StreamHandle`, `Connection`). Uses `Arc>>` internally. First `.take()` extracts the value; subsequent calls (including from clones) return `None`. This allows non-Clone payloads inside Clone message enums required by the actor system's `Message` trait. +- **`StreamManagerMsg`** -- 8-variant enum for messages sent TO the StreamManager actor: + - `Open { target_node, mode, config, reply_to }` -- Request a new stream to a remote node. + - `Accept { stream_id, reply_to }` -- Accept an offered incoming stream. + - `Reject { stream_id }` -- Reject an offered incoming stream. + - `Listen { mode, listener }` -- Register as a stream listener for a given mode. + - `Close { stream_id }` -- Close a stream. + - `IncomingConnection { node_id, stream_id, mode, config, conn }` -- Internal: from accept bridge to StreamManager. + - `OpenCompleted { stream_id, reply_to, result }` -- Internal: async open task completed. + - `AcceptCompleted { stream_id, reply_to, result }` -- Internal: async accept task completed. +- **`StreamNotification`** -- 4-variant enum for notifications sent FROM StreamManager TO user actors: + - `StreamReady { stream_id, handle }` -- Stream is ready for use (open or accept completed). + - `StreamOffer { stream_id, mode, metadata, from_node }` -- A remote node is offering a stream. + - `StreamClosed { stream_id, reason }` -- A stream was closed. + - `StreamFailed { stream_id, error }` -- A stream open/accept failed. + +#### `src/connection.rs` + +Async connection cache for stream QUIC connections, separate from SWIM connections. + +- **`StreamConnectionCache`** -- `HashMap<[u8; 32], Connection>` with health-check-on-access. `get_or_connect()` checks `conn.close_reason().is_none()` before reuse and falls back to connecting via `endpoint.connect(key, ALPN)`. `prune_closed()` for bulk cleanup. Uses the stream ALPN (`swactor/stream/1`). + +#### `src/manager.rs` + +The core StreamManager system actor. + +- **`StreamManager`** -- implements `ActorInterface`. Manages active streams, pending incoming offers, listener registrations, and a connection cache. Holds an `Endpoint`, `tokio::runtime::Handle`, and `Arc` for spawning async tasks and sending messages back to itself. +- **`STREAM_MANAGER_NAME`** -- well-known name `"StreamManager"` for the name registry. +- **Open flow**: Generates `StreamId`, spawns a tokio task that connects, sends header on a control bi-stream, waits for a 1-byte accept/reject response, then creates `StreamHandle` + data-plane tasks, and sends `OpenCompleted` back to the StreamManager. StreamManager then delivers `StreamNotification::StreamReady` to the requesting actor. +- **Incoming flow**: Accept bridge reads header, sends `IncomingConnection` to StreamManager. StreamManager stores as pending, notifies matching listeners with `StreamOffer`. +- **Accept flow**: Takes pending connection, spawns tokio task that sends accept byte, creates `StreamHandle` + data-plane tasks, sends `AcceptCompleted` back. StreamManager delivers `StreamReady` to accepting actor. +- **Reject flow**: Sends reject byte on a uni-stream, drops the connection. +- **Close flow**: Removes stream state; data-plane tasks terminate when channels drop. +- **`handle_down`**: Cleans up streams owned by dead actors and removes dead listeners. +- **Data-plane spawning**: For each stream direction, a single tokio task opens N uni-streams and round-robins data frames across them. Recv tasks accept incoming uni-streams and dispatch each to a `recv_stripe_task`. + +#### `src/accept.rs` + +Bridge between incoming QUIC connections and the StreamManager actor. + +- **`spawn_accept_bridge`** -- spawns a tokio task that reads from a channel of `(node_id, Connection)` pairs, accepting the control bi-stream, reading the stream header via `read_to_end` + `decode_header`, and forwarding `StreamManagerMsg::IncomingConnection` to the StreamManager via `runtime.send_to()`. +- **`handle_incoming`** -- public async function for per-connection header processing. Can also be called directly from the main loop (used by `swactor-node`). + +#### Modified: `crates/distribution/src/iroh_driver.rs` + +Generic ALPN support to route stream connections separately from SWIM. + +- **`IrohDriverConfig`**: Added `additional_alpns: Vec>` field. All existing call sites updated with `additional_alpns: vec![]`. +- **Endpoint creation**: ALPNs now include both SWIM and any additional ALPNs (`vec![ALPN.to_vec()] + additional_alpns`). +- **Accept loop**: After accepting a connection, checks `conn.alpn()`. SWIM ALPN routes to `accepted_conns` (existing behavior). All other ALPNs route to `other_accepted_conns` (new buffer). +- **New field**: `other_accepted_conns: Arc>>`. +- **New methods**: `endpoint() -> &Endpoint` (for outbound stream connections), `drain_other_connections() -> Vec<(NodeId, Connection)>` (polled from main loop). + +#### Modified: `crates/streams/src/types.rs` + +- Added `Hash` derive to `StreamMode` (needed as `HashMap` key in listeners registry). + +#### Modified: `crates/streams/src/lib.rs` + +- Added module declarations and re-exports for `accept`, `connection`, `manager`, `messages`. +- Re-exports: `StreamConnectionCache`, `StreamManager`, `STREAM_MANAGER_NAME`, `OneShot`, `StreamManagerMsg`, `StreamNotification`. + +#### Modified: `crates/streams/Cargo.toml` + +- Added `swactor-std` dependency (for `CtxMonitoring`, `RuntimeNaming`). +- Added `io-util` feature to `tokio` (for `AsyncWriteExt::flush`). + +#### Modified: `crates/swactor-node/src/main.rs` + +Bootstrap integration in `run_iroh()`. + +- Passes `swactor_streams::ALPN.to_vec()` in `IrohDriverConfig::additional_alpns`. +- After driver creation, spawns `StreamManager::new(endpoint, tokio_handle, runtime)` as a named actor under `"StreamManager"`. +- In the main loop, drains `driver.drain_other_connections()` and spawns `handle_incoming` tasks for each, forwarding to the StreamManager. + +#### Modified: `crates/swactor-node/Cargo.toml` + +- Added `swactor-streams` dependency. + +#### Modified: `crates/distribution/tests/common/iroh.rs`, `crates/dashboard/src/bin/swactor-node.rs` + +- Updated all `IrohDriverConfig` construction sites with `additional_alpns: vec![]`. + +### Stage 4: Datastore Stream Integration + +Connects the stream system to the datastore so blob transfers flow over QUIC streams instead of sequential per-chunk actor-message round-trips. A 1GB blob with 1MB chunks that previously required 1,024 round-trips now flows in a single burst. + +#### `crates/datastore/src/blob_transfer.rs` (NEW) + +Async functions for sending/receiving blobs over StreamHandle. Runs inside tokio tasks, NOT actor handlers. + +- **`BlobTransferError`** -- enum: `IncompleteTransfer(String)`, `ChunkVerificationFailed { expected, actual }`, `InvalidManifest(String)`, `Storage(String)`. +- **`ReceivedBlob`** -- `{ manifest: ObjectManifest, chunks: Vec<(ContentHash, Vec)> }`. +- **`send_blob(send, manifest, read_chunk)`** -- generic over an async callback `F: Fn(ContentHash) -> Future>>`. Writes `[4B manifest_json_len][manifest JSON]` preamble, then for each chunk in the manifest calls `read_chunk(hash)` and writes the raw bytes. Chunks are NOT preloaded -- the callback reads one at a time. +- **`recv_blob(recv)`** -- reads manifest preamble, deserializes JSON, then reads + blake3-verifies each chunk against the manifest's `ChunkRef` entries. Returns `ReceivedBlob`. +- **`poll_inbox(inbox, timeout)`** -- async version of `bridge.rs:poll_response`. Yields (`tokio::task::yield_now`) instead of `thread::sleep`, polling the swactor `Inbox` until a message arrives or timeout. +- **Internal helpers**: `write_all` (loops `try_write` + `yield_now`), `read_exact` (loops `try_read` + `yield_now`). + +Wire format: +``` +[4B manifest_json_length (u32 BE)] +[N bytes manifest JSON] +[chunk_0 raw bytes] <- size from manifest.chunks[0].size +[chunk_1 raw bytes] +... +``` + +#### `crates/datastore/src/actors/stream_listener.rs` (NEW) + +Listens for incoming BlobTransfer stream offers and routes them to DatastoreNode. + +- **`StreamListener`** -- `Incoming = StreamNotification`. State: `datastore_node: ActorAddress`, `stream_manager: Option`. +- `on_start`: looks up `"StreamManager"` via `ctx.where_is()`, sends `StreamManagerMsg::Listen { mode: BlobTransfer }`. +- `handle(StreamOffer)`: extracts 32-byte ContentHash from `metadata`, sends `DatastoreNodeMsg::HandleStreamOffer` to DatastoreNode. Rejects if metadata != 32 bytes. + +#### `crates/datastore/src/actors/stream_downloader.rs` (NEW) + +Opens a stream to a remote node and downloads a blob. + +- **`StreamDownloader`** -- `Incoming = StreamNotification`. Constructor takes: `content_hash`, `source_node`, `datastore_node`, `blob_store`, `reply_to`, `stream_manager`, `tokio_handle`, `runtime`. +- `on_start`: sends `StreamManagerMsg::Open { target_node, mode: BlobTransfer, config.metadata: content_hash.0.to_vec() }`. +- `handle(StreamReady)`: takes handle via `OneShot::take()`, spawns tokio task: + - Calls `recv_blob(&mut recv_half)`. + - Writes each chunk to BlobStore via `runtime.send_to(blob_store, WriteChunk)` (fire-and-forget). + - Writes manifest via `runtime.send_to(blob_store, WriteManifest)` (fire-and-forget). + - Sends `DatastoreNodeMsg::StreamDownloadComplete` to DatastoreNode. + - On error: sends `DatastoreNodeMsg::StreamDownloadFailed`. + - Actor calls `ctx.stop_self()` after spawning the task. +- `handle(StreamFailed)`: sends `StreamDownloadFailed`, stops self. + +#### `crates/datastore/src/actors/stream_server.rs` (NEW) + +Serves a blob to a requesting node over a stream, reading chunks on-demand. + +- **`StreamServer`** -- `Incoming = StreamNotification`. Constructor takes: `stream_id`, `content_hash`, `blob_store`, `stream_manager`, `tokio_handle`, `runtime`. +- `on_start`: sends `StreamManagerMsg::Accept { stream_id }`. +- `handle(StreamReady)`: takes handle, spawns tokio task: + - Reads manifest from BlobStore via `runtime.new_inbox()` + `poll_inbox` (async Inbox polling). + - Calls `send_blob(&mut send_half, &manifest, |chunk_hash| { ... })` with a callback that reads each chunk on-demand from BlobStore via a fresh Inbox. + - At most one chunk is in memory at a time. Chunks flow directly from BlobStore to stream. + - Actor calls `ctx.stop_self()`. +- `handle(StreamFailed)`: stops self. + +#### Modified: `crates/datastore/src/messages.rs` + +Added 5 new variants to `DatastoreNodeMsg`: + +- `DownloadViaStream { content_hash, source_node, reply_to }` -- triggers a stream download. +- `HandleStreamOffer { stream_id, content_hash, from_node, stream_manager }` -- routes incoming stream offers. +- `StreamDownloadComplete { content_hash, manifest, reply_to }` -- download succeeded; persist metadata. +- `StreamDownloadFailed { content_hash, reason, reply_to }` -- download failed; notify caller. +- `ConfigureStreams { stream_manager, tokio_handle, runtime }` -- late-binding stream support. + +Changed from `#[derive(Debug, Clone)]` to `#[derive(Clone)]` with manual `Debug` impl (because `Arc` doesn't implement `Debug`). + +#### Modified: `crates/datastore/src/actors/datastore_node.rs` + +Added stream support fields and handlers to the coordinator actor. + +- **New fields**: `runtime: Option>`, `tokio_handle: Option`, `stream_manager: Option` -- all initialized to `None`. +- **`handle_configure_streams`**: stores runtime/tokio_handle/stream_manager. +- **`handle_download_via_stream`**: spawns `StreamDownloader`. Returns `TransferFailed` if streams not configured. +- **`handle_stream_offer`**: spawns `StreamServer`. +- **`handle_stream_download_complete`**: creates `ObjectEntry`, sends `MetadataMsg::PutObject` to metadata actor with the original `reply_to` for direct response routing. +- **`handle_stream_download_failed`**: sends `DatastoreResponse::TransferFailed` to `reply_to`. + +#### Modified: `crates/datastore/src/actors/mod.rs` + +Added module declarations for `stream_downloader`, `stream_listener`, `stream_server`. + +#### Modified: `crates/datastore/src/lib.rs` + +Added `pub mod blob_transfer`. + +#### Modified: `crates/datastore/src/bridge.rs` + +- Added `datastore_addr: ActorAddress` field to `DatastoreGroup` (stored during `spawn()`). +- Added `configure_streams(&self, stream_manager, tokio_handle)` method: sends `ConfigureStreams` to DatastoreNode, spawns and registers `StreamListener` under `"StreamListener"`. + +#### Modified: `crates/datastore/Cargo.toml` + +- Added `swactor-streams = { path = "../streams" }` and `tokio = { version = "1", features = ["sync", "rt", "time"] }` dependencies. +- Added dev-dependencies for testing: `swactor-streams`, `tokio` with `rt-multi-thread`, `macros`, `io-util`. + +#### Modified: `crates/swactor-node/src/main.rs` + +After StreamManager registration, wires stream support into the datastore: +```rust +if let Some(group) = ds_group { + group.configure_streams(stream_mgr_addr, driver.tokio_handle()); +} +``` + +## Test Coverage + +42 tests across all modules: + +| Category | Tests | What they verify | +|----------|-------|------------------| +| `types` | 4 | StreamId uniqueness, Debug/Display formatting, StreamConfig defaults | +| `wire` | 8 | Header round-trip (basic + property-based), bad magic/version/truncation rejection, data frame round-trip (basic + property-based), end-of-stripe sentinel | +| `buffer` | 7 | FrameBuf write/read/reset/load, BufferPool checkout/checkin/exhaustion/recycling/sharing | +| `notify` | 4 | Set returns true first time / false on duplicate, clear re-enables, independent flags, read shows all bits | +| `data_plane` | 9 | Single-stripe end-to-end transfer, multi-chunk ordered delivery (20 chunks), 4-stripe round-robin (100 chunks), graceful close, notification coalescing, backpressure detection | +| `messages` | 5 | OneShot take-once semantics, clone sharing, debug format, StreamManagerMsg is Message, StreamNotification is Message | +| `blob_transfer` | 5 | Small blob round-trip (single chunk), multi-chunk round-trip (4MB / 256KB chunks / 16 chunks), corrupted chunk detection (blake3 verification), truncated stream detection, property-based arbitrary blob round-trips | + +Property-based tests (via `proptest`) cover: +- Arbitrary stream headers (random IDs, stripe counts 1-16, frame sizes 1KB-1MB, metadata 0-256 bytes) +- Arbitrary data frame payloads (0-256KB) +- Arbitrary blob transfers (random data 1-64KB, chunk sizes 256B-8KB) + +## Dependency Footprint + +### `swactor-streams` crate + +- `swactor` (core actor types, with `serde` feature) +- `swactor-std` (for `CtxMonitoring`, `RuntimeNaming`) +- `shared-types` (ContentHash) +- `distribution` (NodeId, iroh re-exports) +- `crossbeam-queue` (lock-free buffer pool -- already a workspace dep) +- `tokio` (mpsc channels, async I/O traits, io-util) +- `iroh` (QUIC transport, connections, endpoints) +- `blake3`, `serde`, `getrandom` + +Dev dependencies: `proptest`, `tokio` (with rt-multi-thread, macros, test-util, io-util). + +### `swactor-datastore` crate (Stage 4 additions) + +- `swactor-streams` (stream primitives, messages, types) +- `tokio` (sync, rt, time -- for spawning async blob transfer tasks and `poll_inbox`) + +Dev dependencies: `swactor-streams`, `tokio` (with rt-multi-thread, macros, io-util). + +## Next Steps + +### Remaining MVP Work + +These items complete the minimum viable stream-based blob transfer: + +1. **Two-node integration test** -- full open/accept/data-transfer/close cycle with real iroh endpoints and two `DatastoreGroup` instances. Verifies StreamListener receives offers, StreamServer serves blobs, StreamDownloader receives and persists them. This is the critical end-to-end validation that all the pieces work together over real QUIC. + +2. **CtxStreams extension trait** (`crates/streams/src/ctx_ext.rs`) -- convenience methods on `Ctx`: `stream_open()`, `stream_listen()`, `stream_accept()`, `stream_reject()`, `stream_close()`. Looks up `"StreamManager"` via `where_is()` and wraps the message construction. Reduces boilerplate for any actor wanting to use streams. + +3. **Resume tokens** -- checkpoint emission every N chunks or N bytes during `send_blob`/`recv_blob`. Stored in `ResumeToken` (already defined in `types.rs`). On reconnect, receiver sends its token in `StreamConfig.metadata` and sender seeks to the right chunk offset. + +### Post-MVP Phases + +- **Dashboard stream metrics** -- expose active streams, bytes transferred, and transfer rates through the existing dashboard infrastructure. +- **Continuous Streams** -- `ContinuousStream` mode for unbounded data (ML gradient streams, data pipelines). Single bidirectional QUIC stream, variable-sized frames, ring buffer backpressure. +- **Unreliable Datagrams** -- QUIC datagram-based mode for latency-sensitive data (voice/video, game state). Sequence-based dropping, jitter buffer. +- **Priority and QoS** -- per-stream priority, write scheduling across concurrent streams, QUIC stream priority hints. +- **Parallel Unordered Transfer** -- independent per-chunk QUIC streams for workloads where any chunk is consumable independently (distributed ML gradient exchange). -- 2.45.2