diff --git a/apps/myelin/src/job_data_plane.rs b/apps/myelin/src/job_data_plane.rs index 3cdaec7..21ac8ad 100644 --- a/apps/myelin/src/job_data_plane.rs +++ b/apps/myelin/src/job_data_plane.rs @@ -1,9 +1,11 @@ use std::collections::BTreeMap; +use std::os::fd::{FromRawFd, OwnedFd}; use std::sync::Arc; use data_plane::arena::{ArenaConfig, ArenaManager, NodeId as ArenaNodeId}; use data_plane::blob_transfer::{BlobTransferReceiver, BlobTransferSender}; use data_plane::bootstrap::{self, BootstrapSpec, ENV_DATA_PLANE_ENDPOINT, JobHandoff}; +use data_plane::data_plane::{DataPlane, DataPlaneBootstrap}; use data_plane::host::{ HostDataPlaneConfig, HostDataPlaneSessionActor, HostRouteRegistrar, install_session_env, }; @@ -11,6 +13,7 @@ use data_plane::namespace::NamespaceClient; use data_plane::path::JobContext; use data_plane::protocol::JobCapability; use data_plane::source::BlobSourcePublisher; +use data_plane::stream_transport::StreamTransport; use distribution::transport_bridge::{OutboxRouteBinder, RouteBinder, RouteView}; use distribution::types::NodeId; use swactor::actor::ActorAddress; @@ -76,6 +79,7 @@ const ARENA_ALIGNMENT: u64 = 64; pub(crate) struct ActorJobDataPlane { handoff: JobHandoff, host_session: ActorAddress, + capability: JobCapability, runtime: Runtime, } @@ -90,6 +94,7 @@ pub(crate) struct ActorJobDataPlaneConfig { pub(crate) source_sender: Option>, pub(crate) source_publisher: Option>, pub(crate) route_registrar: Option>, + pub(crate) stream_transport: Option>, } impl ActorJobDataPlane { @@ -105,6 +110,7 @@ impl ActorJobDataPlane { source_sender, source_publisher, route_registrar, + stream_transport, } = config; let mut arena = ArenaManager::boot(ArenaConfig { node_id: ArenaNodeId(1), @@ -134,6 +140,7 @@ impl ActorJobDataPlane { source_sender, source_publisher, route_registrar, + stream_transport, }) .map_err(|error| format!("configure host data-plane session: {error}"))?, ) @@ -141,11 +148,31 @@ impl ActorJobDataPlane { install_session_env(&mut handoff, host_session, capability); Ok(Self { handoff, - runtime: runtime.clone(), host_session, + capability, + runtime: runtime.clone(), }) } + pub(crate) fn attach_local(&self) -> Result { + let fd = unsafe { libc::dup(std::os::fd::AsRawFd::as_raw_fd(&self.handoff.arena_fd)) }; + if fd < 0 { + return Err(format!( + "duplicate local data-plane arena: {}", + std::io::Error::last_os_error() + )); + } + let owned = unsafe { OwnedFd::from_raw_fd(fd) }; + futures_lite::future::block_on(DataPlaneBootstrap::attach( + owned, + self.runtime.clone(), + self.host_session, + self.capability, + )) + .map(|bootstrap| bootstrap.data_plane) + .map_err(|error| format!("attach local data-plane client: {error}")) + } + pub(crate) fn configure_run(&self, run_id: String) -> Result<(), String> { futures_lite::future::block_on(async { self.runtime diff --git a/apps/myelin/src/job_deploy.rs b/apps/myelin/src/job_deploy.rs index cceb81b..e402e4f 100644 --- a/apps/myelin/src/job_deploy.rs +++ b/apps/myelin/src/job_deploy.rs @@ -4,6 +4,7 @@ //! in the directory, and exchanges its `EndpointAddr` + actor address //! out-of-band so each side can route to the other over the iroh actor plane. +use crate::data_namespace::{DataNamespaceAuthority, install_namespace_client}; use crate::job_data_plane::{ ActorJobDataPlane, ActorJobDataPlaneConfig, MyelinChildRouteRegistrar, }; @@ -11,7 +12,7 @@ use parking_lot::Mutex; use std::collections::BTreeMap; use std::env; use std::io::{BufRead, BufReader}; -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant}; @@ -21,25 +22,25 @@ use swactor::actor::{ActorAddress, ActorInterface}; use swactor::runtime::{Ctx, ExternalSender}; use swactor_engine::{ActorCompletion, Engine, EngineHandle, TokioBackend, TokioConfig}; use swactor_job_runner::{ - INFERENCE_RESULTS_EDGE_ID, Job, JobDataPlanePort, JobDone, NodeJobActor, OUTPUTS_EDGE_ID, - OrchestratorJobActor, OrchestratorJobMsg, WORKSPACE_EDGE_ID, register_job_codecs, + Job, JobDataPlanePort, JobDone, NodeJobActor, OUTPUTS_EDGE_ID, OrchestratorJobActor, + OrchestratorJobMsg, WORKSPACE_EDGE_ID, register_job_codecs, }; use swactor_transport::hex_encode; use data_plane::blob_transfer::{BlobTransferReceiver, BlobTransferSender}; +use data_plane::data_plane::StreamConsumer; use data_plane::edge_wire::WireEvent; use data_plane::namespace::NamespaceClient; use data_plane::path::{DataPath, JobContext}; +use data_plane::protocol::DataPlaneError; use data_plane::protocol::{JobCapability, register_data_plane_codecs}; use data_plane::source::BlobSourcePublisher; use distribution::node::DistributedNodeConfig; use iroh::{EndpointAddr, RelayMode}; use iroh_driver::{ - EDGE_ALPN, EdgeConnector, EdgeSendHandle, EndpointAddrMask, IrohDriver, IrohDriverConfig, - MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint, + EDGE_ALPN, EndpointAddrMask, IrohDriver, IrohDriverConfig, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, + advertised_endpoint, }; -use tokio::io::AsyncReadExt; -use tokio::sync::Notify; use crate::orchestration::distribution_stack::DistributionRuntimeStack; @@ -50,57 +51,37 @@ const RELAY_WAIT_DEADLINE: Duration = Duration::from_secs(30); const MYELIN_IROH_RELAY_MODE_ENV: &str = "MYELIN_IROH_RELAY_MODE"; const MYELIN_IROH_RELAY_URL_ENV: &str = "MYELIN_IROH_RELAY_URL"; const SWACTOR_IROH_RELAY_URL_ENV: &str = "SWACTOR_IROH_RELAY_URL"; -const JOB_OUTPUT_SOCKET: &str = "inference-results.sock"; -const DATA_PLANE_CONNECT_DEADLINE: Duration = Duration::from_secs(30); const JOB_ARENA_BYTES: u64 = 1 << 20; -#[derive(Clone, Debug, Serialize, Deserialize)] -struct EmbeddedDataPlaneAssignment { - result_endpoint: EndpointAddr, -} - -/// Actor-driven finite-blob ingress plus the retained temporary Unix output -/// stream bridge. Remote bytes remain on `EDGE_ALPN`. +/// Actor-driven finite-blob ingress and namespace-addressed result streams. #[derive(Clone)] pub(crate) struct EmbeddedJobDataPlane { - result_sink: Arc>>, - result_ready: Arc, - connector: EdgeConnector, actor_plane: Arc, host_endpoint_json: String, - output_path: PathBuf, } -pub(crate) struct EmbeddedJobDataPlaneConfig<'a> { - pub(crate) engine: EngineHandle, - pub(crate) connector: EdgeConnector, - pub(crate) root: &'a Path, +pub(crate) struct EmbeddedJobDataPlaneConfig { pub(crate) host_endpoint: EndpointAddr, pub(crate) namespace: NamespaceClient, pub(crate) transfer_receiver: Arc, pub(crate) source_sender: Arc, pub(crate) source_publisher: Arc, + pub(crate) stream_transport: Arc, } impl EmbeddedJobDataPlane { pub(crate) fn start( stack: &DistributionRuntimeStack, - config: EmbeddedJobDataPlaneConfig<'_>, + config: EmbeddedJobDataPlaneConfig, ) -> Result { let EmbeddedJobDataPlaneConfig { - engine, - connector, - root, host_endpoint, namespace, transfer_receiver, source_sender, source_publisher, + stream_transport, } = config; - std::fs::create_dir_all(root) - .map_err(|error| format!("create job data-plane root {}: {error}", root.display()))?; - let output_path = root.join(JOB_OUTPUT_SOCKET); - remove_stale_socket(&output_path)?; let capability = JobCapability::new(ActorAddress::new_random().0); let route_registrar = Arc::new(MyelinChildRouteRegistrar::new( stack.route_view.clone(), @@ -124,49 +105,13 @@ impl EmbeddedJobDataPlane { source_sender: Some(source_sender), source_publisher: Some(source_publisher), route_registrar: Some(route_registrar), + stream_transport: Some(stream_transport), }, )?); let host_endpoint_json = serde_json::to_string(&host_endpoint) .map_err(|error| format!("serialize host data-plane endpoint: {error}"))?; - let result_sink: Arc>> = Arc::new(Mutex::new(None)); - let result_ready = Arc::new(Notify::new()); - - let output_slot = Arc::clone(&result_sink); - let output_ready = Arc::clone(&result_ready); - swactor_process::spawn_unix_stream_listener(engine, &output_path, move |mut stream| { - let output_slot = Arc::clone(&output_slot); - let output_ready = Arc::clone(&output_ready); - async move { - let sink = loop { - let notified = output_ready.notified(); - if let Some(sink) = output_slot.lock().take() { - break sink; - } - notified.await; - }; - let mut bytes = vec![0_u8; 64 * 1024]; - loop { - match stream.read(&mut bytes).await { - Ok(0) => break, - Ok(count) => { - if sink.send(bytes[..count].to_vec()).is_err() { - break; - } - } - Err(_) => break, - } - } - drop(sink); - } - }) - .map_err(|error| format!("bind job data-plane output: {error}"))?; - Ok(Self { - result_sink, - result_ready, - connector, - output_path, actor_plane, host_endpoint_json, }) @@ -177,43 +122,17 @@ impl JobDataPlanePort for EmbeddedJobDataPlane { fn configure( &self, job_id: u64, - result_peer: &str, + _result_peer: &str, ) -> Result, String> { - let assignment = serde_json::from_str::(result_peer) - .map_err(|error| format!("parse job data-plane assignment: {error}"))?; - let sink = self.connector.connect( - assignment.result_endpoint, - INFERENCE_RESULTS_EDGE_ID, - DATA_PLANE_CONNECT_DEADLINE, - )?; - *self.result_sink.lock() = Some(sink); - self.result_ready.notify_one(); self.actor_plane.configure_run(job_id.to_string())?; - let mut env = self.actor_plane.handoff_env(&self.host_endpoint_json); - env.insert( - "SWACTOR_DATA_PLANE_OUTPUT".to_owned(), - self.output_path.to_string_lossy().into_owned(), - ); - Ok(env) + Ok(self.actor_plane.handoff_env(&self.host_endpoint_json)) } fn session_ended(&self, _job_id: u64) { - self.result_sink.lock().take(); self.actor_plane.close(); } } -fn remove_stale_socket(path: &Path) -> Result<(), String> { - match std::fs::remove_file(path) { - Ok(()) => Ok(()), - Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), - Err(error) => Err(format!( - "remove stale job data-plane socket {}: {error}", - path.display() - )), - } -} - /// Out-of-band identity one side publishes so the other can route to it. #[derive(Clone, Debug, Serialize, Deserialize)] pub struct NodeIdentity { @@ -221,6 +140,17 @@ pub struct NodeIdentity { pub actor_hex: String, } +struct OrchestratorResultSink { + bytes: Arc>>, +} + +impl StreamConsumer for OrchestratorResultSink { + fn consume(&self, bytes: &[u8]) -> Result<(), String> { + self.bytes.lock().extend_from_slice(bytes); + Ok(()) + } +} + pub(crate) struct JobOrchestratorSession { _engine: Engine, driver: IrohDriver, @@ -229,6 +159,10 @@ pub(crate) struct JobOrchestratorSession { orch: ActorAddress, identity: NodeIdentity, landing: PathBuf, + _namespace_authority: DataNamespaceAuthority, + _result_plane: Arc, + _result_completion: ActorCompletion>, + _result_bytes: Arc>>, } type JobComposition = (Engine, IrohDriver, DistributionRuntimeStack); @@ -611,6 +545,47 @@ fn start_orchestrator_mode( Some(relay_mode) => build_composition_with_relay(relay_mode)?, None => build_composition()?, }; + let namespace_authority = DataNamespaceAuthority::start( + &stack, + &driver, + landing.join(".swactor-data-namespace.json"), + )?; + let namespace = install_namespace_client(&stack, &driver)?; + let result_capability = JobCapability::new(ActorAddress::new_random().0); + let result_transport: Arc = + driver.stream_transport(); + let result_plane = Arc::new(ActorJobDataPlane::new( + &stack.runtime, + ActorJobDataPlaneConfig { + arena_bytes: JOB_ARENA_BYTES, + arena_generation: 2, + session_generation: 2, + capability: result_capability, + job_context: JobContext { + run_id: "unconfigured".to_owned(), + read_prefixes: vec![DataPath::parse("/runs").expect("static run prefix")], + write_prefixes: Vec::new(), + }, + namespace: Some(namespace.client), + transfer_receiver: None, + source_sender: None, + source_publisher: Some(namespace.source_publisher), + route_registrar: None, + stream_transport: Some(result_transport), + }, + )?); + result_plane.configure_run("0".to_owned())?; + let result_client = result_plane.attach_local()?; + let result_bytes = Arc::new(Mutex::new(Vec::new())); + let result_consumer: Arc = Arc::new(OrchestratorResultSink { + bytes: Arc::clone(&result_bytes), + }); + let result_completion = result_client + .collect_stream( + DataPath::parse("/runs/0/results/inference").expect("static result path"), + result_consumer, + ) + .map_err(|error| format!("register inference result sink: {error}"))?; let done = stack .runtime .new_inbox::() @@ -630,6 +605,10 @@ fn start_orchestrator_mode( orch, identity, landing, + _namespace_authority: namespace_authority, + _result_plane: result_plane, + _result_completion: result_completion, + _result_bytes: result_bytes, }) } diff --git a/apps/myelin/src/node/worker_node_runtime.rs b/apps/myelin/src/node/worker_node_runtime.rs index dbcaa51..2595886 100644 --- a/apps/myelin/src/node/worker_node_runtime.rs +++ b/apps/myelin/src/node/worker_node_runtime.rs @@ -1983,14 +1983,12 @@ fn run() -> Result<(), String> { let data_plane = EmbeddedJobDataPlane::start( &stack, crate::job_deploy::EmbeddedJobDataPlaneConfig { - engine: engine.handle(), - connector: driver.edge_connector(), - root: &workdir, host_endpoint: driver.endpoint_addr(), namespace: namespace.client, transfer_receiver, source_sender, source_publisher: namespace.source_publisher, + stream_transport: driver.stream_transport(), }, )?; let job_route_registrar = Arc::new(MyelinChildRouteRegistrar::new( diff --git a/apps/myelin/src/orchestration/manual_control.rs b/apps/myelin/src/orchestration/manual_control.rs index f1ba82c..c6f7c14 100644 --- a/apps/myelin/src/orchestration/manual_control.rs +++ b/apps/myelin/src/orchestration/manual_control.rs @@ -3441,8 +3441,8 @@ mod tests { proptest! { #![proptest_config(ProptestConfig { - cases: 128, - max_shrink_iters: 2_000, + cases: 16, + max_shrink_iters: 256, ..ProptestConfig::default() })] diff --git a/apps/myelin/src/tests/data_namespace_guarantees.rs b/apps/myelin/src/tests/data_namespace_guarantees.rs index 4a50636..db1ac75 100644 --- a/apps/myelin/src/tests/data_namespace_guarantees.rs +++ b/apps/myelin/src/tests/data_namespace_guarantees.rs @@ -99,6 +99,7 @@ impl DataNode { source_sender: Some(Arc::clone(&self.sender)), source_publisher: Some(Arc::clone(&self.namespace.source_publisher)), route_registrar: None, + stream_transport: None, }) .expect("host session"), ) diff --git a/apps/myelin/src/tests/job_data_plane_guarantees.rs b/apps/myelin/src/tests/job_data_plane_guarantees.rs index 421e60c..aba0225 100755 --- a/apps/myelin/src/tests/job_data_plane_guarantees.rs +++ b/apps/myelin/src/tests/job_data_plane_guarantees.rs @@ -1,14 +1,26 @@ #![cfg(target_os = "linux")] use std::collections::BTreeSet; +use std::sync::Arc; +use std::time::Duration; use data_plane::bootstrap::{ ENV_ARENA_FD, ENV_DATA_PLANE_ACTOR, ENV_DATA_PLANE_ENDPOINT, ENV_JOB_CAPABILITY, }; +use data_plane::data_plane::StreamConsumer; +use data_plane::namespace::{ + DataDirectoryActor, NamespaceClient, NamespaceClientActor, NamespaceDiscovery, +}; use data_plane::path::{DataPath, JobContext}; use data_plane::protocol::JobCapability; +use data_plane::source::BlobSourcePublisher; +use data_plane::stream_transport::{LocalStreamTransport, StreamTransport}; +use futures_lite::future; +use parking_lot::Mutex; +use swactor::actor::ActorAddress; use swactor::config::RuntimeConfig; use swactor::runtime::RuntimeParts; +use swactor_engine::{Engine, TokioBackend, TokioConfig}; use crate::job_data_plane::ActorJobDataPlane; @@ -39,6 +51,7 @@ fn plane() -> ActorJobDataPlane { source_sender: None, source_publisher: None, route_registrar: None, + stream_transport: None, }, ) .expect("actor data-plane") @@ -69,3 +82,121 @@ fn handoff_contains_one_descriptor_and_private_actor_metadata() { assert!(flags >= 0); assert_eq!(flags & libc::FD_CLOEXEC, 0); } + +struct StaticDiscovery(ActorAddress); + +impl NamespaceDiscovery for StaticDiscovery { + fn current_directory(&self) -> Option { + Some(self.0) + } +} + +struct LocalPublisher; + +impl BlobSourcePublisher for LocalPublisher { + fn publish_source(&self, _source: ActorAddress) -> Result<(), String> { + Ok(()) + } +} + +struct BytesConsumer(Arc>>); + +impl StreamConsumer for BytesConsumer { + fn consume(&self, bytes: &[u8]) -> Result<(), String> { + self.0.lock().extend_from_slice(bytes); + Ok(()) + } +} + +#[test] +fn canonical_inference_result_uses_native_stream_end_to_end() { + let state_root = std::env::temp_dir().join(format!( + "myelin-stream-test-{}", + ActorAddress::new_random().to_full_hex() + )); + let parts = RuntimeParts::new(RuntimeConfig::default()); + let runtime = parts.runtime().clone(); + let engine = Engine::new( + parts, + TokioBackend::new(TokioConfig::default()).expect("tokio backend"), + ) + .expect("engine"); + let directory = runtime + .spawn( + DataDirectoryActor::recover(state_root.join("namespace.json"), |_recovery, _length| { + Err(data_plane::namespace::NamespaceError::SourceRecovery( + "no recovered blob sources".to_owned(), + )) + }) + .expect("directory"), + ) + .expect("spawn directory"); + let proxy = runtime + .spawn(NamespaceClientActor::new( + engine.handle(), + runtime.create_sender(), + Arc::new(StaticDiscovery(directory)), + Duration::from_millis(5), + )) + .expect("namespace proxy"); + let namespace = NamespaceClient::new(runtime.clone(), proxy); + let transport: Arc = Arc::new(LocalStreamTransport::new()); + let publisher: Arc = Arc::new(LocalPublisher); + + let make_plane = |read: bool| { + ActorJobDataPlane::new( + &runtime, + crate::job_data_plane::ActorJobDataPlaneConfig { + arena_bytes: 1 << 20, + arena_generation: if read { 21 } else { 22 }, + session_generation: if read { 31 } else { 32 }, + capability: CAPABILITY, + job_context: JobContext { + run_id: "0".to_owned(), + read_prefixes: if read { + vec![path("/runs/0/results")] + } else { + Vec::new() + }, + write_prefixes: if read { + Vec::new() + } else { + vec![path("/runs/0/results")] + }, + }, + namespace: Some(namespace.clone()), + transfer_receiver: None, + source_sender: None, + source_publisher: Some(Arc::clone(&publisher)), + route_registrar: None, + stream_transport: Some(Arc::clone(&transport)), + }, + ) + .expect("actor plane") + }; + let reader_plane = make_plane(true); + let writer_plane = make_plane(false); + let reader = reader_plane.attach_local().expect("reader attachment"); + let writer = writer_plane.attach_local().expect("writer attachment"); + let logical = path("/runs/0/results/inference"); + let observed = Arc::new(Mutex::new(Vec::new())); + let consumer: Arc = Arc::new(BytesConsumer(Arc::clone(&observed))); + let completed = reader + .collect_stream(logical.clone(), consumer) + .expect("register result sink"); + + future::block_on(async { + let mut writer = writer.write_stream(&logical).await.expect("open writer"); + writer + .write(br#"{"device":"CUDA:0","output":[2.75,-8.75]}"#) + .await + .expect("write result"); + writer.close().await.expect("close writer"); + }); + completed.wait().expect("result sink completed"); + assert_eq!( + &*observed.lock(), + br#"{"device":"CUDA:0","output":[2.75,-8.75]}"# + ); + let _ = std::fs::remove_dir_all(state_root); +} diff --git a/crates/bindings/python/src/job.rs b/crates/bindings/python/src/job.rs index 3ee64c8..919ffe3 100644 --- a/crates/bindings/python/src/job.rs +++ b/crates/bindings/python/src/job.rs @@ -2,12 +2,15 @@ use std::collections::HashMap; use std::ffi::{CString, c_int, c_void}; use std::os::fd::{FromRawFd, OwnedFd, RawFd}; use std::ptr; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, RwLock}; use std::time::Duration; use data_plane::blob::{Blob, BlobView, ContentDigest, WritableArenaView}; use data_plane::bootstrap as dp_bootstrap; -use data_plane::data_plane::{BlobWriter, DataPlane, DataPlaneBootstrap, parse_actor_address}; +use data_plane::data_plane::{ + BlobWriter, DataPlane, DataPlaneBootstrap, StreamReader, StreamWriter, parse_actor_address, +}; use data_plane::path::DataPath; use data_plane::protocol::{ BlobFailure, DataPlaneError, JobCapability, register_data_plane_codecs, @@ -23,24 +26,22 @@ use parking_lot::Mutex as ParkingMutex; use pyo3::exceptions::{PyBufferError, PyPermissionError, PyRuntimeError}; use pyo3::ffi; use pyo3::prelude::*; -use pyo3::types::{PyAny, PyModule}; +use pyo3::types::{PyAny, PyBytes, PyModule}; use swactor::actor::{ActorAddress, ActorInterface, Ctx}; use swactor::config::RuntimeConfig; use swactor::runtime::{Runtime, RuntimeParts}; use swactor_engine::{ActorCompletion, Engine, EngineHandle, TokioBackend, TokioConfig}; use swactor_transport::{CodecRegistry, CodecRemoteSink, TransportRouter}; -use tokio::io::{AsyncWriteExt, BufWriter}; -use tokio::net::UnixStream; const ROUTE_POLL: Duration = Duration::from_millis(5); const ROUTE_DEADLINE: Duration = Duration::from_secs(5); -const LEGACY_OUTPUT_ENV: &str = "SWACTOR_DATA_PLANE_OUTPUT"; pyo3::create_exception!(swactor, SwactorError, pyo3::exceptions::PyException); pyo3::create_exception!(swactor, BootstrapError, SwactorError); pyo3::create_exception!(swactor, DataPathError, SwactorError); pyo3::create_exception!(swactor, BlobError, SwactorError); pyo3::create_exception!(swactor, SessionError, SwactorError); +pyo3::create_exception!(swactor, StreamError, SwactorError); fn bootstrap_error(message: impl Into) -> PyErr { PyErr::new::(message.into()) @@ -56,6 +57,11 @@ fn data_plane_error(error: DataPlaneError) -> PyErr { PyErr::new::(format!("data path not found: {path}")) } DataPlaneError::Blob(reason) => PyErr::new::(format!("{reason:?}")), + DataPlaneError::WrongEntryType { .. } + | DataPlaneError::PathReplaced(_) + | DataPlaneError::PeerLost + | DataPlaneError::StreamFault(_) + | DataPlaneError::StreamClosed => PyErr::new::(error.to_string()), DataPlaneError::Attachment(reason) => { PyErr::new::(format!("attachment failed: {reason:?}")) } @@ -229,7 +235,6 @@ fn build_child_routing( pub struct PyDataPlane { inner: Arc, _routing: Arc, - legacy_output: Option, } #[pymethods] @@ -257,20 +262,34 @@ impl PyDataPlane { }) } - fn read_stream(&self, _path: String) -> PyResult<()> { - Err(PyRuntimeError::new_err( - "actor-driven stream reads are not installed", - )) + fn read_stream<'py>(&self, py: Python<'py>, path: String) -> PyResult> { + let path = DataPath::parse(path) + .map_err(|error| PyErr::new::(error.to_string()))?; + let data_plane = self.inner.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let reader = data_plane + .read_stream(&path) + .await + .map_err(data_plane_error)?; + Python::with_gil(|py| { + Py::new( + py, + PyStreamReader { + reader: Arc::new(tokio::sync::Mutex::new(reader)), + }, + ) + }) + }) } fn write_stream(&self, path: String) -> PyResult { - DataPath::parse(path).map_err(|error| PyErr::new::(error.to_string()))?; - let socket = self.legacy_output.clone().ok_or_else(|| { - PyRuntimeError::new_err("temporary deployment stream bridge is not configured") - })?; + let path = DataPath::parse(path) + .map_err(|error| PyErr::new::(error.to_string()))?; Ok(PyStreamContext { - socket, + data_plane: self.inner.clone(), + path, stream: Arc::new(tokio::sync::Mutex::new(None)), + entered: Arc::new(AtomicBool::new(false)), }) } } @@ -618,9 +637,25 @@ unsafe fn release_buffer_format(view: *mut ffi::Py_buffer) { } } +#[pyclass(name = "StreamReader")] +pub struct PyStreamReader { + reader: Arc>, +} + +#[pymethods] +impl PyStreamReader { + fn read<'py>(&self, py: Python<'py>) -> PyResult> { + let reader = self.reader.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let bytes = reader.lock().await.read().await.map_err(data_plane_error)?; + Python::with_gil(|py| Ok(bytes.map(|bytes| PyBytes::new(py, &bytes).unbind()))) + }) + } +} + #[pyclass(name = "StreamWriter")] pub struct PyStreamWriter { - stream: Arc>>>, + stream: Arc>>, } #[pymethods] @@ -629,56 +664,67 @@ impl PyStreamWriter { let stream = self.stream.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { let mut stream = stream.lock().await; - let stream = stream + let writer = stream .as_mut() .ok_or_else(|| PyRuntimeError::new_err("stream writer is closed"))?; - stream - .write_all(&bytes) - .await - .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; - Ok(()) + writer.write(&bytes).await.map_err(data_plane_error) }) } } #[pyclass(name = "_StreamWriteContext")] pub struct PyStreamContext { - socket: String, - stream: Arc>>>, + data_plane: Arc, + path: DataPath, + stream: Arc>>, + entered: Arc, } #[pymethods] impl PyStreamContext { fn __aenter__<'py>(&self, py: Python<'py>) -> PyResult> { - let socket = self.socket.clone(); + if self.entered.swap(true, Ordering::AcqRel) { + return Err(PyRuntimeError::new_err( + "stream write context cannot be entered twice", + )); + } + let data_plane = self.data_plane.clone(); + let path = self.path.clone(); let state = self.stream.clone(); + let entered = self.entered.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - let stream = UnixStream::connect(&socket).await.map_err(|error| { - PyRuntimeError::new_err(format!("connect output stream: {error}")) - })?; - *state.lock().await = Some(BufWriter::new(stream)); - Python::with_gil(|py| Py::new(py, PyStreamWriter { stream: state })) + match data_plane.write_stream(&path).await { + Ok(writer) => { + *state.lock().await = Some(writer); + Python::with_gil(|py| Py::new(py, PyStreamWriter { stream: state })) + } + Err(error) => { + entered.store(false, Ordering::Release); + Err(data_plane_error(error)) + } + } }) } fn __aexit__<'py>( &self, py: Python<'py>, - _exception_type: &Bound<'_, PyAny>, + exception_type: &Bound<'_, PyAny>, _exception: &Bound<'_, PyAny>, _traceback: &Bound<'_, PyAny>, ) -> PyResult> { + let clean = exception_type.is_none(); let state = self.stream.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { - if let Some(mut stream) = state.lock().await.take() { - stream - .flush() - .await - .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; - stream - .shutdown() - .await - .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + let mut writer = state + .lock() + .await + .take() + .ok_or_else(|| PyRuntimeError::new_err("stream write context is not active"))?; + if clean { + writer.close().await.map_err(data_plane_error)?; + } else { + writer.abort().map_err(data_plane_error)?; } Ok(false) }) @@ -736,7 +782,6 @@ fn run(py: Python<'_>, main: Bound<'_, PyAny>) -> PyResult<()> { PyDataPlane { inner: Arc::new(bootstrap.data_plane), _routing: Arc::new(routing), - legacy_output: std::env::var(LEGACY_OUTPUT_ENV).ok(), }, )?; let context = Py::new(py, PyContext { data })?; @@ -848,6 +893,19 @@ impl data_plane::source::BlobSourcePublisher for DebugSourceRegistrar { } } +#[cfg(all(debug_assertions, target_os = "linux"))] +struct DebugStreamCaptureConsumer { + bytes: Arc>>, +} + +#[cfg(all(debug_assertions, target_os = "linux"))] +impl data_plane::data_plane::StreamConsumer for DebugStreamCaptureConsumer { + fn consume(&self, bytes: &[u8]) -> Result<(), String> { + self.bytes.lock().extend_from_slice(bytes); + Ok(()) + } +} + #[cfg(all(debug_assertions, target_os = "linux"))] #[pyclass(name = "_TestDataPlaneHost")] struct PyTestDataPlaneHost { @@ -855,6 +913,7 @@ struct PyTestDataPlaneHost { _engine: Engine, handoff: data_plane::bootstrap::JobHandoff, namespace_root: std::path::PathBuf, + stream_data_plane: Arc, } #[cfg(all(debug_assertions, target_os = "linux"))] @@ -879,6 +938,57 @@ impl PyTestDataPlaneHost { fn arena_fd(&self) -> RawFd { std::os::fd::AsRawFd::as_raw_fd(&self.handoff.arena_fd) } + + fn read_stream<'py>(&self, py: Python<'py>, path: String) -> PyResult> { + let path = DataPath::parse(path) + .map_err(|error| PyErr::new::(error.to_string()))?; + let data_plane = self.stream_data_plane.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let reader = data_plane + .read_stream(&path) + .await + .map_err(data_plane_error)?; + Python::with_gil(|py| { + Py::new( + py, + PyStreamReader { + reader: Arc::new(tokio::sync::Mutex::new(reader)), + }, + ) + }) + }) + } + + fn capture_stream(&self, path: String) -> PyResult { + let path = DataPath::parse(path) + .map_err(|error| PyErr::new::(error.to_string()))?; + let bytes = Arc::new(ParkingMutex::new(Vec::new())); + let consumer: Arc = + Arc::new(DebugStreamCaptureConsumer { + bytes: Arc::clone(&bytes), + }); + let completion = self + .stream_data_plane + .collect_stream(path, consumer) + .map_err(data_plane_error)?; + Ok(PyTestStreamCapture { completion, bytes }) + } +} + +#[cfg(all(debug_assertions, target_os = "linux"))] +#[pyclass(name = "_TestStreamCapture")] +struct PyTestStreamCapture { + completion: ActorCompletion>, + bytes: Arc>>, +} + +#[cfg(all(debug_assertions, target_os = "linux"))] +#[pymethods] +impl PyTestStreamCapture { + fn result<'py>(&self, py: Python<'py>) -> PyResult> { + self.completion.wait().map_err(data_plane_error)?; + Ok(PyBytes::new(py, &self.bytes.lock())) + } } #[cfg(all(debug_assertions, target_os = "linux"))] @@ -976,6 +1086,8 @@ fn _test_data_plane_host() -> PyResult { )) .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; let namespace = data_plane::namespace::NamespaceClient::new(runtime.clone(), namespace_proxy); + let stream_transport: Arc = + Arc::new(data_plane::stream_transport::LocalStreamTransport::new()); let host_session = runtime .spawn( data_plane::host::HostDataPlaneSessionActor::new( @@ -998,11 +1110,12 @@ fn _test_data_plane_host() -> PyResult { .map_err(|error| PyRuntimeError::new_err(error.to_string()))?, ], }, - namespace: Some(namespace), + namespace: Some(namespace.clone()), transfer_receiver: Some(Arc::new(DebugBlobReceiver)), - source_sender: Some(source_sender), - source_publisher: Some(source_publisher), + source_sender: Some(Arc::clone(&source_sender)), + source_publisher: Some(Arc::clone(&source_publisher)), route_registrar: Some(registrar), + stream_transport: Some(Arc::clone(&stream_transport)), }, ) .map_err(|error| PyRuntimeError::new_err(error.to_string()))?, @@ -1010,6 +1123,59 @@ fn _test_data_plane_host() -> PyResult { .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; data_plane::host::install_session_env(&mut handoff, host_session, capability); + let mut sink_arena = data_plane::arena::ArenaManager::boot(data_plane::arena::ArenaConfig { + node_id: data_plane::arena::NodeId(2), + reservation_ceiling: 1 << 20, + base_alignment: 64, + }) + .map_err(|error| PyRuntimeError::new_err(format!("debug sink arena: {error:?}")))?; + let sink_handoff = data_plane::bootstrap::write_bootstrap( + &mut sink_arena, + data_plane::bootstrap::BootstrapSpec { + arena_generation: 2, + alignment: 64, + }, + ) + .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + let sink_session = runtime + .spawn( + data_plane::host::HostDataPlaneSessionActor::new( + data_plane::host::HostDataPlaneConfig { + runtime: runtime.clone(), + arena: sink_arena, + arena_generation: 2, + session_generation: 2, + capability, + job_context: data_plane::path::JobContext { + run_id: "test-run".to_owned(), + read_prefixes: vec![ + DataPath::parse("/runs/test-run/results") + .map_err(|error| PyRuntimeError::new_err(error.to_string()))?, + ], + write_prefixes: Vec::new(), + }, + namespace: Some(namespace), + transfer_receiver: None, + source_sender: None, + source_publisher: None, + route_registrar: None, + stream_transport: Some(stream_transport), + }, + ) + .map_err(|error| PyRuntimeError::new_err(error.to_string()))?, + ) + .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + let stream_data_plane = Arc::new( + future::block_on(DataPlaneBootstrap::attach( + sink_handoff.arena_fd, + runtime.clone(), + sink_session, + capability, + )) + .map_err(data_plane_error)? + .data_plane, + ); + driver.enable_actor_bridge(iroh_driver::ActorBridgeConfig { runtime: runtime.clone(), codec: codecs, @@ -1031,6 +1197,7 @@ fn _test_data_plane_host() -> PyResult { _engine: engine, handoff, namespace_root, + stream_data_plane, }) } @@ -1040,14 +1207,17 @@ pub fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("DataPathError", module.py().get_type::())?; module.add("BlobError", module.py().get_type::())?; module.add("SessionError", module.py().get_type::())?; + module.add("StreamError", module.py().get_type::())?; module.add_class::()?; module.add_class::()?; module.add_class::()?; module.add_class::()?; + module.add_class::()?; module.add_class::()?; module.add_function(wrap_pyfunction!(run, module)?)?; #[cfg(all(debug_assertions, target_os = "linux"))] { + module.add_class::()?; module.add_class::()?; module.add_function(wrap_pyfunction!(_test_data_plane_host, module)?)?; } diff --git a/crates/bindings/python/tests/test_bootstrap.py b/crates/bindings/python/tests/test_bootstrap.py index a12cdf8..15bf055 100644 --- a/crates/bindings/python/tests/test_bootstrap.py +++ b/crates/bindings/python/tests/test_bootstrap.py @@ -2,16 +2,15 @@ from __future__ import annotations +import asyncio import ctypes import json import os import struct import subprocess import importlib -import socket import sys import tempfile -import threading from pathlib import Path import pytest @@ -215,39 +214,29 @@ def test_missing_path_and_authorization_are_typed(monkeypatch, host): swactor.run(main) -def test_temporary_output_stream_bridge_remains_available(monkeypatch, host, tmp_path): +def test_native_stream_round_trip_has_ordered_eof(monkeypatch, host): install_host_env(monkeypatch, host) - socket_path = tmp_path / "output.sock" - listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - listener.bind(str(socket_path)) - listener.listen(1) - monkeypatch.setenv("SWACTOR_DATA_PLANE_OUTPUT", str(socket_path)) received = [] - def receive(): - connection, _ = listener.accept() - with connection: - chunks = [] - while chunk := connection.recv(4096): - chunks.append(chunk) - received.append(b"".join(chunks)) - - receiver = threading.Thread(target=receive) - receiver.start() - async def main(ctx): + async def receive(): + reader = await ctx.data.read_stream( + "/runs/self/results/predictions" + ) + while (chunk := await reader.read()) is not None: + received.append(chunk) + + receiver = asyncio.create_task(receive()) async with ctx.data.write_stream( "/runs/self/results/predictions" ) as stream: - await stream.write(b"temporary-result") + await stream.write(b"native-") + await stream.write(b"result") + await receiver - try: - swactor.run(main) - receiver.join(timeout=2) - assert not receiver.is_alive() - assert received == [b"temporary-result"] - finally: - listener.close() + swactor.run(main) + assert b"".join(received) == b"native-result" + assert "SWACTOR_DATA_PLANE_OUTPUT" not in host.env() def test_invalid_capability_prevents_main(monkeypatch, host): @@ -350,34 +339,18 @@ def test_real_exec_attachment_and_blob_mapping(host): not Path("/dev/nvidia0").exists(), reason="CUDA device is unavailable", ) -def test_real_exec_tinygrad_cuda_scenario(host, tmp_path): - socket_path = tmp_path / "cuda-output.sock" - listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - listener.bind(str(socket_path)) - listener.listen(1) - received = [] - - def receive(): - connection, _ = listener.accept() - with connection: - chunks = [] - while chunk := connection.recv(4096): - chunks.append(chunk) - received.append(b"".join(chunks)) - - receiver = threading.Thread(target=receive) - receiver.start() +def test_real_exec_tinygrad_cuda_scenario(host): + capture = host.capture_stream("/runs/self/results/inference") env = { key: value for key, value in os.environ.items() - if key not in (*BOOTSTRAP_ENV, *OLD_WAKE_ENV) + if key not in (*BOOTSTRAP_ENV, *OLD_WAKE_ENV, "SWACTOR_DATA_PLANE_OUTPUT") } env.update(dict(host.env())) env.update( { "CUDA_PTX": "1", "DEV": "CUDA", - "SWACTOR_DATA_PLANE_OUTPUT": str(socket_path), } ) script = ( @@ -387,21 +360,16 @@ def test_real_exec_tinygrad_cuda_scenario(host, tmp_path): / "jobs" / "tiny_linear_inference.py" ) - try: - result = subprocess.run( - [sys.executable, str(script)], - env=env, - pass_fds=(host.arena_fd(),), - text=True, - capture_output=True, - timeout=45, - check=False, - ) - receiver.join(timeout=2) - assert result.returncode == 0, result.stderr - assert not receiver.is_alive() - payload = json.loads(received[0]) - assert payload["device"].startswith("CUDA") - assert payload["output"] == pytest.approx([2.75, -8.75]) - finally: - listener.close() + result = subprocess.run( + [sys.executable, str(script)], + env=env, + pass_fds=(host.arena_fd(),), + text=True, + capture_output=True, + timeout=45, + check=False, + ) + assert result.returncode == 0, result.stderr + payload = json.loads(capture.result()) + assert payload["device"].startswith("CUDA") + assert payload["output"] == pytest.approx([2.75, -8.75]) diff --git a/crates/data-plane/Cargo.toml b/crates/data-plane/Cargo.toml index 12cfb6a..32509bf 100644 --- a/crates/data-plane/Cargo.toml +++ b/crates/data-plane/Cargo.toml @@ -14,6 +14,7 @@ swactor = { path = "../..", features = ["serde", "transport"] } swactor-transport = { path = "../transport" } swactor-engine = { path = "../engine" } +parking_lot = "0.12" [dev-dependencies] parking_lot = "0.12" diff --git a/crates/data-plane/src/byte_ring.rs b/crates/data-plane/src/byte_ring.rs index e3a465a..9840521 100644 --- a/crates/data-plane/src/byte_ring.rs +++ b/crates/data-plane/src/byte_ring.rs @@ -46,6 +46,7 @@ //! - P11 (binding slice) fast path performs no syscalls; //! - P12 one copy per side per byte. +use serde::{Deserialize, Serialize}; use std::ptr::NonNull; use std::sync::atomic::{AtomicU64, Ordering}; @@ -85,7 +86,7 @@ pub struct ByteRingSpec { } /// A located, installed ring: what `attach` needs to find it again. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct RingHandle { /// Lease start (the header lives here). pub offset: u64, @@ -93,9 +94,11 @@ pub struct RingHandle { pub capacity: u64, /// Ring generation. pub generation: u64, + /// Arena lease identity used for deterministic release. + pub lease_id: u64, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] pub enum Role { Producer, Consumer, @@ -112,15 +115,139 @@ pub struct Reservation { } /// Record framing layered on the byte stream (property P9). -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] pub enum RecordKind { Data, Eof, Fault, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct RecordMeta { + pub kind: RecordKind, + pub len: u64, +} + +/// Borrowed committed record. Dropping the view releases the complete framed +/// record back to the producer. +pub struct PinnedRecord<'a> { + endpoint: &'a mut Endpoint, + kind: RecordKind, + payload_start: u64, + payload_len: u64, + record_start: u64, + record_len: u64, + generation: u64, + released: bool, +} + +impl std::fmt::Debug for PinnedRecord<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PinnedRecord") + .field("kind", &self.kind) + .field("payload_len", &self.payload_len) + .field("record_start", &self.record_start) + .finish_non_exhaustive() + } +} + +impl PinnedRecord<'_> { + pub fn kind(&self) -> RecordKind { + self.kind + } + + pub fn len(&self) -> usize { + self.payload_len as usize + } + + pub fn is_empty(&self) -> bool { + self.payload_len == 0 + } + + pub fn spans(&self) -> (&[u8], &[u8]) { + self.endpoint + .read_spans(self.payload_start, self.payload_len) + } + + pub fn release(mut self) -> Result<(), FlowError> { + self.release_inner()?; + self.released = true; + Ok(()) + } + + fn release_inner(&mut self) -> Result<(), FlowError> { + self.endpoint.check("release_record", Role::Consumer)?; + let generation = self.endpoint.field_u64(OFF_GENERATION); + if generation != self.generation { + return Err(FlowError::StaleReservation { + reservation: self.generation, + ring: generation, + }); + } + let consume = self.endpoint.consume_cursor(); + if consume != self.record_start { + return Err(FlowError::PinnedRecordMoved { + expected: self.record_start, + found: consume, + }); + } + self.endpoint + .atomic(OFF_CONSUME) + .store(self.record_start + self.record_len, Ordering::Release); + Ok(()) + } +} + +impl Drop for PinnedRecord<'_> { + fn drop(&mut self) { + if !self.released && self.release_inner().is_ok() { + self.released = true; + } + } +} + +/// Producer reservation for a complete framed record. Payload bytes are +/// written directly into the ring and remain invisible until `commit`. +pub struct WritableRecord<'a> { + endpoint: &'a mut Endpoint, + reservation: Option, + payload_start: u64, + payload_len: u64, +} + +impl std::fmt::Debug for WritableRecord<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WritableRecord") + .field("payload_len", &self.payload_len) + .finish_non_exhaustive() + } +} + +impl WritableRecord<'_> { + pub fn len(&self) -> usize { + self.payload_len as usize + } + + pub fn is_empty(&self) -> bool { + self.payload_len == 0 + } + + pub fn spans_mut(&mut self) -> (&mut [u8], &mut [u8]) { + self.endpoint + .write_spans(self.payload_start, self.payload_len) + } + + pub fn commit(mut self) -> Result<(), FlowError> { + let reservation = self + .reservation + .take() + .expect("writable record commits at most once"); + self.endpoint.commit(reservation) + } +} + impl RecordKind { - fn to_byte(self) -> u8 { + pub fn to_byte(self) -> u8 { match self { Self::Data => 1, Self::Eof => 2, @@ -128,7 +255,7 @@ impl RecordKind { } } - fn from_byte(byte: u8) -> Option { + pub fn from_byte(byte: u8) -> Option { match byte { 1 => Some(Self::Data), 2 => Some(Self::Eof), @@ -194,6 +321,7 @@ pub enum AttachError { pub enum RecordError { LengthExceedsCapacity { len: u64, capacity: u64 }, InvalidKind(u8), + LengthExceedsWire { len: u64 }, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -218,6 +346,10 @@ pub enum FlowError { }, /// The header stopped validating mid-protocol (property P5). Corrupt(HeaderError), + PinnedRecordMoved { + expected: u64, + found: u64, + }, BadRecord(RecordError), Io, } @@ -230,6 +362,59 @@ pub struct Endpoint { role: Role, } +#[derive(Clone, Copy)] +pub struct RingProbe { + header: NonNull, + info: RingHandle, +} + +unsafe impl Send for RingProbe {} +unsafe impl Sync for RingProbe {} + +impl RingProbe { + fn atomic(&self, off: u64) -> &AtomicU64 { + // SAFETY: probes are created only from a successfully attached, + // aligned endpoint and remain bounded by that ring's live lease. + unsafe { &*(self.header.as_ptr().add(off as usize) as *const AtomicU64) } + } + + pub fn positions(&self) -> Result<(u64, u64), FlowError> { + let generation = self.atomic(OFF_GENERATION).load(Ordering::Relaxed); + if generation != self.info.generation { + return Err(FlowError::StaleReservation { + reservation: self.info.generation, + ring: generation, + }); + } + let commit = self.atomic(OFF_COMMIT).load(Ordering::Acquire); + let consume = self.atomic(OFF_CONSUME).load(Ordering::Acquire); + if commit < consume { + return Err(FlowError::Corrupt(HeaderError::CommitBelowConsume { + commit, + consume, + })); + } + if commit - consume > self.info.capacity { + return Err(FlowError::Corrupt(HeaderError::ReadableExceedsCapacity { + commit, + consume, + capacity: self.info.capacity, + })); + } + Ok((commit, consume)) + } + + pub fn has_data(&self) -> bool { + self.positions() + .is_ok_and(|(commit, consume)| commit > consume) + } + + pub fn has_capacity(&self) -> bool { + self.positions() + .is_ok_and(|(commit, consume)| commit - consume < self.info.capacity) + } +} + // SAFETY: an endpoint touches only its role-owned cursor field and the data // region under the single-writer protocol (properties P3/P4); the raw // pointer is never dereferenced outside `[header, header + DATA_OFFSET + @@ -295,6 +480,7 @@ pub fn install(arena: &mut ArenaManager, spec: ByteRingSpec) -> Result Result { + let total = DATA_OFFSET + .checked_add(handle.capacity) + .ok_or(AttachError::OutOfBounds { + end: u64::MAX, + arena_len: arena.len() as u64, + })?; + let range = + arena + .checked_range(handle.offset, total) + .map_err(|_| AttachError::OutOfBounds { + end: handle.offset.saturating_add(total), + arena_len: arena.len() as u64, + })?; + let endpoint = Endpoint { + header: arena.ptr_at(range.start), + info: handle, + role, + }; + endpoint.validate_fixed().map_err(AttachError::Header)?; + endpoint + .validate_generation(handle.generation) + .map_err(AttachError::Header)?; + endpoint.validate_cursors().map_err(AttachError::Header)?; + Ok(endpoint) +} + fn lease_ring( arena: &mut ArenaManager, request_id: u64, @@ -360,6 +578,32 @@ fn lease_ring( } impl Endpoint { + pub fn capacity(&self) -> u64 { + self.info.capacity + } + + pub fn role(&self) -> Role { + self.role + } + + /// Current published producer and consumer positions. This role-neutral + /// observation is used only to wait for an already-committed clean close + /// to enter the downstream bounded transport. + pub fn positions(&self) -> Result<(u64, u64), FlowError> { + self.validate_fixed().map_err(FlowError::Corrupt)?; + self.validate_generation(self.info.generation) + .map_err(FlowError::Corrupt)?; + self.validate_cursors().map_err(FlowError::Corrupt)?; + Ok((self.commit_cursor(), self.consume_cursor())) + } + + pub fn probe(&self) -> RingProbe { + RingProbe { + header: self.header, + info: self.info, + } + } + // ─── field access ──────────────────────────────────────────────────────── /// SAFETY: callers keep `self` alive; the pointer is bounds-checked at @@ -490,6 +734,36 @@ impl Endpoint { } } + fn read_spans(&self, stream_pos: u64, len: u64) -> (&[u8], &[u8]) { + let capacity = self.info.capacity as usize; + let start = (stream_pos % self.info.capacity) as usize; + let len = len as usize; + let first = len.min(capacity - start); + // SAFETY: `stream_pos` and `len` describe a validated committed + // record no larger than the ring. The split ranges do not overlap. + unsafe { + ( + std::slice::from_raw_parts(self.data_ptr().add(start), first), + std::slice::from_raw_parts(self.data_ptr(), len - first), + ) + } + } + + fn write_spans(&mut self, stream_pos: u64, len: u64) -> (&mut [u8], &mut [u8]) { + let capacity = self.info.capacity as usize; + let start = (stream_pos % self.info.capacity) as usize; + let len = len as usize; + let first = len.min(capacity - start); + // SAFETY: the producer exclusively owns this uncommitted reservation; + // it is no larger than the ring and the split ranges do not overlap. + unsafe { + ( + std::slice::from_raw_parts_mut(self.data_ptr().add(start), first), + std::slice::from_raw_parts_mut(self.data_ptr(), len - first), + ) + } + } + fn copy_out(&self, stream_pos: u64, len: u64) -> Vec { let capacity = self.info.capacity as usize; let start = (stream_pos % self.info.capacity) as usize; @@ -610,6 +884,34 @@ impl Endpoint { // ─── records (P9) ──────────────────────────────────────────────────────── + /// Reserve one complete record for direct producer-side payload writes. + pub fn reserve_record( + &mut self, + kind: RecordKind, + len: u64, + ) -> Result, FlowError> { + if len > u64::from(u32::MAX) { + return Err(FlowError::BadRecord(RecordError::LengthExceedsWire { len })); + } + if RECORD_HEADER_LEN + len > self.info.capacity { + return Err(FlowError::BadRecord(RecordError::LengthExceedsCapacity { + len, + capacity: self.info.capacity, + })); + } + let reservation = self.reserve(RECORD_HEADER_LEN + len)?; + let mut prefix = [0_u8; RECORD_HEADER_LEN as usize]; + prefix[0] = kind.to_byte(); + prefix[1..5].copy_from_slice(&(len as u32).to_le_bytes()); + self.copy_into(reservation.start, &prefix); + let payload_start = reservation.start + RECORD_HEADER_LEN; + Ok(WritableRecord { + endpoint: self, + reservation: Some(reservation), + payload_start, + payload_len: len, + }) + } /// Frame and commit one record in a single step. pub fn send_record(&mut self, kind: RecordKind, bytes: &[u8]) -> Result<(), FlowError> { let len = bytes.len() as u64; @@ -628,6 +930,72 @@ impl Endpoint { self.commit(reservation) } + /// Inspect the next complete committed record without pinning or consuming + /// it. Used by local direct transfer to reserve destination capacity before + /// borrowing the source payload. + pub fn next_record_meta(&self) -> Result, FlowError> { + self.check("next_record_meta", Role::Consumer)?; + self.validate_generation(self.info.generation) + .map_err(FlowError::Corrupt)?; + let (commit, consume) = (self.commit_cursor(), self.consume_cursor()); + let readable = commit - consume; + if readable < RECORD_HEADER_LEN { + return Ok(None); + } + let prefix = self.copy_out(consume, RECORD_HEADER_LEN); + let kind = RecordKind::from_byte(prefix[0]) + .ok_or(FlowError::BadRecord(RecordError::InvalidKind(prefix[0])))?; + let len = u64::from(u32::from_le_bytes(prefix[1..5].try_into().unwrap())); + if RECORD_HEADER_LEN + len > self.info.capacity { + return Err(FlowError::BadRecord(RecordError::LengthExceedsCapacity { + len, + capacity: self.info.capacity, + })); + } + if readable < RECORD_HEADER_LEN + len { + return Ok(None); + } + Ok(Some(RecordMeta { kind, len })) + } + /// Borrow one complete committed record without copying its payload. + /// The record remains pinned until the returned view is released or + /// dropped. + pub fn peek_record(&mut self) -> Result>, FlowError> { + self.check("peek_record", Role::Consumer)?; + self.validate_generation(self.info.generation) + .map_err(FlowError::Corrupt)?; + let (commit, consume) = (self.commit_cursor(), self.consume_cursor()); + let readable = commit - consume; + if readable < RECORD_HEADER_LEN { + return Ok(None); + } + let prefix = self.copy_out(consume, RECORD_HEADER_LEN); + let kind = RecordKind::from_byte(prefix[0]) + .ok_or(FlowError::BadRecord(RecordError::InvalidKind(prefix[0])))?; + let len = u64::from(u32::from_le_bytes(prefix[1..5].try_into().unwrap())); + if RECORD_HEADER_LEN + len > self.info.capacity { + return Err(FlowError::BadRecord(RecordError::LengthExceedsCapacity { + len, + capacity: self.info.capacity, + })); + } + let record_len = RECORD_HEADER_LEN + len; + if readable < record_len { + return Ok(None); + } + let generation = self.info.generation; + Ok(Some(PinnedRecord { + endpoint: self, + kind, + payload_start: consume + RECORD_HEADER_LEN, + payload_len: len, + record_start: consume, + record_len, + generation, + released: false, + })) + } + /// Receive one complete record; `Ok(None)` when nothing (or only a /// torn, uncommitted prefix) is readable. pub fn recv_record(&mut self) -> Result)>, FlowError> { diff --git a/crates/data-plane/src/data_plane.rs b/crates/data-plane/src/data_plane.rs index 833f771..28e8ee7 100644 --- a/crates/data-plane/src/data_plane.rs +++ b/crates/data-plane/src/data_plane.rs @@ -1,6 +1,6 @@ //! Child-side data-plane session, per-operation actors, and native API. -use std::collections::{HashMap, HashSet}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::os::fd::OwnedFd; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; @@ -8,14 +8,15 @@ use std::time::Duration; use swactor::actor::{ActorAddress, ActorInterface, Ctx}; use swactor::runtime::{ExternalSender, Runtime}; -use swactor_engine::EngineHandle; +use swactor_engine::{ActorCompletion, EngineHandle}; use crate::blob::{ Blob, BlobLease, BlobMetadata, LeaseReleaser, WritableArenaView, WritableBlobLease, }; +use crate::byte_ring::{Endpoint, FlowError, RecordKind, RingHandle, Role, attach_mapped}; use crate::mapped_arena::MappedArena; use crate::path::DataPath; -use crate::protocol::{ChildSessionIn, DataPlaneError, HostSessionIn, JobCapability}; +use crate::protocol::{ChildSessionIn, DataPlaneError, HostSessionIn, HostStreamIn, JobCapability}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ChildSessionState { @@ -140,6 +141,9 @@ impl DataPlaneBootstrap { operations: HashSet::new(), read_operations: HashMap::new(), state: ChildSessionState::Attaching, + stream_operations: HashMap::new(), + pending_blob_releases: 0, + deferred_blob_opens: VecDeque::new(), }) .map_err(|error| DataPlaneError::SessionFailed(error.to_string()))?; if let Some((engine, sender, timeout)) = deadline { @@ -178,6 +182,26 @@ impl Drop for ReadCancellation { } } +struct StreamOpenCancellation { + runtime: Runtime, + child_session: ActorAddress, + reply_to: ActorAddress, + armed: bool, +} + +impl Drop for StreamOpenCancellation { + fn drop(&mut self) { + if self.armed { + let _ = self.runtime.send_to( + self.child_session, + ChildSessionIn::CancelStream { + reply_to: self.reply_to, + }, + ); + } + } +} + #[derive(Clone)] pub struct DataPlane { runtime: Runtime, @@ -264,12 +288,94 @@ impl DataPlane { self.write_blob(&path, length).await } - pub async fn read_stream(&self, _path: &DataPath) -> Result<(), DataPlaneError> { - Err(DataPlaneError::StreamsDeferred) + async fn open_stream( + &self, + path: &DataPath, + role: Role, + replace: bool, + ) -> Result { + let inbox = self + .runtime + .new_inbox::() + .map_err(|error| DataPlaneError::SessionFailed(error.to_string()))?; + let reply_to = *inbox.addr(); + let message = match role { + Role::Consumer => ChildSessionIn::OpenReadStream { + path: path.clone(), + reply_to, + replace, + }, + Role::Producer => ChildSessionIn::OpenWriteStream { + path: path.clone(), + reply_to, + replace, + }, + }; + self.runtime + .send_to(self.child_session, message) + .map_err(|error| DataPlaneError::SessionFailed(error.to_string()))?; + let mut cancellation = StreamOpenCancellation { + runtime: self.runtime.clone(), + child_session: self.child_session, + reply_to, + armed: true, + }; + let result = match inbox.recv().await { + ChildStreamIn::Opened(result) => result, + ChildStreamIn::Wake(_) => Err(DataPlaneError::StreamFault( + "received stream wake before open completed".to_owned(), + )), + }; + cancellation.armed = false; + result } - pub async fn write_stream(&self, _path: &DataPath) -> Result<(), DataPlaneError> { - Err(DataPlaneError::StreamsDeferred) + pub async fn read_stream(&self, path: &DataPath) -> Result { + let grant = self.open_stream(path, Role::Consumer, false).await?; + let endpoint = attach_mapped(&self.arena, grant.ring, Role::Consumer).map_err(|error| { + DataPlaneError::StreamFault(format!("attach stream reader: {error:?}")) + })?; + Ok(StreamReader { + runtime: self.runtime.clone(), + child_session: self.child_session, + operation: grant.operation, + host_binding: grant.host_binding, + endpoint, + terminal: None, + }) + } + + pub async fn write_stream(&self, path: &DataPath) -> Result { + let grant = self.open_stream(path, Role::Producer, false).await?; + let endpoint = attach_mapped(&self.arena, grant.ring, Role::Producer).map_err(|error| { + DataPlaneError::StreamFault(format!("attach stream writer: {error:?}")) + })?; + Ok(StreamWriter { + runtime: self.runtime.clone(), + child_session: self.child_session, + operation: grant.operation, + host_binding: grant.host_binding, + endpoint, + closed: false, + }) + } + + pub async fn write_stream_replacing( + &self, + path: &DataPath, + ) -> Result { + let grant = self.open_stream(path, Role::Producer, true).await?; + let endpoint = attach_mapped(&self.arena, grant.ring, Role::Producer).map_err(|error| { + DataPlaneError::StreamFault(format!("attach stream writer: {error:?}")) + })?; + Ok(StreamWriter { + runtime: self.runtime.clone(), + child_session: self.child_session, + operation: grant.operation, + host_binding: grant.host_binding, + endpoint, + closed: false, + }) } pub fn close(&self) -> Result<(), DataPlaneError> { @@ -279,6 +385,34 @@ impl DataPlane { } } +pub trait StreamConsumer: Send + Sync + 'static { + fn consume(&self, bytes: &[u8]) -> Result<(), String>; +} + +impl DataPlane { + pub fn collect_stream( + &self, + path: DataPath, + consumer: Arc, + ) -> Result>, DataPlaneError> { + let completion = ActorCompletion::new(); + self.runtime + .spawn(StreamConsumerActor { + child_session: self.child_session, + arena: self.arena.clone(), + path, + consumer, + completion: completion.clone(), + operation: None, + host_binding: None, + endpoint: None, + pending_result: None, + finished: false, + }) + .map_err(|error| DataPlaneError::SessionFailed(error.to_string()))?; + Ok(completion) + } +} pub struct BlobWriter { runtime: Runtime, operation: ActorAddress, @@ -350,6 +484,608 @@ impl Drop for BlobWriter { } } +#[derive(Clone)] +pub(crate) struct StreamOpenGrant { + operation: ActorAddress, + host_binding: ActorAddress, + ring: RingHandle, +} + +#[derive(Clone)] +pub(crate) enum ChildStreamIn { + Opened(Result), + Wake(Result<(), DataPlaneError>), +} + +pub struct StreamWriter { + runtime: Runtime, + child_session: ActorAddress, + operation: ActorAddress, + host_binding: ActorAddress, + endpoint: Endpoint, + closed: bool, +} + +impl StreamWriter { + pub fn capacity(&self) -> u64 { + self.endpoint.capacity() + } + fn send_control(&self, message: HostStreamIn) -> Result<(), DataPlaneError> { + self.runtime + .send_to( + self.child_session, + ChildSessionIn::StreamControl { + binding: self.host_binding, + message, + }, + ) + .map_err(|error| DataPlaneError::SessionFailed(error.to_string())) + } + + async fn send_one(&mut self, kind: RecordKind, bytes: &[u8]) -> Result<(), DataPlaneError> { + loop { + match self.endpoint.send_record(kind, bytes) { + Ok(()) => { + self.send_control(HostStreamIn::DataAvailable)?; + return Ok(()); + } + Err(FlowError::InsufficientSpace { .. }) => { + let inbox = self + .runtime + .new_inbox::() + .map_err(|error| DataPlaneError::SessionFailed(error.to_string()))?; + self.send_control(HostStreamIn::WaitCapacity { + reply_to: *inbox.addr(), + })?; + match self.endpoint.send_record(kind, bytes) { + Ok(()) => { + self.send_control(HostStreamIn::DataAvailable)?; + return Ok(()); + } + Err(FlowError::InsufficientSpace { .. }) => match inbox.recv().await { + ChildStreamIn::Wake(result) => result?, + ChildStreamIn::Opened(_) => { + return Err(DataPlaneError::StreamFault( + "received stream-open result while waiting for capacity" + .to_owned(), + )); + } + }, + Err(error) => { + return Err(DataPlaneError::StreamFault(format!( + "write stream ring: {error:?}" + ))); + } + } + } + Err(error) => { + return Err(DataPlaneError::StreamFault(format!( + "write stream ring: {error:?}" + ))); + } + } + } + } + + pub async fn flush(&mut self) -> Result<(), DataPlaneError> { + let target = self + .endpoint + .positions() + .map_err(|error| { + DataPlaneError::StreamFault(format!("observe stream flush position: {error:?}")) + })? + .0; + loop { + let consumed = self + .endpoint + .positions() + .map_err(|error| { + DataPlaneError::StreamFault(format!("observe stream flush progress: {error:?}")) + })? + .1; + if consumed >= target { + return Ok(()); + } + let inbox = self + .runtime + .new_inbox::() + .map_err(|error| DataPlaneError::SessionFailed(error.to_string()))?; + self.send_control(HostStreamIn::WaitCapacity { + reply_to: *inbox.addr(), + })?; + self.send_control(HostStreamIn::DataAvailable)?; + if self + .endpoint + .positions() + .map_err(|error| { + DataPlaneError::StreamFault(format!("observe stream flush progress: {error:?}")) + })? + .1 + >= target + { + return Ok(()); + } + match inbox.recv().await { + ChildStreamIn::Wake(result) => { + if let Err(error) = result + && self + .endpoint + .positions() + .map_err(|flow| { + DataPlaneError::StreamFault(format!( + "observe terminal flush progress: {flow:?}" + )) + })? + .1 + < target + { + return Err(error); + } + } + ChildStreamIn::Opened(_) => { + return Err(DataPlaneError::StreamFault( + "received stream-open result while flushing".to_owned(), + )); + } + } + } + } + + pub async fn write(&mut self, bytes: &[u8]) -> Result<(), DataPlaneError> { + if self.closed { + return Err(DataPlaneError::StreamClosed); + } + if bytes.is_empty() { + return Ok(()); + } + let max_payload = usize::try_from(self.endpoint.capacity().saturating_sub(5)) + .map_err(|_| DataPlaneError::StreamFault("stream capacity exceeds usize".to_owned()))?; + if max_payload == 0 { + return Err(DataPlaneError::StreamFault( + "stream ring cannot hold a framed byte".to_owned(), + )); + } + for chunk in bytes.chunks(max_payload) { + self.send_one(RecordKind::Data, chunk).await?; + } + Ok(()) + } + + pub async fn close(&mut self) -> Result<(), DataPlaneError> { + if self.closed { + return Ok(()); + } + self.send_one(RecordKind::Eof, &[]).await?; + self.flush().await?; + self.closed = true; + let _ = self.runtime.send_to( + self.child_session, + ChildSessionIn::OperationDone { + operation: self.operation, + }, + ); + Ok(()) + } + + pub fn abort(&mut self) -> Result<(), DataPlaneError> { + if self.closed { + return Ok(()); + } + self.closed = true; + self.send_control(HostStreamIn::Close { + clean: false, + reply_to: None, + }) + } +} + +impl Drop for StreamWriter { + fn drop(&mut self) { + if !self.closed { + let _ = self.send_control(HostStreamIn::Close { + clean: false, + reply_to: None, + }); + } + let _ = self.runtime.send_to( + self.child_session, + ChildSessionIn::OperationDone { + operation: self.operation, + }, + ); + } +} + +#[derive(Clone)] +enum StreamReadTerminal { + Eof, + Error(DataPlaneError), +} + +pub struct StreamReader { + runtime: Runtime, + child_session: ActorAddress, + operation: ActorAddress, + host_binding: ActorAddress, + endpoint: Endpoint, + terminal: Option, +} + +impl StreamReader { + pub fn capacity(&self) -> u64 { + self.endpoint.capacity() + } + + fn send_control(&self, message: HostStreamIn) -> Result<(), DataPlaneError> { + self.runtime + .send_to( + self.child_session, + ChildSessionIn::StreamControl { + binding: self.host_binding, + message, + }, + ) + .map_err(|error| DataPlaneError::SessionFailed(error.to_string())) + } + async fn close_clean(&mut self) -> Result<(), DataPlaneError> { + let inbox = self + .runtime + .new_inbox::() + .map_err(|error| DataPlaneError::SessionFailed(error.to_string()))?; + self.send_control(HostStreamIn::Close { + clean: true, + reply_to: Some(*inbox.addr()), + })?; + match inbox.recv().await { + ChildStreamIn::Wake(result) => result, + ChildStreamIn::Opened(_) => Err(DataPlaneError::StreamFault( + "received stream-open result while closing reader".to_owned(), + )), + } + } + + async fn finish_read( + &mut self, + result: Option>, + ) -> Result>, DataPlaneError> { + if result.is_none() { + self.close_clean().await?; + let _ = self.runtime.send_to( + self.child_session, + ChildSessionIn::OperationDone { + operation: self.operation, + }, + ); + } + Ok(result) + } + + fn terminal_result(&self) -> Option>, DataPlaneError>> { + self.terminal.as_ref().map(|terminal| match terminal { + StreamReadTerminal::Eof => Ok(None), + StreamReadTerminal::Error(error) => Err(error.clone()), + }) + } + + fn try_read(&mut self) -> Result>>, DataPlaneError> { + let Some(view) = self + .endpoint + .peek_record() + .map_err(|error| DataPlaneError::StreamFault(format!("read stream ring: {error:?}")))? + else { + return Ok(None); + }; + let kind = view.kind(); + let (first, second) = view.spans(); + let mut bytes = Vec::with_capacity(first.len() + second.len()); + bytes.extend_from_slice(first); + bytes.extend_from_slice(second); + view.release().map_err(|error| { + DataPlaneError::StreamFault(format!("consume stream ring: {error:?}")) + })?; + self.send_control(HostStreamIn::CapacityAvailable)?; + match kind { + RecordKind::Data => Ok(Some(Some(bytes))), + RecordKind::Eof => { + self.terminal = Some(StreamReadTerminal::Eof); + Ok(Some(None)) + } + RecordKind::Fault => { + let error = + DataPlaneError::StreamFault(String::from_utf8_lossy(&bytes).into_owned()); + self.terminal = Some(StreamReadTerminal::Error(error.clone())); + let _ = self.send_control(HostStreamIn::Close { + clean: false, + reply_to: None, + }); + Err(error) + } + } + } + + pub async fn read(&mut self) -> Result>, DataPlaneError> { + if let Some(result) = self.terminal_result() { + return result; + } + loop { + if let Some(result) = self.try_read()? { + return self.finish_read(result).await; + } + let inbox = self + .runtime + .new_inbox::() + .map_err(|error| DataPlaneError::SessionFailed(error.to_string()))?; + self.send_control(HostStreamIn::WaitData { + reply_to: *inbox.addr(), + })?; + if let Some(result) = self.try_read()? { + return self.finish_read(result).await; + } + match inbox.recv().await { + ChildStreamIn::Wake(Ok(())) => {} + ChildStreamIn::Wake(Err(error)) => { + self.terminal = Some(StreamReadTerminal::Error(error.clone())); + return Err(error); + } + ChildStreamIn::Opened(_) => { + return Err(DataPlaneError::StreamFault( + "received stream-open result while waiting for data".to_owned(), + )); + } + } + } + } +} + +impl Drop for StreamReader { + fn drop(&mut self) { + if self.terminal.is_none() { + let _ = self.send_control(HostStreamIn::Close { + clean: false, + reply_to: None, + }); + } + let _ = self.runtime.send_to( + self.child_session, + ChildSessionIn::OperationDone { + operation: self.operation, + }, + ); + } +} + +struct StreamConsumerActor { + child_session: ActorAddress, + arena: Arc, + path: DataPath, + consumer: Arc, + completion: ActorCompletion>, + operation: Option, + host_binding: Option, + endpoint: Option, + pending_result: Option>, + finished: bool, +} + +enum ConsumerDrainStep { + Empty, + Data, + Eof, + Fault(String), +} + +fn consume_next_record( + endpoint: &mut Endpoint, + consumer: &dyn StreamConsumer, +) -> Result { + let Some(view) = endpoint + .peek_record() + .map_err(|error| DataPlaneError::StreamFault(format!("collect stream ring: {error:?}")))? + else { + return Ok(ConsumerDrainStep::Empty); + }; + let kind = view.kind(); + let (first, second) = view.spans(); + let fault = (kind == RecordKind::Fault).then(|| { + let mut reason = Vec::with_capacity(first.len() + second.len()); + reason.extend_from_slice(first); + reason.extend_from_slice(second); + String::from_utf8_lossy(&reason).into_owned() + }); + if kind == RecordKind::Data { + consumer + .consume(first) + .and_then(|()| consumer.consume(second)) + .map_err(DataPlaneError::StreamFault)?; + } + view.release().map_err(|error| { + DataPlaneError::StreamFault(format!("release collected stream ring: {error:?}")) + })?; + Ok(match kind { + RecordKind::Data => ConsumerDrainStep::Data, + RecordKind::Eof => ConsumerDrainStep::Eof, + RecordKind::Fault => ConsumerDrainStep::Fault(fault.expect("fault payload captured")), + }) +} + +impl StreamConsumerActor { + fn send_control(&self, ctx: &Ctx<'_>, binding: ActorAddress, message: HostStreamIn) { + let _ = ctx.send( + self.child_session, + ChildSessionIn::StreamControl { binding, message }, + ); + } + + fn complete_now(&mut self, ctx: &Ctx<'_>, result: Result<(), DataPlaneError>) { + if self.finished { + return; + } + self.finished = true; + if let Some(operation) = self.operation { + let _ = ctx.send( + self.child_session, + ChildSessionIn::OperationDone { operation }, + ); + } + let _ = self.completion.complete(result); + ctx.stop_self(); + } + + fn finish(&mut self, ctx: &Ctx<'_>, result: Result<(), DataPlaneError>, clean: bool) { + if self.finished || self.pending_result.is_some() { + return; + } + if clean && let Some(host_binding) = self.host_binding { + self.pending_result = Some(result); + self.send_control( + ctx, + host_binding, + HostStreamIn::Close { + clean: true, + reply_to: Some(ctx.self_addr()), + }, + ); + return; + } + if let Some(host_binding) = self.host_binding { + self.send_control( + ctx, + host_binding, + HostStreamIn::Close { + clean: false, + reply_to: None, + }, + ); + } + self.complete_now(ctx, result); + } + + fn drain(&mut self, ctx: &Ctx<'_>) { + loop { + let step = consume_next_record( + self.endpoint + .as_mut() + .expect("collector endpoint is installed before drain"), + self.consumer.as_ref(), + ); + match step { + Ok(ConsumerDrainStep::Empty) => { + if let Some(host_binding) = self.host_binding { + self.send_control( + ctx, + host_binding, + HostStreamIn::WaitData { + reply_to: ctx.self_addr(), + }, + ); + } + return; + } + Ok(ConsumerDrainStep::Data) => { + if let Some(host_binding) = self.host_binding { + self.send_control(ctx, host_binding, HostStreamIn::CapacityAvailable); + } + } + Ok(ConsumerDrainStep::Eof) => { + if let Some(host_binding) = self.host_binding { + self.send_control(ctx, host_binding, HostStreamIn::CapacityAvailable); + } + self.finish(ctx, Ok(()), true); + return; + } + Ok(ConsumerDrainStep::Fault(reason)) => { + self.finish(ctx, Err(DataPlaneError::StreamFault(reason)), false); + return; + } + Err(error) => { + self.finish(ctx, Err(error), false); + return; + } + } + } + } +} + +impl ActorInterface for StreamConsumerActor { + type Incoming = ChildStreamIn; + type Response = (); + + fn on_start(&mut self, ctx: &Ctx<'_>) { + let _ = ctx.send( + self.child_session, + ChildSessionIn::OpenReadStream { + path: self.path.clone(), + reply_to: ctx.self_addr(), + replace: false, + }, + ); + } + + fn handle(&mut self, ctx: &Ctx<'_>, message: ChildStreamIn) { + match message { + ChildStreamIn::Opened(Ok(grant)) => { + match attach_mapped(&self.arena, grant.ring, Role::Consumer) { + Ok(endpoint) => { + self.operation = Some(grant.operation); + self.host_binding = Some(grant.host_binding); + self.endpoint = Some(endpoint); + self.drain(ctx); + } + Err(error) => self.finish( + ctx, + Err(DataPlaneError::StreamFault(format!( + "attach stream collector: {error:?}" + ))), + false, + ), + } + } + ChildStreamIn::Opened(Err(error)) => { + self.finish(ctx, Err(error), false); + } + ChildStreamIn::Wake(result) => { + if let Some(pending) = self.pending_result.take() { + let completed = match result { + Ok(()) => pending, + Err(error) => Err(error), + }; + self.complete_now(ctx, completed); + } else { + match result { + Ok(()) => self.drain(ctx), + Err(error) => self.finish(ctx, Err(error), false), + } + } + } + } + } + + fn on_stop(&mut self, ctx: &Ctx<'_>) { + if !self.finished { + if let Some(host_binding) = self.host_binding { + self.send_control( + ctx, + host_binding, + HostStreamIn::Close { + clean: false, + reply_to: None, + }, + ); + } else { + let _ = ctx.send( + self.child_session, + ChildSessionIn::CancelStream { + reply_to: ctx.self_addr(), + }, + ); + } + let _ = self + .completion + .complete(Err(DataPlaneError::OperationCancelled)); + } + } +} + pub struct ChildDataPlaneSessionActor { runtime: Runtime, host_session: ActorAddress, @@ -362,6 +1098,9 @@ pub struct ChildDataPlaneSessionActor { operations: HashSet, read_operations: HashMap, state: ChildSessionState, + stream_operations: HashMap, + pending_blob_releases: usize, + deferred_blob_opens: VecDeque, } impl ChildDataPlaneSessionActor { @@ -372,6 +1111,44 @@ impl ChildDataPlaneSessionActor { fn fail_local_open(&self, ctx: &Ctx<'_>, reply_to: ActorAddress, error: DataPlaneError) { let _ = ctx.send(reply_to, Err::(error)); } + + fn start_stream_open( + &mut self, + ctx: &Ctx<'_>, + path: DataPath, + reply_to: ActorAddress, + role: Role, + replace: bool, + ) { + if self.state != ChildSessionState::Running { + let _ = ctx.send( + reply_to, + Err::(DataPlaneError::SessionNotRunning), + ); + return; + } + let actor = StreamOpenOperationActor { + host_session: self.host_session, + child_session: ctx.self_addr(), + path, + role, + replace, + reply_to, + replied: false, + }; + match ctx.spawn(actor) { + Ok(operation) => { + self.operations.insert(operation); + self.stream_operations.insert(reply_to, operation); + } + Err(error) => { + let _ = ctx.send( + reply_to, + Err::(DataPlaneError::SessionFailed(error.to_string())), + ); + } + } + } } impl ActorInterface for ChildDataPlaneSessionActor { @@ -435,6 +1212,11 @@ impl ActorInterface for ChildDataPlaneSessionActor { ctx.stop_self(); } ChildSessionIn::ReadBlob { path, reply_to } => { + if self.pending_blob_releases != 0 { + self.deferred_blob_opens + .push_back(ChildSessionIn::ReadBlob { path, reply_to }); + return; + } if self.state != ChildSessionState::Running { self.fail_local_open(ctx, reply_to, DataPlaneError::SessionNotRunning); return; @@ -471,6 +1253,15 @@ impl ActorInterface for ChildDataPlaneSessionActor { length, reply_to, } => { + if self.pending_blob_releases != 0 { + self.deferred_blob_opens + .push_back(ChildSessionIn::OpenWriteBlob { + path, + length, + reply_to, + }); + return; + } if self.state != ChildSessionState::Running { let _ = ctx.send( reply_to, @@ -503,6 +1294,62 @@ impl ActorInterface for ChildDataPlaneSessionActor { } } } + ChildSessionIn::OpenReadStream { + path, + reply_to, + replace, + } => { + self.start_stream_open(ctx, path, reply_to, Role::Consumer, replace); + } + ChildSessionIn::OpenWriteStream { + path, + reply_to, + replace, + } => { + self.start_stream_open(ctx, path, reply_to, Role::Producer, replace); + } + ChildSessionIn::CancelStream { reply_to } => { + if let Some(operation) = self.stream_operations.remove(&reply_to) { + self.operations.remove(&operation); + let _ = ctx.send(self.host_session, HostSessionIn::CancelStream { operation }); + let _ = ctx.stop_actor(operation); + } + } + ChildSessionIn::StreamWake { reply_to, result } => { + let _ = ctx.send(reply_to, ChildStreamIn::Wake(result)); + } + ChildSessionIn::StreamControl { binding, message } => { + let _ = ctx.send( + self.host_session, + HostSessionIn::StreamControl { binding, message }, + ); + } + ChildSessionIn::BlobReleased => { + self.pending_blob_releases = self.pending_blob_releases.saturating_sub(1); + if self.pending_blob_releases == 0 { + while let Some(deferred) = self.deferred_blob_opens.pop_front() { + self.handle(ctx, deferred); + if self.pending_blob_releases != 0 { + break; + } + } + } + } + ChildSessionIn::ReleaseBlob { + binding, + lease_id, + generation, + } => { + self.pending_blob_releases = self.pending_blob_releases.saturating_add(1); + let _ = ctx.send( + self.host_session, + HostSessionIn::ReleaseBlob { + binding, + lease_id, + generation, + }, + ); + } ChildSessionIn::BlobOpened { operation, host_binding, @@ -537,6 +1384,23 @@ impl ActorInterface for ChildDataPlaneSessionActor { ); } } + ChildSessionIn::StreamOpened { + operation, + host_binding, + ring, + role, + } => { + if self.operations.contains(&operation) { + let _ = ctx.send( + operation, + ChildOperationIn::StreamOpened { + host_binding, + ring, + role, + }, + ); + } + } ChildSessionIn::OperationFailed { operation, error } => { if self.operations.contains(&operation) { let _ = ctx.send(operation, ChildOperationIn::Failed(error)); @@ -556,6 +1420,8 @@ impl ActorInterface for ChildDataPlaneSessionActor { self.operations.remove(&operation); self.read_operations .retain(|_, read_operation| *read_operation != operation); + self.stream_operations + .retain(|_, stream_operation| *stream_operation != operation); } ChildSessionIn::Close => { if matches!( @@ -569,6 +1435,7 @@ impl ActorInterface for ChildDataPlaneSessionActor { let _ = ctx.stop_actor(operation); } self.read_operations.clear(); + self.stream_operations.clear(); let _ = ctx.send(self.host_session, HostSessionIn::Close); self.state = ChildSessionState::Closed; } @@ -589,6 +1456,11 @@ enum ChildOperationIn { lease: BlobLease, metadata: BlobMetadata, }, + StreamOpened { + host_binding: ActorAddress, + ring: RingHandle, + role: Role, + }, Failed(DataPlaneError), SealRequested { reply_to: ActorAddress, @@ -603,17 +1475,106 @@ enum ChildOperationIn { WriteAborted, } +struct StreamOpenOperationActor { + host_session: ActorAddress, + child_session: ActorAddress, + path: DataPath, + role: Role, + reply_to: ActorAddress, + replace: bool, + replied: bool, +} + +impl StreamOpenOperationActor { + fn finish(&mut self, ctx: &Ctx<'_>, result: Result) { + self.replied = true; + let _ = ctx.send(self.reply_to, ChildStreamIn::Opened(result)); + let _ = ctx.send( + self.child_session, + ChildSessionIn::OperationDone { + operation: ctx.self_addr(), + }, + ); + ctx.stop_self(); + } +} + +impl ActorInterface for StreamOpenOperationActor { + type Incoming = ChildOperationIn; + type Response = (); + + fn on_start(&mut self, ctx: &Ctx<'_>) { + let message = match self.role { + Role::Consumer => HostSessionIn::OpenReadStream { + path: self.path.clone(), + child_session: self.child_session, + operation: ctx.self_addr(), + replace: self.replace, + }, + Role::Producer => HostSessionIn::OpenWriteStream { + path: self.path.clone(), + child_session: self.child_session, + operation: ctx.self_addr(), + replace: self.replace, + }, + }; + if let Err(error) = ctx.send(self.host_session, message) { + self.finish(ctx, Err(DataPlaneError::SessionFailed(error.to_string()))); + } + } + + fn handle(&mut self, ctx: &Ctx<'_>, message: ChildOperationIn) { + match message { + ChildOperationIn::StreamOpened { + host_binding, + ring, + role, + } if role == self.role => self.finish( + ctx, + Ok(StreamOpenGrant { + operation: ctx.self_addr(), + host_binding, + ring, + }), + ), + ChildOperationIn::StreamOpened { .. } => self.finish( + ctx, + Err(DataPlaneError::StreamFault( + "host opened stream with the wrong ring role".to_owned(), + )), + ), + ChildOperationIn::Failed(error) => self.finish(ctx, Err(error)), + _ => {} + } + } + + fn on_stop(&mut self, ctx: &Ctx<'_>) { + if !self.replied { + let _ = ctx.send( + self.reply_to, + ChildStreamIn::Opened(Err(DataPlaneError::OperationCancelled)), + ); + let _ = ctx.send( + self.host_session, + HostSessionIn::CancelStream { + operation: ctx.self_addr(), + }, + ); + } + } +} + struct RuntimeLeaseReleaser { runtime: Runtime, - host_session: ActorAddress, + child_session: ActorAddress, host_binding: ActorAddress, } impl LeaseReleaser for RuntimeLeaseReleaser { fn release(&self, lease: BlobLease) { let _ = self.runtime.send_to( - self.host_session, - HostSessionIn::ReleaseBlob { + self.child_session, + ChildSessionIn::ReleaseBlob { binding: self.host_binding, lease_id: lease.lease_id, generation: lease.generation, @@ -680,7 +1641,7 @@ impl ActorInterface for ReadBlobOperationActor { } => { let releaser: Arc = Arc::new(RuntimeLeaseReleaser { runtime: self.runtime.clone(), - host_session: self.host_session, + child_session: self.child_session, host_binding, }); let result = Blob::from_sealed_lease(self.arena.clone(), lease, metadata, releaser) diff --git a/crates/data-plane/src/host.rs b/crates/data-plane/src/host.rs index a5a8137..eaab5cb 100644 --- a/crates/data-plane/src/host.rs +++ b/crates/data-plane/src/host.rs @@ -1,6 +1,6 @@ //! Host-side session, binding, and arena-allocation actors. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::os::fd::{FromRawFd, OwnedFd}; use std::sync::Arc; @@ -20,18 +20,27 @@ use crate::blob_transfer::{ BlobTransferEvent, BlobTransferId, BlobTransferOffer, BlobTransferReceiver, BlobTransferSender, }; use crate::bootstrap::JobHandoff; +use crate::byte_ring::{self, ByteRingSpec, RingHandle, Role}; +use crate::mapped_arena::MappedArena; use crate::namespace::{ BlobBinding as NamespaceBlobBinding, DataDirectoryOut, NamespaceClient, NamespaceClientIn, - NamespaceError, NamespaceRequest, OperationId, SourceRecovery, + NamespaceError, NamespaceRequest, OperationId, SourceRecovery, StreamIncarnation, StreamMatch, + StreamRole, }; use crate::path::{DataPath, JobContext}; use crate::protocol::{ - AttachmentFailure, ChildSessionIn, DataOperation, DataPlaneError, HostSessionIn, JobCapability, + AttachmentFailure, ChildSessionIn, DataOperation, DataPlaneError, HostSessionIn, HostStreamIn, + JobCapability, }; use crate::source::{BlobSourceIn, BlobSourcePublisher, BlobSourceRetirement, FileBlobSourceActor}; +use crate::stream_transport::{ + StreamPeerDescriptor, StreamSinkRequest, StreamSourceRequest, StreamTransport, + StreamTransportEvent, StreamTransportNotifier, +}; const BLOB_ALIGNMENT: u64 = 64; const FIRST_BLOB_REQUEST_ID: u64 = 2; +const STREAM_RING_CAPACITY: u64 = 256 * 1024; pub trait HostRouteRegistrar: Send + Sync + 'static { fn register_child( @@ -53,6 +62,7 @@ pub struct HostDataPlaneConfig { pub source_sender: Option>, pub source_publisher: Option>, pub route_registrar: Option>, + pub stream_transport: Option>, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -75,9 +85,12 @@ pub struct HostDataPlaneSessionActor { source_sender: Option>, source_publisher: Option>, route_registrar: Option>, + stream_arena: Option>, + stream_transport: Option>, child_session: Option, allocator: Option, active_bindings: HashSet, + stream_bindings: HashMap, state: HostSessionState, } @@ -92,6 +105,21 @@ impl HostDataPlaneSessionActor { .job_context .validate() .map_err(|error| DataPlaneError::InvalidPath(error.to_string()))?; + let stream_arena = if config.stream_transport.is_some() { + let fd = unsafe { libc::dup(config.arena.arena_fd()) }; + if fd < 0 { + return Err(DataPlaneError::SessionFailed(format!( + "duplicate arena backing for streams: {}", + std::io::Error::last_os_error() + ))); + } + let owned = unsafe { OwnedFd::from_raw_fd(fd) }; + let (mapped, _) = MappedArena::map(owned) + .map_err(|error| DataPlaneError::SessionFailed(error.to_string()))?; + Some(Arc::new(mapped)) + } else { + None + }; Ok(Self { arena: Some(config.arena), arena_generation: config.arena_generation, @@ -104,9 +132,12 @@ impl HostDataPlaneSessionActor { source_sender: config.source_sender, source_publisher: config.source_publisher, route_registrar: config.route_registrar, + stream_arena, + stream_transport: config.stream_transport, child_session: None, allocator: None, active_bindings: HashSet::new(), + stream_bindings: HashMap::new(), state: HostSessionState::AwaitingAttachment, }) } @@ -159,7 +190,10 @@ impl HostDataPlaneSessionActor { } fn maybe_finish_close(&mut self) { - if self.state == HostSessionState::Closing && self.active_bindings.is_empty() { + if self.state == HostSessionState::Closing + && self.active_bindings.is_empty() + && self.stream_bindings.is_empty() + { self.state = HostSessionState::Closed; } } @@ -329,6 +363,133 @@ impl ActorInterface for HostDataPlaneSessionActor { ), } } + ref message @ (HostSessionIn::OpenReadStream { + ref path, + child_session, + operation, + replace, + } + | HostSessionIn::OpenWriteStream { + ref path, + child_session, + operation, + replace, + }) => { + let role = match message { + HostSessionIn::OpenReadStream { .. } => StreamRole::Sink, + HostSessionIn::OpenWriteStream { .. } => StreamRole::Source, + _ => unreachable!(), + }; + let data_operation = match role { + StreamRole::Source => DataOperation::WriteStream, + StreamRole::Sink => DataOperation::ReadStream, + }; + let resolved = match self.validate_open(child_session, path, data_operation) { + Ok(path) => path, + Err(error) => { + self.send_open_failure(ctx, child_session, operation, error); + return; + } + }; + let (Some(namespace), Some(arena), Some(transport)) = ( + self.namespace.clone(), + self.stream_arena.clone(), + self.stream_transport.clone(), + ) else { + self.send_open_failure( + ctx, + child_session, + operation, + DataPlaneError::SessionFailed( + "stream namespace or transport service is unavailable".to_owned(), + ), + ); + return; + }; + let local_descriptor = match transport.descriptor() { + Ok(descriptor) => descriptor, + Err(reason) => { + self.send_open_failure( + ctx, + child_session, + operation, + DataPlaneError::StreamFault(reason), + ); + return; + } + }; + let binding = HostStreamBindingActor { + runtime: self.runtime.clone(), + host_session: ctx.self_addr(), + allocator: self.allocator.expect("allocator started"), + child_session, + operation, + path: resolved, + replace, + role, + namespace, + arena, + transport, + local_descriptor, + ring: None, + matched: None, + peer_descriptor: None, + transport_installed: false, + transport_ready: false, + transport_quiesced: false, + opened: false, + terminal: None, + data_waiters: Vec::new(), + capacity_waiters: Vec::new(), + close_waiters: Vec::new(), + release_started: false, + }; + match ctx.spawn(binding) { + Ok(binding) => { + if let Some(publisher) = &self.source_publisher + && let Err(error) = publisher.publish_source(binding) + { + let _ = ctx.stop_actor(binding); + self.send_open_failure( + ctx, + child_session, + operation, + DataPlaneError::SessionFailed(format!( + "publish stream endpoint: {error}" + )), + ); + return; + } + self.stream_bindings.insert(operation, binding); + } + Err(error) => self.send_open_failure( + ctx, + child_session, + operation, + DataPlaneError::SessionFailed(error.to_string()), + ), + } + } + HostSessionIn::CancelStream { operation } => { + if let Some(binding) = self.stream_bindings.get(&operation).copied() { + let _ = ctx.send( + binding, + HostStreamIn::Close { + clean: false, + reply_to: None, + }, + ); + } + } + HostSessionIn::StreamControl { binding, message } => { + if self + .stream_bindings + .values() + .any(|stream_binding| *stream_binding == binding) + { + let _ = ctx.send(binding, message); + } + } HostSessionIn::ReleaseBlob { binding, lease_id, @@ -391,6 +552,8 @@ impl ActorInterface for HostDataPlaneSessionActor { } HostSessionIn::BindingDone { binding } | HostSessionIn::BindingDetached { binding } => { self.active_bindings.remove(&binding); + self.stream_bindings + .retain(|_, stream_binding| *stream_binding != binding); self.maybe_finish_close(); } HostSessionIn::ConfigureRun { run_id, reply_to } => { @@ -422,6 +585,15 @@ impl ActorInterface for HostDataPlaneSessionActor { for binding in self.active_bindings.iter().copied() { let _ = ctx.send(binding, HostBindingIn::SessionClosed); } + for binding in self.stream_bindings.values().copied() { + let _ = ctx.send( + binding, + HostStreamIn::Close { + clean: false, + reply_to: None, + }, + ); + } self.maybe_finish_close(); } } @@ -465,6 +637,14 @@ enum ArenaAllocatorIn { transfer: ActorAddress, kind: AllocationKind, }, + AllocateStream { + binding: ActorAddress, + capacity: u64, + }, + ReleaseStream { + binding: ActorAddress, + ring: RingHandle, + }, ValidateSealed { binding: ActorAddress, lease: BlobLease, @@ -664,6 +844,65 @@ impl ArenaAllocatorActor { ctx.spawn(source) .map_err(|error| DataPlaneError::SessionFailed(error.to_string())) } + + fn allocate_stream(&mut self, capacity: u64) -> Result { + let request_id = self.next_request_id; + self.next_request_id = self.next_request_id.checked_add(1).ok_or_else(|| { + DataPlaneError::SessionFailed("stream request id exhausted".to_owned()) + })?; + let generation = self.next_generation; + self.next_generation = self + .next_generation + .checked_add(1) + .filter(|next| *next != 0) + .ok_or_else(|| { + DataPlaneError::SessionFailed("stream generation exhausted".to_owned()) + })?; + byte_ring::install( + &mut self.arena, + ByteRingSpec { + capacity, + generation, + alignment: BLOB_ALIGNMENT, + request_id, + }, + ) + .map_err(|error| match error { + byte_ring::InstallError::LeaseRejected(_) | byte_ring::InstallError::LeaseQueued => { + DataPlaneError::ArenaExhausted + } + other => { + DataPlaneError::SessionFailed(format!("stream ring installation failed: {other:?}")) + } + }) + } + + fn release_stream(&mut self, ring: RingHandle) -> Result<(), DataPlaneError> { + let ring_id = RingId(ring.lease_id); + let Some(allocation) = self.arena.lookup_lease(ring_id) else { + return Err(DataPlaneError::StreamFault( + "stream arena lease is no longer live".to_owned(), + )); + }; + if allocation.layout.start_offset != ring.offset + || allocation.layout.data_bytes != ring.capacity + { + return Err(DataPlaneError::StreamFault( + "stream arena lease does not match ring handle".to_owned(), + )); + } + let events = self.arena.request(ArenaRequest::ReleaseRing { + ring_id, + proof: QuiescenceProof::verified(), + }); + if matches!(events.as_slice(), [ArenaEvent::RingReleased { .. }]) { + Ok(()) + } else { + Err(DataPlaneError::SessionFailed( + "arena rejected stream release".to_owned(), + )) + } + } } impl ActorInterface for ArenaAllocatorActor { @@ -699,6 +938,14 @@ impl ActorInterface for ArenaAllocatorActor { let _ = ctx.send(transfer, BlobTransferEvent::AllocatorFailed(error.into())); } } + ArenaAllocatorIn::AllocateStream { binding, capacity } => { + let result = self.allocate_stream(capacity); + let _ = ctx.send(binding, HostStreamIn::Allocated(result)); + } + ArenaAllocatorIn::ReleaseStream { binding, ring } => { + let result = self.release_stream(ring); + let _ = ctx.send(binding, HostStreamIn::ReleaseComplete(result)); + } ArenaAllocatorIn::SealTransfer { transfer, lease, @@ -1056,7 +1303,6 @@ impl ActorInterface for NamespacePublishActor { Err(error) => HostBindingIn::PublicationRejected(namespace_error(error)), }; let _ = ctx.send(self.binding, response); - ctx.stop_self(); } } @@ -1118,9 +1364,557 @@ impl ActorInterface for NamespaceUnpublishActor { } } +struct NamespaceStreamOpenActor { + proxy: ActorAddress, + parent: ActorAddress, + path: DataPath, + role: StreamRole, + replace: bool, + descriptor: Vec, + operation_id: OperationId, + completed: bool, +} + +impl ActorInterface for NamespaceStreamOpenActor { + type Incoming = DataDirectoryOut; + type Response = (); + + fn on_start(&mut self, ctx: &Ctx<'_>) { + let _ = ctx.send( + self.proxy, + NamespaceClientIn::Request { + request: NamespaceRequest::OpenStream { + path: self.path.clone(), + role: self.role, + endpoint: self.parent, + descriptor: self.descriptor.clone(), + replace: self.replace, + operation_id: self.operation_id, + }, + reply_to: ctx.self_addr(), + }, + ); + } + + fn handle(&mut self, ctx: &Ctx<'_>, message: DataDirectoryOut) { + let result = match message { + DataDirectoryOut::StreamOpened { result, .. } => result, + other => Err(NamespaceError::Protocol(format!( + "expected stream-open reply, received {other:?}" + ))), + }; + self.completed = true; + let _ = ctx.send(self.parent, HostStreamIn::NamespaceMatched(result)); + ctx.stop_self(); + } + + fn on_stop(&mut self, ctx: &Ctx<'_>) { + if !self.completed { + let _ = ctx.send( + self.proxy, + NamespaceClientIn::Cancel { + reply_to: ctx.self_addr(), + }, + ); + } + } +} + +struct NamespaceStreamCloseActor { + proxy: ActorAddress, + path: DataPath, + incarnation: StreamIncarnation, +} + +impl ActorInterface for NamespaceStreamCloseActor { + type Incoming = DataDirectoryOut; + type Response = (); + + fn on_start(&mut self, ctx: &Ctx<'_>) { + let _ = ctx.send( + self.proxy, + NamespaceClientIn::Request { + request: NamespaceRequest::CloseStream { + path: self.path.clone(), + incarnation: self.incarnation, + }, + reply_to: ctx.self_addr(), + }, + ); + } + + fn handle(&mut self, ctx: &Ctx<'_>, _message: DataDirectoryOut) { + ctx.stop_self(); + } +} + +struct RuntimeStreamNotifier { + runtime: Runtime, + target: ActorAddress, +} + +impl StreamTransportNotifier for RuntimeStreamNotifier { + fn notify(&self, event: StreamTransportEvent) { + let _ = self + .runtime + .send_to(self.target, HostStreamIn::Transport(event)); + } +} + +struct HostStreamBindingActor { + runtime: Runtime, + host_session: ActorAddress, + allocator: ActorAddress, + child_session: ActorAddress, + operation: ActorAddress, + path: DataPath, + replace: bool, + role: StreamRole, + namespace: NamespaceClient, + arena: Arc, + transport: Arc, + local_descriptor: StreamPeerDescriptor, + ring: Option, + matched: Option, + peer_descriptor: Option, + transport_installed: bool, + transport_ready: bool, + opened: bool, + transport_quiesced: bool, + terminal: Option, + data_waiters: Vec, + capacity_waiters: Vec, + close_waiters: Vec, + release_started: bool, +} + +impl HostStreamBindingActor { + fn operation_id(address: ActorAddress) -> OperationId { + let mut bytes = [0_u8; 16]; + bytes.copy_from_slice(&address.0[..16]); + OperationId::from_u128(u128::from_le_bytes(bytes)) + } + + fn fail_open(&mut self, ctx: &Ctx<'_>, error: DataPlaneError) { + if !self.opened { + let _ = ctx.send( + self.child_session, + ChildSessionIn::OperationFailed { + operation: self.operation, + error: error.clone(), + }, + ); + } + self.begin_terminal(ctx, error, true); + } + + fn try_install_transport(&mut self, ctx: &Ctx<'_>) { + if self.transport_installed || self.terminal.is_some() { + return; + } + let (Some(ring), Some(matched)) = (self.ring, self.matched.clone()) else { + return; + }; + let endpoint_role = match self.role { + StreamRole::Source => Role::Consumer, + StreamRole::Sink => Role::Producer, + }; + let endpoint = match byte_ring::attach_mapped(&self.arena, ring, endpoint_role) { + Ok(endpoint) => endpoint, + Err(error) => { + self.fail_open( + ctx, + DataPlaneError::StreamFault(format!("attach host stream ring: {error:?}")), + ); + return; + } + }; + let notifier: Arc = Arc::new(RuntimeStreamNotifier { + runtime: self.runtime.clone(), + target: ctx.self_addr(), + }); + let install = match self.role { + StreamRole::Source => { + let Some(peer) = self.peer_descriptor.clone() else { + return; + }; + self.transport.install_source(StreamSourceRequest { + incarnation: matched.incarnation, + peer, + endpoint, + notifier, + }) + } + StreamRole::Sink => self.transport.install_sink(StreamSinkRequest { + incarnation: matched.incarnation, + endpoint, + notifier, + }), + }; + match install { + Ok(()) => { + self.transport_installed = true; + if self.role == StreamRole::Sink && matched.sink_descriptor.is_empty() { + let _ = ctx.send( + matched.source, + HostStreamIn::PeerOffer { + incarnation: matched.incarnation, + descriptor: self.local_descriptor.clone(), + }, + ); + } + } + Err(reason) => self.fail_open(ctx, DataPlaneError::StreamFault(reason)), + } + } + + fn open_if_ready(&mut self, ctx: &Ctx<'_>) { + if self.opened || !self.transport_ready || self.terminal.is_some() { + return; + } + let Some(ring) = self.ring else { + return; + }; + self.opened = true; + let _ = ctx.send( + self.child_session, + ChildSessionIn::StreamOpened { + operation: self.operation, + host_binding: ctx.self_addr(), + ring, + role: match self.role { + StreamRole::Source => Role::Producer, + StreamRole::Sink => Role::Consumer, + }, + }, + ); + } + + fn wake_waiters( + ctx: &Ctx<'_>, + child_session: ActorAddress, + waiters: &mut Vec, + result: Result<(), DataPlaneError>, + ) { + for reply_to in waiters.drain(..) { + let _ = ctx.send( + child_session, + ChildSessionIn::StreamWake { + reply_to, + result: result.clone(), + }, + ); + } + } + + fn begin_terminal(&mut self, ctx: &Ctx<'_>, error: DataPlaneError, notify_peer: bool) { + if self.terminal.is_some() { + return; + } + self.terminal = Some(error.clone()); + Self::wake_waiters( + ctx, + self.child_session, + &mut self.data_waiters, + Err(error.clone()), + ); + Self::wake_waiters( + ctx, + self.child_session, + &mut self.capacity_waiters, + Err(error.clone()), + ); + if let Some(matched) = &self.matched { + if notify_peer { + let peer = match self.role { + StreamRole::Source => matched.sink, + StreamRole::Sink => matched.source, + }; + let peer_error = if matches!(error, DataPlaneError::StreamClosed) { + DataPlaneError::StreamClosed + } else { + DataPlaneError::PeerLost + }; + let _ = ctx.send( + peer, + HostStreamIn::PeerTerminated { + incarnation: matched.incarnation, + error: peer_error, + }, + ); + } + let _ = ctx.spawn(NamespaceStreamCloseActor { + proxy: self.namespace.proxy(), + path: self.path.clone(), + incarnation: matched.incarnation, + }); + if self.transport_installed { + self.transport.terminate(matched.incarnation); + if !self.transport_quiesced { + return; + } + } + } + self.release_ring(ctx); + } + + fn complete_release(&mut self, ctx: &Ctx<'_>, result: Result<(), DataPlaneError>) { + for reply_to in self.close_waiters.drain(..) { + let _ = ctx.send( + self.child_session, + ChildSessionIn::StreamWake { + reply_to, + result: result.clone(), + }, + ); + } + let _ = ctx.send( + self.host_session, + HostSessionIn::BindingDone { + binding: ctx.self_addr(), + }, + ); + ctx.stop_self(); + } + + fn release_ring(&mut self, ctx: &Ctx<'_>) { + if self.release_started { + return; + } + self.release_started = true; + if let Some(ring) = self.ring { + let _ = ctx.send( + self.allocator, + ArenaAllocatorIn::ReleaseStream { + binding: ctx.self_addr(), + ring, + }, + ); + } else { + self.complete_release(ctx, Ok(())); + } + } +} + +impl ActorInterface for HostStreamBindingActor { + type Incoming = HostStreamIn; + type Response = (); + + fn on_start(&mut self, ctx: &Ctx<'_>) { + let _ = ctx.send( + self.allocator, + ArenaAllocatorIn::AllocateStream { + binding: ctx.self_addr(), + capacity: STREAM_RING_CAPACITY, + }, + ); + let _ = ctx.spawn(NamespaceStreamOpenActor { + proxy: self.namespace.proxy(), + parent: ctx.self_addr(), + path: self.path.clone(), + role: self.role, + descriptor: self.local_descriptor.0.clone(), + replace: self.replace, + operation_id: Self::operation_id(ctx.self_addr()), + completed: false, + }); + } + + fn handle(&mut self, ctx: &Ctx<'_>, message: HostStreamIn) { + match message { + HostStreamIn::NamespaceMatched(result) => match result { + Ok(matched) => { + let expected = match self.role { + StreamRole::Source => matched.source, + StreamRole::Sink => matched.sink, + }; + if expected != ctx.self_addr() { + self.fail_open( + ctx, + DataPlaneError::StreamFault( + "namespace matched the wrong stream endpoint".to_owned(), + ), + ); + return; + } + if self.role == StreamRole::Source && !matched.sink_descriptor.is_empty() { + self.peer_descriptor = + Some(StreamPeerDescriptor(matched.sink_descriptor.clone())); + } + self.matched = Some(matched); + self.try_install_transport(ctx); + } + Err(error) => self.fail_open(ctx, namespace_error(error)), + }, + HostStreamIn::Allocated(result) => match result { + Ok(ring) => { + self.ring = Some(ring); + self.try_install_transport(ctx); + } + Err(error) => self.fail_open(ctx, error), + }, + HostStreamIn::PeerOffer { + incarnation, + descriptor, + } => { + if self.role != StreamRole::Source { + return; + } + if self + .matched + .as_ref() + .is_some_and(|matched| matched.incarnation == incarnation) + { + self.peer_descriptor = Some(descriptor); + self.try_install_transport(ctx); + } + } + HostStreamIn::Transport(StreamTransportEvent::Ready) => { + self.transport_ready = true; + self.open_if_ready(ctx); + } + HostStreamIn::Transport(StreamTransportEvent::DataAvailable) => { + Self::wake_waiters(ctx, self.child_session, &mut self.data_waiters, Ok(())); + } + HostStreamIn::Transport(StreamTransportEvent::CapacityAvailable) => { + Self::wake_waiters(ctx, self.child_session, &mut self.capacity_waiters, Ok(())); + } + HostStreamIn::Transport(StreamTransportEvent::Fault(reason)) => { + self.fail_open(ctx, DataPlaneError::StreamFault(reason)); + } + HostStreamIn::Transport(StreamTransportEvent::Quiesced) => { + self.transport_quiesced = true; + if self.terminal.is_some() { + self.release_ring(ctx); + } + } + HostStreamIn::DataAvailable => { + if let Some(matched) = &self.matched { + self.transport.source_progress(matched.incarnation); + } + } + HostStreamIn::CapacityAvailable => { + if let Some(matched) = &self.matched { + self.transport.sink_progress(matched.incarnation); + } + } + HostStreamIn::WaitData { reply_to } => { + let result = self + .terminal + .as_ref() + .map(|error| Err(error.clone())) + .or_else(|| { + self.matched + .as_ref() + .filter(|matched| self.transport.sink_has_data(matched.incarnation)) + .map(|_| Ok(())) + }); + if let Some(result) = result { + let _ = ctx.send( + self.child_session, + ChildSessionIn::StreamWake { reply_to, result }, + ); + } else { + self.data_waiters.push(reply_to); + } + } + HostStreamIn::WaitCapacity { reply_to } => { + let result = self + .terminal + .as_ref() + .map(|error| Err(error.clone())) + .or_else(|| { + self.matched + .as_ref() + .filter(|matched| { + self.transport.source_has_capacity(matched.incarnation) + }) + .map(|_| Ok(())) + }); + if let Some(result) = result { + let _ = ctx.send( + self.child_session, + ChildSessionIn::StreamWake { reply_to, result }, + ); + } else { + self.capacity_waiters.push(reply_to); + } + } + HostStreamIn::Close { clean, reply_to } => { + if let Some(reply_to) = reply_to { + self.close_waiters.push(reply_to); + } + let error = if clean { + DataPlaneError::StreamClosed + } else { + DataPlaneError::OperationCancelled + }; + self.fail_open(ctx, error); + } + HostStreamIn::PeerTerminated { incarnation, error } => { + if self + .matched + .as_ref() + .is_some_and(|matched| matched.incarnation == incarnation) + { + if !self.opened { + let _ = ctx.send( + self.child_session, + ChildSessionIn::OperationFailed { + operation: self.operation, + error: error.clone(), + }, + ); + } + self.begin_terminal(ctx, error, false); + } + } + HostStreamIn::ReleaseComplete(result) => { + if let Err(error) = &result + && !self.opened + { + let _ = ctx.send( + self.child_session, + ChildSessionIn::OperationFailed { + operation: self.operation, + error: error.clone(), + }, + ); + } + self.complete_release(ctx, result); + } + } + } + + fn on_stop(&mut self, _ctx: &Ctx<'_>) { + if let Some(matched) = &self.matched { + self.transport.terminate(matched.incarnation); + } + } +} + fn namespace_error(error: NamespaceError) -> DataPlaneError { match error { NamespaceError::PathNotFound(path) => DataPlaneError::PathNotFound(path), + NamespaceError::WrongEntryType { + path, + expected, + found, + } => DataPlaneError::WrongEntryType { + path, + expected, + found, + }, + NamespaceError::PathReplaced(path) => DataPlaneError::PathReplaced(path), + NamespaceError::DuplicateStreamRole { path, role } => { + DataPlaneError::SessionFailed(format!("stream path {path} already has a {role:?}")) + } + NamespaceError::StaleIncarnation { path, incarnation } => { + DataPlaneError::SessionFailed(format!( + "stream path {path} no longer names incarnation {}:{}", + incarnation.authority_epoch, incarnation.revision + )) + } NamespaceError::SourceRecovery(reason) => DataPlaneError::SourceFailure(reason), NamespaceError::DirectoryUnavailable(reason) | NamespaceError::Storage(reason) @@ -1345,6 +2139,7 @@ impl HostBlobBindingActor { fn finish_without_lease(&mut self, ctx: &Ctx<'_>, outcome: ReleaseOutcome) { self.state = HostBindingState::Released; + let read_released = matches!(outcome, ReleaseOutcome::ReadReleased); if let Some(auxiliary) = self.auxiliary.take() { let _ = ctx.stop_actor(auxiliary); } @@ -1354,6 +2149,9 @@ impl HostBlobBindingActor { ChildSessionIn::WriteAborted { operation }, ); } + if read_released { + let _ = ctx.send(self.child_session, ChildSessionIn::BlobReleased); + } let _ = ctx.send( self.host_session, HostSessionIn::BindingDone { diff --git a/crates/data-plane/src/lib.rs b/crates/data-plane/src/lib.rs index 30c8f01..3847faf 100644 --- a/crates/data-plane/src/lib.rs +++ b/crates/data-plane/src/lib.rs @@ -25,3 +25,4 @@ pub mod path; pub mod protocol; pub mod ring; pub mod source; +pub mod stream_transport; diff --git a/crates/data-plane/src/namespace.rs b/crates/data-plane/src/namespace.rs index 8316435..9523ca3 100644 --- a/crates/data-plane/src/namespace.rs +++ b/crates/data-plane/src/namespace.rs @@ -33,9 +33,51 @@ pub struct BlobBinding { pub revision: u64, } +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum EntryKind { + Blob, + Stream, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum StreamRole { + Source, + Sink, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct StreamIncarnation { + pub authority_epoch: u64, + pub revision: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct StreamMatch { + pub incarnation: StreamIncarnation, + pub source: ActorAddress, + pub source_descriptor: Vec, + pub sink_descriptor: Vec, + pub sink: ActorAddress, + pub revision: u64, +} + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub enum NamespaceError { PathNotFound(DataPath), + WrongEntryType { + path: DataPath, + expected: EntryKind, + found: EntryKind, + }, + PathReplaced(DataPath), + DuplicateStreamRole { + path: DataPath, + role: StreamRole, + }, + StaleIncarnation { + path: DataPath, + incarnation: StreamIncarnation, + }, OperationConflict(OperationId), Storage(String), SourceRecovery(String), @@ -47,6 +89,25 @@ impl fmt::Display for NamespaceError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::PathNotFound(path) => write!(f, "data path not found: {path}"), + Self::WrongEntryType { + path, + expected, + found, + } => write!( + f, + "data path {path} has entry kind {found:?}, expected {expected:?}" + ), + Self::PathReplaced(path) => { + write!(f, "pending data path was replaced: {path}") + } + Self::DuplicateStreamRole { path, role } => { + write!(f, "stream path {path} already has a {role:?}") + } + Self::StaleIncarnation { path, incarnation } => write!( + f, + "stream path {path} no longer names incarnation {}:{}", + incarnation.authority_epoch, incarnation.revision + ), Self::OperationConflict(operation) => write!( f, "namespace operation ID {:02x?} was reused for a different request", @@ -92,6 +153,26 @@ pub enum DataDirectoryIn { operation_id: OperationId, reply_to: ActorAddress, }, + OpenStream { + request_id: DirectoryRequestId, + path: DataPath, + role: StreamRole, + descriptor: Vec, + endpoint: ActorAddress, + replace: bool, + operation_id: OperationId, + reply_to: ActorAddress, + }, + CancelStream { + path: DataPath, + operation_id: OperationId, + }, + CloseStream { + request_id: DirectoryRequestId, + path: DataPath, + incarnation: StreamIncarnation, + reply_to: ActorAddress, + }, } impl NetworkMessage for DataDirectoryIn { @@ -117,14 +198,25 @@ pub enum DataDirectoryOut { authority_epoch: u64, result: Result, }, + StreamOpened { + request_id: DirectoryRequestId, + authority_epoch: u64, + result: Result, + }, + StreamClosed { + request_id: DirectoryRequestId, + authority_epoch: u64, + result: Result<(), NamespaceError>, + }, } - impl DataDirectoryOut { pub fn request_id(&self) -> DirectoryRequestId { match self { Self::Registered { request_id, .. } | Self::Resolved { request_id, .. } - | Self::Unregistered { request_id, .. } => *request_id, + | Self::Unregistered { request_id, .. } + | Self::StreamOpened { request_id, .. } + | Self::StreamClosed { request_id, .. } => *request_id, } } } @@ -145,8 +237,19 @@ pub enum NamespaceRequest { path: DataPath, operation_id: OperationId, }, + OpenStream { + path: DataPath, + role: StreamRole, + endpoint: ActorAddress, + descriptor: Vec, + replace: bool, + operation_id: OperationId, + }, + CloseStream { + path: DataPath, + incarnation: StreamIncarnation, + }, } - #[derive(Clone, Debug, Serialize, Deserialize)] pub enum NamespaceClientIn { Request { @@ -177,9 +280,42 @@ enum RuntimeSource { Unavailable(String), } +struct PendingStream { + role: StreamRole, + endpoint: ActorAddress, + operation_id: OperationId, + request_id: DirectoryRequestId, + descriptor: Vec, + reply_to: ActorAddress, + revision: u64, +} + +struct StreamOpenRequest { + request_id: DirectoryRequestId, + path: DataPath, + role: StreamRole, + endpoint: ActorAddress, + replace: bool, + descriptor: Vec, + operation_id: OperationId, + reply_to: ActorAddress, +} + +struct ActiveStream { + binding: StreamMatch, + source_operation: OperationId, + sink_operation: OperationId, +} + +enum RuntimeStream { + Pending(PendingStream), + Active(ActiveStream), +} + pub struct DataDirectoryActor { store: NamespaceStore, sources: BTreeMap, + streams: BTreeMap, authority_epoch: u64, } @@ -202,6 +338,7 @@ impl DataDirectoryActor { Ok(Self { store, sources, + streams: BTreeMap::new(), authority_epoch, }) } @@ -232,6 +369,257 @@ impl DataDirectoryActor { }) } + fn bind_stream( + &mut self, + path: DataPath, + operation_id: OperationId, + retired: Option, + ) -> Result { + let request = MutationRequest::BindStream { path: path.clone() }; + if let Some(replayed) = self.replay(operation_id, &request) { + return replayed; + } + let revision = self.store.snapshot().next_revision; + let next_revision = revision + .checked_add(1) + .filter(|revision| *revision != 0) + .ok_or(NamespaceStoreError::RevisionExhausted)?; + let receipt = MutationReceipt { revision }; + let mut next = self.store.snapshot().clone(); + next.next_revision = next_revision; + next.bindings.remove(&path); + if let Some(retired) = retired + && !next.retirements.contains(&retired) + { + next.retirements.push(retired); + } + next.operations.insert( + operation_id, + PersistedOperation { + request, + result: PersistedMutationResult::Committed(receipt), + }, + ); + self.store.commit(next)?; + self.sources.remove(&path); + Ok(receipt) + } + + fn send_stream_result( + &self, + ctx: &Ctx<'_>, + request_id: DirectoryRequestId, + reply_to: ActorAddress, + result: Result, + ) { + let _ = ctx.send( + reply_to, + NamespaceClientIn::DirectoryReply(DataDirectoryOut::StreamOpened { + request_id, + authority_epoch: self.authority_epoch, + result, + }), + ); + } + + fn displace_stream(&mut self, ctx: &Ctx<'_>, path: &DataPath) { + if let Some(RuntimeStream::Pending(pending)) = self.streams.remove(path) { + self.send_stream_result( + ctx, + pending.request_id, + pending.reply_to, + Err(NamespaceError::PathReplaced(path.clone())), + ); + } + } + + fn open_stream(&mut self, ctx: &Ctx<'_>, request: StreamOpenRequest) { + let StreamOpenRequest { + request_id, + path, + role, + endpoint, + replace, + descriptor, + operation_id, + reply_to, + } = request; + if let Some(RuntimeStream::Pending(pending)) = self.streams.get_mut(&path) + && pending.operation_id == operation_id + { + if pending.role != role || pending.endpoint != endpoint { + self.send_stream_result( + ctx, + request_id, + reply_to, + Err(NamespaceError::OperationConflict(operation_id)), + ); + } else { + pending.request_id = request_id; + pending.reply_to = reply_to; + } + return; + } + if let Some(RuntimeStream::Active(active)) = self.streams.get(&path) + && (active.source_operation == operation_id || active.sink_operation == operation_id) + { + self.send_stream_result(ctx, request_id, reply_to, Ok(active.binding.clone())); + return; + } + + let compatible_pending = matches!( + self.streams.get(&path), + Some(RuntimeStream::Pending(pending)) if pending.role != role + ); + if replace && self.streams.contains_key(&path) && !compatible_pending { + self.displace_stream(ctx, &path); + } + + if let Some(RuntimeStream::Pending(pending)) = self.streams.remove(&path) { + if pending.role == role { + self.streams + .insert(path.clone(), RuntimeStream::Pending(pending)); + self.send_stream_result( + ctx, + request_id, + reply_to, + Err(NamespaceError::DuplicateStreamRole { path, role }), + ); + return; + } + let ( + source, + sink, + source_descriptor, + sink_descriptor, + source_operation, + sink_operation, + ) = match role { + StreamRole::Source => ( + endpoint, + pending.endpoint, + descriptor, + pending.descriptor, + operation_id, + pending.operation_id, + ), + StreamRole::Sink => ( + pending.endpoint, + endpoint, + pending.descriptor, + descriptor, + pending.operation_id, + operation_id, + ), + }; + let binding = StreamMatch { + incarnation: StreamIncarnation { + authority_epoch: self.authority_epoch, + revision: pending.revision, + }, + source, + source_descriptor, + sink_descriptor, + sink, + revision: pending.revision, + }; + self.send_stream_result( + ctx, + pending.request_id, + pending.reply_to, + Ok(binding.clone()), + ); + self.send_stream_result(ctx, request_id, reply_to, Ok(binding.clone())); + self.streams.insert( + path, + RuntimeStream::Active(ActiveStream { + binding, + source_operation, + sink_operation, + }), + ); + return; + } + + if self.streams.contains_key(&path) { + self.send_stream_result( + ctx, + request_id, + reply_to, + Err(NamespaceError::DuplicateStreamRole { path, role }), + ); + return; + } + + if self.store.snapshot().bindings.contains_key(&path) && !replace { + self.send_stream_result( + ctx, + request_id, + reply_to, + Err(NamespaceError::WrongEntryType { + path, + expected: EntryKind::Stream, + found: EntryKind::Blob, + }), + ); + return; + } + let retired = self.sources.get(&path).and_then(|source| match source { + RuntimeSource::Available(actor) => Some(*actor), + RuntimeSource::Unavailable(_) => None, + }); + match self.bind_stream(path.clone(), operation_id, retired) { + Ok(receipt) => { + if let Some(retired) = retired { + let _ = ctx.send(retired, BlobSourceIn::Retire); + } + self.streams.insert( + path, + RuntimeStream::Pending(PendingStream { + role, + descriptor, + endpoint, + operation_id, + request_id, + reply_to, + revision: receipt.revision, + }), + ); + } + Err(error) => self.send_stream_result(ctx, request_id, reply_to, Err(error)), + } + } + + fn cancel_stream(&mut self, path: &DataPath, operation_id: OperationId) { + let should_remove = matches!( + self.streams.get(path), + Some(RuntimeStream::Pending(pending)) if pending.operation_id == operation_id + ); + if should_remove { + self.streams.remove(path); + } + } + + fn close_stream( + &mut self, + path: &DataPath, + incarnation: StreamIncarnation, + ) -> Result<(), NamespaceError> { + let matches = matches!( + self.streams.get(path), + Some(RuntimeStream::Active(active)) if active.binding.incarnation == incarnation + ); + if matches { + self.streams.remove(path); + Ok(()) + } else { + Err(NamespaceError::StaleIncarnation { + path: path.clone(), + incarnation, + }) + } + } + fn register( &mut self, path: DataPath, @@ -283,6 +671,13 @@ impl DataDirectoryActor { } fn resolve(&self, path: &DataPath) -> Result { + if self.streams.contains_key(path) { + return Err(NamespaceError::WrongEntryType { + path: path.clone(), + expected: EntryKind::Blob, + found: EntryKind::Stream, + }); + } let persisted = self .store .snapshot() @@ -376,6 +771,7 @@ impl ActorInterface for DataDirectoryActor { operation_id, reply_to, } => { + let logical = path.clone(); let replayed = self.store.snapshot().operations.contains_key(&operation_id); let retired = (!replayed) .then(|| self.sources.get(&path)) @@ -390,6 +786,9 @@ impl ActorInterface for DataDirectoryActor { { let _ = ctx.send(retired, BlobSourceIn::Retire); } + if result.is_ok() { + self.displace_stream(ctx, &logical); + } let _ = ctx.send( reply_to, NamespaceClientIn::DirectoryReply(DataDirectoryOut::Registered { @@ -420,6 +819,21 @@ impl ActorInterface for DataDirectoryActor { operation_id, reply_to, } => { + if self.streams.contains_key(&path) { + let _ = ctx.send( + reply_to, + NamespaceClientIn::DirectoryReply(DataDirectoryOut::Unregistered { + request_id, + authority_epoch: self.authority_epoch, + result: Err(NamespaceError::WrongEntryType { + path, + expected: EntryKind::Blob, + found: EntryKind::Stream, + }), + }), + ); + return; + } let replayed = self.store.snapshot().operations.contains_key(&operation_id); let retired = (!replayed) .then(|| self.sources.get(&path)) @@ -443,6 +857,47 @@ impl ActorInterface for DataDirectoryActor { }), ); } + DataDirectoryIn::OpenStream { + request_id, + path, + role, + endpoint, + descriptor, + replace, + operation_id, + reply_to, + } => self.open_stream( + ctx, + StreamOpenRequest { + request_id, + path, + role, + endpoint, + replace, + descriptor, + operation_id, + reply_to, + }, + ), + DataDirectoryIn::CancelStream { path, operation_id } => { + self.cancel_stream(&path, operation_id); + } + DataDirectoryIn::CloseStream { + request_id, + path, + incarnation, + reply_to, + } => { + let result = self.close_stream(&path, incarnation); + let _ = ctx.send( + reply_to, + NamespaceClientIn::DirectoryReply(DataDirectoryOut::StreamClosed { + request_id, + authority_epoch: self.authority_epoch, + result, + }), + ); + } } } } @@ -513,6 +968,29 @@ impl NamespaceClientActor { operation_id: *operation_id, reply_to: ctx.self_addr(), }, + NamespaceRequest::OpenStream { + path, + role, + endpoint, + descriptor, + replace, + operation_id, + } => DataDirectoryIn::OpenStream { + request_id, + path: path.clone(), + role: *role, + endpoint: *endpoint, + descriptor: descriptor.clone(), + replace: *replace, + operation_id: *operation_id, + reply_to: ctx.self_addr(), + }, + NamespaceRequest::CloseStream { path, incarnation } => DataDirectoryIn::CloseStream { + request_id, + path: path.clone(), + incarnation: *incarnation, + reply_to: ctx.self_addr(), + }, }; let _ = ctx.send(directory, message); } @@ -558,6 +1036,25 @@ impl ActorInterface for NamespaceClientActor { .insert(request_id, PendingRequest { request, reply_to }); } NamespaceClientIn::Cancel { reply_to } => { + let cancelled: Vec<(DataPath, OperationId)> = self + .pending + .values() + .filter(|pending| pending.reply_to == reply_to) + .filter_map(|pending| match &pending.request { + NamespaceRequest::OpenStream { + path, operation_id, .. + } => Some((path.clone(), *operation_id)), + _ => None, + }) + .collect(); + if let Some(directory) = self.discovery.current_directory() { + for (path, operation_id) in cancelled { + let _ = ctx.send( + directory, + DataDirectoryIn::CancelStream { path, operation_id }, + ); + } + } self.pending .retain(|_, pending| pending.reply_to != reply_to); } @@ -692,6 +1189,92 @@ impl NamespaceClient { ))), } } + + async fn open_stream_inner( + &self, + path: DataPath, + role: StreamRole, + endpoint: ActorAddress, + replace: bool, + operation_id: OperationId, + ) -> Result { + match self + .request(NamespaceRequest::OpenStream { + path, + role, + endpoint, + replace, + descriptor: Vec::new(), + operation_id, + }) + .await? + { + DataDirectoryOut::StreamOpened { result, .. } => result, + other => Err(NamespaceError::Protocol(format!( + "expected stream-open reply, received {other:?}" + ))), + } + } + + pub async fn open_stream( + &self, + path: DataPath, + role: StreamRole, + endpoint: ActorAddress, + operation_id: OperationId, + ) -> Result { + self.open_stream_inner(path, role, endpoint, false, operation_id) + .await + } + + pub async fn replace_with_stream( + &self, + path: DataPath, + role: StreamRole, + endpoint: ActorAddress, + operation_id: OperationId, + ) -> Result { + self.open_stream_inner(path, role, endpoint, true, operation_id) + .await + } + + pub async fn close_stream( + &self, + path: DataPath, + incarnation: StreamIncarnation, + ) -> Result<(), NamespaceError> { + match self + .request(NamespaceRequest::CloseStream { path, incarnation }) + .await? + { + DataDirectoryOut::StreamClosed { result, .. } => result, + other => Err(NamespaceError::Protocol(format!( + "expected stream-close reply, received {other:?}" + ))), + } + } +} + +struct DirectoryStreamCancellation { + runtime: Runtime, + directory: ActorAddress, + path: DataPath, + operation_id: OperationId, + armed: bool, +} + +impl Drop for DirectoryStreamCancellation { + fn drop(&mut self) { + if self.armed { + let _ = self.runtime.send_to( + self.directory, + DataDirectoryIn::CancelStream { + path: self.path.clone(), + operation_id: self.operation_id, + }, + ); + } + } } #[derive(Clone)] @@ -814,6 +1397,100 @@ impl DirectoryClient { ))), } } + + async fn open_stream_inner( + &self, + path: DataPath, + role: StreamRole, + endpoint: ActorAddress, + replace: bool, + operation_id: OperationId, + ) -> Result { + let inbox = self + .runtime + .new_inbox::() + .map_err(|error| NamespaceError::DirectoryUnavailable(error.to_string()))?; + let mut cancellation = DirectoryStreamCancellation { + runtime: self.runtime.clone(), + directory: self.directory, + path: path.clone(), + operation_id, + armed: true, + }; + self.runtime + .send_to( + self.directory, + DataDirectoryIn::OpenStream { + request_id: self.request_id(), + path, + role, + endpoint, + replace, + operation_id, + descriptor: Vec::new(), + reply_to: *inbox.addr(), + }, + ) + .map_err(|error| NamespaceError::DirectoryUnavailable(error.to_string()))?; + let result = match self.receive(&inbox).await? { + DataDirectoryOut::StreamOpened { result, .. } => result, + other => Err(NamespaceError::Protocol(format!( + "expected stream-open reply, received {other:?}" + ))), + }; + cancellation.armed = false; + result + } + + pub async fn open_stream( + &self, + path: DataPath, + role: StreamRole, + endpoint: ActorAddress, + operation_id: OperationId, + ) -> Result { + self.open_stream_inner(path, role, endpoint, false, operation_id) + .await + } + + pub async fn replace_with_stream( + &self, + path: DataPath, + role: StreamRole, + endpoint: ActorAddress, + operation_id: OperationId, + ) -> Result { + self.open_stream_inner(path, role, endpoint, true, operation_id) + .await + } + + pub async fn close_stream( + &self, + path: DataPath, + incarnation: StreamIncarnation, + ) -> Result<(), NamespaceError> { + let inbox = self + .runtime + .new_inbox::() + .map_err(|error| NamespaceError::DirectoryUnavailable(error.to_string()))?; + self.runtime + .send_to( + self.directory, + DataDirectoryIn::CloseStream { + request_id: self.request_id(), + path, + incarnation, + reply_to: *inbox.addr(), + }, + ) + .map_err(|error| NamespaceError::DirectoryUnavailable(error.to_string()))?; + match self.receive(&inbox).await? { + DataDirectoryOut::StreamClosed { result, .. } => result, + other => Err(NamespaceError::Protocol(format!( + "expected stream-close reply, received {other:?}" + ))), + } + } } pub fn register_namespace_codecs(registry: &mut CodecRegistry) { diff --git a/crates/data-plane/src/namespace_store.rs b/crates/data-plane/src/namespace_store.rs index 94ed599..164f0aa 100644 --- a/crates/data-plane/src/namespace_store.rs +++ b/crates/data-plane/src/namespace_store.rs @@ -72,6 +72,9 @@ pub enum MutationRequest { length: u64, recovery: SourceRecovery, }, + BindStream { + path: DataPath, + }, Unregister { path: DataPath, }, diff --git a/crates/data-plane/src/protocol.rs b/crates/data-plane/src/protocol.rs index b34d175..3ad0294 100644 --- a/crates/data-plane/src/protocol.rs +++ b/crates/data-plane/src/protocol.rs @@ -7,8 +7,11 @@ use swactor::actor::ActorAddress; use swactor_transport::{CodecRegistry, JsonCodec, NetworkMessage}; use crate::blob::{BlobError, BlobLease, BlobMetadata}; +use crate::byte_ring::{RingHandle, Role}; use crate::ids::BlobLeaseId; +use crate::namespace::{EntryKind, NamespaceError, StreamIncarnation, StreamMatch}; use crate::path::DataPath; +use crate::stream_transport::{StreamPeerDescriptor, StreamTransportEvent}; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct JobCapability([u8; 32]); @@ -108,7 +111,15 @@ pub enum DataPlaneError { ArenaExhausted, Blob(BlobFailure), OperationCancelled, - StreamsDeferred, + WrongEntryType { + path: DataPath, + expected: EntryKind, + found: EntryKind, + }, + PathReplaced(DataPath), + PeerLost, + StreamFault(String), + StreamClosed, } impl fmt::Display for DataPlaneError { @@ -127,7 +138,18 @@ impl fmt::Display for DataPlaneError { Self::ArenaExhausted => f.write_str("data-plane arena is exhausted"), Self::Blob(reason) => write!(f, "blob lease failure: {reason:?}"), Self::OperationCancelled => f.write_str("data-plane operation was cancelled"), - Self::StreamsDeferred => f.write_str("actor-driven streams are not installed"), + Self::WrongEntryType { + path, + expected, + found, + } => write!( + f, + "data path {path} has entry kind {found:?}, expected {expected:?}" + ), + Self::PathReplaced(path) => write!(f, "data path was replaced: {path}"), + Self::PeerLost => f.write_str("stream peer was lost"), + Self::StreamFault(reason) => write!(f, "stream fault: {reason}"), + Self::StreamClosed => f.write_str("stream is closed"), } } } @@ -162,6 +184,25 @@ pub enum HostSessionIn { child_session: ActorAddress, operation: ActorAddress, }, + OpenReadStream { + path: DataPath, + child_session: ActorAddress, + operation: ActorAddress, + replace: bool, + }, + OpenWriteStream { + path: DataPath, + child_session: ActorAddress, + operation: ActorAddress, + replace: bool, + }, + CancelStream { + operation: ActorAddress, + }, + StreamControl { + binding: ActorAddress, + message: HostStreamIn, + }, ReleaseBlob { binding: ActorAddress, lease_id: BlobLeaseId, @@ -224,6 +265,33 @@ pub enum ChildSessionIn { length: u64, reply_to: ActorAddress, }, + OpenReadStream { + path: DataPath, + reply_to: ActorAddress, + replace: bool, + }, + OpenWriteStream { + path: DataPath, + reply_to: ActorAddress, + replace: bool, + }, + CancelStream { + reply_to: ActorAddress, + }, + StreamWake { + reply_to: ActorAddress, + result: Result<(), DataPlaneError>, + }, + StreamControl { + binding: ActorAddress, + message: HostStreamIn, + }, + ReleaseBlob { + binding: ActorAddress, + lease_id: BlobLeaseId, + generation: u64, + }, + BlobReleased, BlobOpened { operation: ActorAddress, host_binding: ActorAddress, @@ -236,6 +304,12 @@ pub enum ChildSessionIn { lease: BlobLease, metadata: BlobMetadata, }, + StreamOpened { + operation: ActorAddress, + host_binding: ActorAddress, + ring: RingHandle, + role: Role, + }, OperationFailed { operation: ActorAddress, error: DataPlaneError, @@ -258,9 +332,44 @@ impl NetworkMessage for ChildSessionIn { } } +#[derive(Clone, Debug, Serialize, Deserialize)] +pub enum HostStreamIn { + NamespaceMatched(Result), + Allocated(Result), + PeerOffer { + incarnation: StreamIncarnation, + descriptor: StreamPeerDescriptor, + }, + Transport(StreamTransportEvent), + DataAvailable, + CapacityAvailable, + WaitData { + reply_to: ActorAddress, + }, + WaitCapacity { + reply_to: ActorAddress, + }, + Close { + clean: bool, + reply_to: Option, + }, + PeerTerminated { + incarnation: StreamIncarnation, + error: DataPlaneError, + }, + ReleaseComplete(Result<(), DataPlaneError>), +} + +impl NetworkMessage for HostStreamIn { + fn type_tag() -> &'static str { + "data-plane.host-stream.v1" + } +} + pub fn register_data_plane_codecs(registry: &mut CodecRegistry) { registry.register::(JsonCodec::default()); registry.register::(JsonCodec::default()); + registry.register::(JsonCodec::default()); crate::namespace::register_namespace_codecs(registry); crate::blob_transfer::register_blob_transfer_codecs(registry); crate::source::register_blob_source_codecs(registry); diff --git a/crates/data-plane/src/stream_transport.rs b/crates/data-plane/src/stream_transport.rs new file mode 100644 index 0000000..51f4abf --- /dev/null +++ b/crates/data-plane/src/stream_transport.rs @@ -0,0 +1,311 @@ +//! Replaceable bounded transport port for SPSC stream incarnations. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; + +use crate::byte_ring::{Endpoint, FlowError}; +use crate::namespace::StreamIncarnation; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct StreamPeerDescriptor(pub Vec); + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum StreamTransportEvent { + Ready, + DataAvailable, + CapacityAvailable, + Quiesced, + Fault(String), +} + +pub trait StreamTransportNotifier: Send + Sync + 'static { + fn notify(&self, event: StreamTransportEvent); +} + +pub struct StreamSourceRequest { + pub incarnation: StreamIncarnation, + pub peer: StreamPeerDescriptor, + pub endpoint: Endpoint, + pub notifier: Arc, +} + +pub struct StreamSinkRequest { + pub incarnation: StreamIncarnation, + pub endpoint: Endpoint, + pub notifier: Arc, +} + +/// Transport effect boundary. Namespace, terminal, reconnection, and lease +/// policy stay in the data-plane actors that invoke this port. +pub trait StreamTransport: Send + Sync + 'static { + fn descriptor(&self) -> Result; + fn install_source(&self, request: StreamSourceRequest) -> Result<(), String>; + fn install_sink(&self, request: StreamSinkRequest) -> Result<(), String>; + fn source_progress(&self, incarnation: StreamIncarnation); + fn sink_progress(&self, incarnation: StreamIncarnation); + fn source_has_capacity(&self, _incarnation: StreamIncarnation) -> bool { + false + } + fn sink_has_data(&self, _incarnation: StreamIncarnation) -> bool { + false + } + fn terminate(&self, incarnation: StreamIncarnation); +} + +struct LocalSource { + endpoint: Endpoint, + notifier: Arc, +} + +struct LocalSink { + endpoint: Endpoint, + notifier: Arc, +} + +#[derive(Default)] +struct LocalTransfer { + source: Option, + sink: Option, + ready: bool, +} + +#[derive(Default)] +struct LocalState { + transfers: BTreeMap, +} + +static NEXT_LOCAL_TRANSPORT: AtomicU64 = AtomicU64::new(1); + +/// Direct in-process adapter. Host composition shares one instance between +/// sessions on a node. Payload bytes copy once from the source ring spans to +/// the destination ring spans; no payload-sized staging allocation exists. +pub struct LocalStreamTransport { + id: u64, + state: Mutex, +} + +impl Default for LocalStreamTransport { + fn default() -> Self { + Self::new() + } +} + +impl LocalStreamTransport { + pub fn new() -> Self { + Self { + id: NEXT_LOCAL_TRANSPORT.fetch_add(1, Ordering::Relaxed), + state: Mutex::new(LocalState::default()), + } + } + + fn drive(&self, incarnation: StreamIncarnation) { + let mut notifications = Vec::new(); + let mut fault = None; + { + let mut state = self.state.lock(); + let Some(transfer) = state.transfers.get_mut(&incarnation) else { + return; + }; + let (Some(source), Some(sink)) = (&mut transfer.source, &mut transfer.sink) else { + return; + }; + if !transfer.ready { + transfer.ready = true; + notifications.push((Arc::clone(&source.notifier), StreamTransportEvent::Ready)); + notifications.push((Arc::clone(&sink.notifier), StreamTransportEvent::Ready)); + } + + let mut moved = false; + loop { + let meta = match source.endpoint.next_record_meta() { + Ok(Some(meta)) => meta, + Ok(None) => break, + Err(error) => { + fault = Some(format_flow_error(error)); + break; + } + }; + let mut destination = match sink.endpoint.reserve_record(meta.kind, meta.len) { + Ok(destination) => destination, + Err(FlowError::InsufficientSpace { .. }) => break, + Err(error) => { + fault = Some(format_flow_error(error)); + break; + } + }; + let source_view = match source.endpoint.peek_record() { + Ok(Some(view)) => view, + Ok(None) => { + fault = Some("source record disappeared after inspection".to_owned()); + break; + } + Err(error) => { + fault = Some(format_flow_error(error)); + break; + } + }; + let (source_first, source_second) = source_view.spans(); + let (destination_first, destination_second) = destination.spans_mut(); + copy_spans( + source_first, + source_second, + destination_first, + destination_second, + ); + if let Err(error) = destination.commit() { + fault = Some(format_flow_error(error)); + break; + } + if let Err(error) = source_view.release() { + fault = Some(format_flow_error(error)); + break; + } + moved = true; + } + + if moved { + notifications.push(( + Arc::clone(&source.notifier), + StreamTransportEvent::CapacityAvailable, + )); + notifications.push(( + Arc::clone(&sink.notifier), + StreamTransportEvent::DataAvailable, + )); + } + if let Some(reason) = &fault { + notifications.push(( + Arc::clone(&source.notifier), + StreamTransportEvent::Fault(reason.clone()), + )); + notifications.push(( + Arc::clone(&sink.notifier), + StreamTransportEvent::Fault(reason.clone()), + )); + } + } + for (notifier, event) in notifications { + notifier.notify(event); + } + if fault.is_some() { + self.state.lock().transfers.remove(&incarnation); + } + } +} + +impl StreamTransport for LocalStreamTransport { + fn descriptor(&self) -> Result { + Ok(StreamPeerDescriptor(self.id.to_le_bytes().to_vec())) + } + + fn install_source(&self, request: StreamSourceRequest) -> Result<(), String> { + if request.peer != self.descriptor()? { + return Err("local stream peer belongs to a different transport instance".to_owned()); + } + let incarnation = request.incarnation; + let mut state = self.state.lock(); + let transfer = state.transfers.entry(incarnation).or_default(); + if transfer.source.is_some() { + return Err("stream source is already installed".to_owned()); + } + transfer.source = Some(LocalSource { + endpoint: request.endpoint, + notifier: request.notifier, + }); + drop(state); + self.drive(incarnation); + Ok(()) + } + + fn install_sink(&self, request: StreamSinkRequest) -> Result<(), String> { + let incarnation = request.incarnation; + let mut state = self.state.lock(); + let transfer = state.transfers.entry(incarnation).or_default(); + if transfer.sink.is_some() { + return Err("stream sink is already installed".to_owned()); + } + transfer.sink = Some(LocalSink { + endpoint: request.endpoint, + notifier: request.notifier, + }); + drop(state); + self.drive(incarnation); + Ok(()) + } + + fn source_progress(&self, incarnation: StreamIncarnation) { + self.drive(incarnation); + } + + fn sink_progress(&self, incarnation: StreamIncarnation) { + self.drive(incarnation); + } + + fn source_has_capacity(&self, incarnation: StreamIncarnation) -> bool { + self.state + .lock() + .transfers + .get(&incarnation) + .and_then(|transfer| transfer.source.as_ref()) + .is_some_and(|source| source.endpoint.probe().has_capacity()) + } + + fn sink_has_data(&self, incarnation: StreamIncarnation) -> bool { + self.state + .lock() + .transfers + .get(&incarnation) + .and_then(|transfer| transfer.sink.as_ref()) + .is_some_and(|sink| sink.endpoint.probe().has_data()) + } + + fn terminate(&self, incarnation: StreamIncarnation) { + let transfer = self.state.lock().transfers.remove(&incarnation); + if let Some(transfer) = transfer { + if let Some(source) = transfer.source { + source.notifier.notify(StreamTransportEvent::Quiesced); + } + if let Some(sink) = transfer.sink { + sink.notifier.notify(StreamTransportEvent::Quiesced); + } + } + } +} + +fn copy_spans( + source_first: &[u8], + source_second: &[u8], + destination_first: &mut [u8], + destination_second: &mut [u8], +) { + debug_assert_eq!( + source_first.len() + source_second.len(), + destination_first.len() + destination_second.len() + ); + let sources = [source_first, source_second]; + let mut source_index = 0; + let mut source_offset = 0; + for destination in [destination_first, destination_second] { + let mut destination_offset = 0; + while destination_offset < destination.len() { + while source_index < sources.len() && source_offset == sources[source_index].len() { + source_index += 1; + source_offset = 0; + } + let source = sources[source_index]; + let take = (destination.len() - destination_offset).min(source.len() - source_offset); + destination[destination_offset..destination_offset + take] + .copy_from_slice(&source[source_offset..source_offset + take]); + destination_offset += take; + source_offset += take; + } + } +} + +fn format_flow_error(error: FlowError) -> String { + format!("byte ring transport fault: {error:?}") +} diff --git a/crates/data-plane/tests/actor_blob_guarantees.rs b/crates/data-plane/tests/actor_blob_guarantees.rs index 93f9f0b..2337839 100755 --- a/crates/data-plane/tests/actor_blob_guarantees.rs +++ b/crates/data-plane/tests/actor_blob_guarantees.rs @@ -24,6 +24,7 @@ const CAPABILITY: JobCapability = JobCapability::new([9; 32]); const ARENA_GENERATION: u64 = 17; const SESSION_GENERATION: u64 = 29; const WEIGHTS: &[u8] = b"0123456789abcdefghijklmn"; +static STREAM_TEST_LOCK: parking_lot::Mutex<()> = parking_lot::Mutex::new(()); struct DirectRuntimeSink { destination: Runtime, @@ -167,7 +168,54 @@ impl data_plane::source::BlobSourcePublisher for NoopSourceRegistrar { } } +struct RejectingStreamTransport; + +impl data_plane::stream_transport::StreamTransport for RejectingStreamTransport { + fn descriptor(&self) -> Result { + Ok(data_plane::stream_transport::StreamPeerDescriptor(vec![1])) + } + + fn install_source( + &self, + _request: data_plane::stream_transport::StreamSourceRequest, + ) -> Result<(), String> { + Err("injected source transport failure".to_owned()) + } + + fn install_sink( + &self, + _request: data_plane::stream_transport::StreamSinkRequest, + ) -> Result<(), String> { + Err("injected sink transport failure".to_owned()) + } + + fn source_progress(&self, _incarnation: data_plane::namespace::StreamIncarnation) {} + + fn sink_progress(&self, _incarnation: data_plane::namespace::StreamIncarnation) {} + + fn terminate(&self, _incarnation: data_plane::namespace::StreamIncarnation) {} +} + +struct CollectBytes(Arc>>); + +impl data_plane::data_plane::StreamConsumer for CollectBytes { + fn consume(&self, bytes: &[u8]) -> Result<(), String> { + self.0.lock().extend_from_slice(bytes); + Ok(()) + } +} + fn harness(arena_bytes: u64) -> Harness { + harness_with_transport( + arena_bytes, + Arc::new(data_plane::stream_transport::LocalStreamTransport::new()), + ) +} + +fn harness_with_transport( + arena_bytes: u64, + stream_transport: Arc, +) -> Harness { let temp = TempState::new(); let mut arena = ArenaManager::boot(ArenaConfig { node_id: NodeId(1), @@ -194,12 +242,20 @@ fn harness(arena_bytes: u64) -> Harness { })); let host_engine = Engine::new( host_parts, - TokioBackend::new(TokioConfig::default()).expect("host backend"), + TokioBackend::new(TokioConfig { + worker_threads: 1, + ..TokioConfig::default() + }) + .expect("host backend"), ) .expect("host engine"); let child_engine = Engine::new( child_parts, - TokioBackend::new(TokioConfig::default()).expect("child backend"), + TokioBackend::new(TokioConfig { + worker_threads: 1, + ..TokioConfig::default() + }) + .expect("child backend"), ) .expect("child engine"); @@ -269,6 +325,7 @@ fn harness(arena_bytes: u64) -> Harness { source_sender: Some(sender), source_publisher: Some(Arc::new(NoopSourceRegistrar)), route_registrar: None, + stream_transport: Some(stream_transport), }) .expect("host session config"), ) @@ -314,7 +371,15 @@ fn attachment_without_a_host_reply_fails_on_actor_deadline() { .unwrap(); let (parts, runtime) = runtime_parts(); runtime.set_remote_sink(Arc::new(BlackHoleSink)); - let engine = Engine::new(parts, TokioBackend::new(TokioConfig::default()).unwrap()).unwrap(); + let engine = Engine::new( + parts, + TokioBackend::new(TokioConfig { + worker_threads: 1, + ..TokioConfig::default() + }) + .unwrap(), + ) + .unwrap(); let (mapped, resolved) = DataPlaneBootstrap::map_arena(handoff.arena_fd).unwrap(); let result = future::block_on(DataPlaneBootstrap::attach_mapped_with_deadline( mapped, @@ -618,3 +683,194 @@ fn write_blob_seals_once_and_abort_publishes_nothing() { Err(DataPlaneError::PathNotFound(_)) )); } + +#[test] +fn stream_endpoints_open_only_after_match_and_deliver_eof_in_order() { + let _stream_test = STREAM_TEST_LOCK.lock(); + let harness = harness(2 << 20); + let data_plane = harness.bootstrap.data_plane.clone(); + let logical = path("/runs/self/results/inference"); + + future::block_on(async { + let mut reader_open = Box::pin(data_plane.read_stream(&logical)); + assert!( + future::poll_once(reader_open.as_mut()).await.is_none(), + "reader open waits for its source" + ); + let mut writer_open = Box::pin(data_plane.write_stream(&logical)); + let mut writer = writer_open.as_mut().await.expect("writer opens"); + let mut reader = reader_open.await.expect("reader opens"); + + writer.write(b"first").await.expect("write first"); + writer.write(b"second").await.expect("write second"); + assert_eq!( + reader.read().await.expect("read first"), + Some(b"first".to_vec()) + ); + assert_eq!( + reader.read().await.expect("read second"), + Some(b"second".to_vec()) + ); + writer.close().await.expect("clean writer close"); + assert!(matches!( + writer.write(b"late").await, + Err(DataPlaneError::StreamClosed) + )); + assert_eq!(reader.read().await.expect("read eof"), None); + assert_eq!(reader.read().await.expect("sticky eof"), None); + }); +} + +#[test] +fn stream_writer_suspends_until_reader_releases_bounded_capacity() { + let _stream_test = STREAM_TEST_LOCK.lock(); + let harness = harness(2 << 20); + let data_plane = harness.bootstrap.data_plane.clone(); + let logical = path("/runs/self/results/backpressure"); + + future::block_on(async { + let mut reader_open = Box::pin(data_plane.read_stream(&logical)); + assert!(future::poll_once(reader_open.as_mut()).await.is_none()); + let mut writer = data_plane + .write_stream(&logical) + .await + .expect("writer opens"); + let mut reader = reader_open.await.expect("reader opens"); + + let capacity = writer.capacity() as usize; + let payload: Vec = (0..(capacity * 2 + 97)) + .map(|index| (index % 251) as u8) + .collect(); + let mut writing = Box::pin(writer.write(&payload)); + assert!( + future::poll_once(writing.as_mut()).await.is_none(), + "bounded source and destination rings must eventually suspend the writer" + ); + + let mut observed = reader + .read() + .await + .expect("read releases destination capacity") + .expect("first data"); + writing.await.expect("writer resumes"); + while observed.len() < payload.len() { + observed.extend( + reader + .read() + .await + .expect("read remaining") + .expect("remaining data"), + ); + } + assert_eq!(observed, payload); + writer.close().await.expect("close"); + assert_eq!(reader.read().await.expect("eof"), None); + }); +} +#[test] +fn transport_startup_failure_faults_both_pending_opens() { + let _stream_test = STREAM_TEST_LOCK.lock(); + let harness = harness_with_transport(2 << 20, Arc::new(RejectingStreamTransport)); + let data_plane = harness.bootstrap.data_plane.clone(); + let logical = path("/runs/self/results/faulted"); + + future::block_on(async { + let mut reader_open = Box::pin(data_plane.read_stream(&logical)); + assert!(future::poll_once(reader_open.as_mut()).await.is_none()); + let writer_error = match data_plane.write_stream(&logical).await { + Ok(_) => panic!("writer must not open when transport setup fails"), + Err(error) => error, + }; + let reader_error = match reader_open.await { + Ok(_) => panic!("reader must not open when transport setup fails"), + Err(error) => error, + }; + assert!(matches!( + writer_error, + DataPlaneError::PeerLost | DataPlaneError::StreamFault(_) + )); + assert!(matches!( + reader_error, + DataPlaneError::PeerLost | DataPlaneError::StreamFault(_) + )); + }); +} + +#[test] +fn peer_replacement_requires_and_supports_a_fresh_incarnation() { + let _stream_test = STREAM_TEST_LOCK.lock(); + let harness = harness(2 << 20); + let data_plane = harness.bootstrap.data_plane.clone(); + let logical = path("/runs/self/results/failover"); + + future::block_on(async { + let mut first_reader_open = Box::pin(data_plane.read_stream(&logical)); + assert!( + future::poll_once(first_reader_open.as_mut()) + .await + .is_none() + ); + let mut first_writer = data_plane + .write_stream(&logical) + .await + .expect("first writer"); + let mut first_reader = first_reader_open.await.expect("first reader"); + first_writer.write(b"old").await.expect("old write"); + assert_eq!( + first_reader.read().await.expect("old read"), + Some(b"old".to_vec()) + ); + first_writer.abort().expect("abort first incarnation"); + assert!(matches!( + first_reader.read().await, + Err(DataPlaneError::PeerLost) + )); + + let mut replacement_writer_open = Box::pin(data_plane.write_stream_replacing(&logical)); + assert!( + future::poll_once(replacement_writer_open.as_mut()) + .await + .is_none(), + "replacement writer waits for an explicit new reader" + ); + let mut replacement_reader = data_plane + .read_stream(&logical) + .await + .expect("replacement reader"); + let mut replacement_writer = replacement_writer_open.await.expect("replacement writer"); + replacement_writer.write(b"new").await.expect("new write"); + assert_eq!( + replacement_reader.read().await.expect("new read"), + Some(b"new".to_vec()) + ); + replacement_writer.close().await.expect("new close"); + assert_eq!(replacement_reader.read().await.expect("new eof"), None); + }); +} + +#[test] +fn actor_stream_consumer_registers_before_writer_and_collects_to_eof() { + let _stream_test = STREAM_TEST_LOCK.lock(); + let harness = harness(2 << 20); + let data_plane = harness.bootstrap.data_plane.clone(); + let logical = path("/runs/self/results/collector"); + let observed = Arc::new(parking_lot::Mutex::new(Vec::new())); + let consumer: Arc = + Arc::new(CollectBytes(Arc::clone(&observed))); + let completion = data_plane + .collect_stream(logical.clone(), consumer) + .expect("spawn collector"); + + future::block_on(async { + let mut writer = data_plane + .write_stream(&logical) + .await + .expect("writer matches collector"); + writer.write(b"actor-").await.expect("first write"); + writer.write(b"consumer").await.expect("second write"); + writer.close().await.expect("close"); + }); + + completion.wait().expect("collector completes"); + assert_eq!(&*observed.lock(), b"actor-consumer"); +} diff --git a/crates/data-plane/tests/byte_ring_guarantees.rs b/crates/data-plane/tests/byte_ring_guarantees.rs index 8809b7f..47f8b35 100644 --- a/crates/data-plane/tests/byte_ring_guarantees.rs +++ b/crates/data-plane/tests/byte_ring_guarantees.rs @@ -470,3 +470,90 @@ fn operations_revalidate_cursors_and_never_panic() { }) ); } + +#[test] +fn pinned_record_blocks_capacity_until_release() { + let (arena, handle) = installed(32, 1); + let mut producer = attach(&arena, handle, Role::Producer).expect("producer"); + let mut consumer = attach(&arena, handle, Role::Consumer).expect("consumer"); + + producer + .send_record(RecordKind::Data, &[7; 20]) + .expect("send"); + let view = consumer + .peek_record() + .expect("peek") + .expect("record must be visible"); + assert_eq!(view.kind(), RecordKind::Data); + assert_eq!(view.len(), 20); + assert!(view.spans().1.is_empty()); + assert_eq!( + producer.reserve(8).unwrap_err(), + FlowError::InsufficientSpace { + requested: 8, + free: 7, + } + ); + + drop(view); + assert_eq!(producer.reserve(8).expect("capacity released").len, 8); +} + +#[test] +fn pinned_record_exposes_wrapped_payload_as_two_spans() { + let (arena, handle) = installed(32, 1); + let mut producer = attach(&arena, handle, Role::Producer).expect("producer"); + let mut consumer = attach(&arena, handle, Role::Consumer).expect("consumer"); + + producer + .send_record(RecordKind::Data, &[1; 18]) + .expect("first"); + assert_eq!( + consumer.recv_record().expect("consume first"), + Some((RecordKind::Data, vec![1; 18])) + ); + let expected: Vec = (0..15).collect(); + producer + .send_record(RecordKind::Data, &expected) + .expect("wrapped record"); + + let view = consumer + .peek_record() + .expect("peek") + .expect("wrapped record visible"); + let (first, second) = view.spans(); + assert!(!first.is_empty()); + assert!(!second.is_empty()); + let observed: Vec = first.iter().chain(second).copied().collect(); + assert_eq!(observed, expected); +} + +#[test] +fn writable_record_is_invisible_until_commit() { + let (arena, handle) = installed(64, 1); + let mut producer = attach(&arena, handle, Role::Producer).expect("producer"); + let mut consumer = attach(&arena, handle, Role::Consumer).expect("consumer"); + + let mut reservation = producer + .reserve_record(RecordKind::Data, 17) + .expect("reserve record"); + let (first, second) = reservation.spans_mut(); + for (index, byte) in first.iter_mut().chain(second).enumerate() { + *byte = index as u8; + } + assert!(consumer.peek_record().expect("peek uncommitted").is_none()); + reservation.commit().expect("commit"); + + let view = consumer + .peek_record() + .expect("peek") + .expect("committed record"); + let observed: Vec = view + .spans() + .0 + .iter() + .chain(view.spans().1) + .copied() + .collect(); + assert_eq!(observed, (0..17).collect::>()); +} diff --git a/crates/data-plane/tests/namespace_guarantees.rs b/crates/data-plane/tests/namespace_guarantees.rs index 599c393..cad1cad 100644 --- a/crates/data-plane/tests/namespace_guarantees.rs +++ b/crates/data-plane/tests/namespace_guarantees.rs @@ -5,8 +5,8 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use data_plane::namespace::{ - DataDirectoryActor, DirectoryClient, NamespaceClient, NamespaceClientActor, NamespaceClientIn, - NamespaceDiscovery, NamespaceError, OperationId, SourceRecovery, + DataDirectoryActor, DirectoryClient, EntryKind, NamespaceClient, NamespaceClientActor, + NamespaceClientIn, NamespaceDiscovery, NamespaceError, OperationId, SourceRecovery, StreamRole, }; use data_plane::path::DataPath; use futures_lite::future; @@ -80,7 +80,11 @@ fn spawn_directory(store: &Path) -> DirectoryHarness { let directory = runtime.spawn(actor).expect("spawn directory actor"); let engine = Engine::new( parts, - TokioBackend::new(TokioConfig::default()).expect("tokio backend"), + TokioBackend::new(TokioConfig { + worker_threads: 1, + ..TokioConfig::default() + }) + .expect("tokio backend"), ) .expect("directory engine"); DirectoryHarness { @@ -168,6 +172,222 @@ fn namespace_mutations_are_linearizable_and_durable() { )); } +#[test] +fn stream_rendezvous_is_symmetric_and_incarnations_are_isolated() { + let state = TempState::new("stream-rendezvous"); + let directory = spawn_directory(&state.store()); + let logical = path("/runs/7/results"); + + future::block_on(async { + let mut source_open = Box::pin(directory.client.open_stream( + logical.clone(), + StreamRole::Source, + source(10), + OperationId::from_u128(10), + )); + assert!(future::poll_once(source_open.as_mut()).await.is_none()); + + let sink_match = directory + .client + .open_stream( + logical.clone(), + StreamRole::Sink, + source(11), + OperationId::from_u128(11), + ) + .await + .expect("sink matches source"); + let source_match = source_open.await.expect("source matches sink"); + assert_eq!(source_match, sink_match); + assert_eq!(source_match.source, source(10)); + assert_eq!(source_match.sink, source(11)); + + directory + .client + .close_stream(logical.clone(), source_match.incarnation) + .await + .expect("close first incarnation"); + + let mut sink_open = Box::pin(directory.client.open_stream( + logical.clone(), + StreamRole::Sink, + source(12), + OperationId::from_u128(12), + )); + assert!(future::poll_once(sink_open.as_mut()).await.is_none()); + let second_source = directory + .client + .open_stream( + logical, + StreamRole::Source, + source(13), + OperationId::from_u128(13), + ) + .await + .expect("source matches waiting sink"); + let second_sink = sink_open.await.expect("sink matches source"); + assert_eq!(second_source, second_sink); + assert_ne!(source_match.incarnation, second_source.incarnation); + }); +} + +#[test] +fn typed_paths_require_explicit_rebinding() { + let state = TempState::new("typed-path"); + let directory = spawn_directory(&state.store()); + let logical = path("/typed/value"); + let blob_source = source(20); + + future::block_on(async { + directory + .client + .register( + logical.clone(), + blob_source, + 4, + recovery(blob_source), + OperationId::from_u128(20), + ) + .await + .expect("register blob"); + + assert!(matches!( + directory + .client + .open_stream( + logical.clone(), + StreamRole::Source, + source(21), + OperationId::from_u128(21), + ) + .await, + Err(NamespaceError::WrongEntryType { + expected: EntryKind::Stream, + found: EntryKind::Blob, + .. + }) + )); + + let mut source_open = Box::pin(directory.client.replace_with_stream( + logical.clone(), + StreamRole::Source, + source(22), + OperationId::from_u128(22), + )); + assert!(future::poll_once(source_open.as_mut()).await.is_none()); + assert!(matches!( + directory.client.resolve(logical.clone()).await, + Err(NamespaceError::WrongEntryType { + expected: EntryKind::Blob, + found: EntryKind::Stream, + .. + }) + )); + let sink_match = directory + .client + .open_stream( + logical, + StreamRole::Sink, + source(23), + OperationId::from_u128(23), + ) + .await + .expect("match rebound stream"); + assert_eq!( + source_open.await.expect("rebound source matched"), + sink_match + ); + }); +} + +#[test] +fn replacing_waiting_stream_displaces_old_open() { + let state = TempState::new("stream-displacement"); + let directory = spawn_directory(&state.store()); + let logical = path("/replace/waiting"); + + future::block_on(async { + let mut old_open = Box::pin(directory.client.open_stream( + logical.clone(), + StreamRole::Source, + source(30), + OperationId::from_u128(30), + )); + assert!(future::poll_once(old_open.as_mut()).await.is_none()); + + let mut replacement = Box::pin(directory.client.replace_with_stream( + logical.clone(), + StreamRole::Source, + source(31), + OperationId::from_u128(31), + )); + assert!(future::poll_once(replacement.as_mut()).await.is_none()); + assert!(matches!( + old_open.await, + Err(NamespaceError::PathReplaced(found)) if found == logical + )); + + let sink_match = directory + .client + .open_stream( + logical, + StreamRole::Sink, + source(32), + OperationId::from_u128(32), + ) + .await + .expect("sink matches replacement"); + assert_eq!( + replacement.await.expect("replacement source matched"), + sink_match + ); + assert_eq!(sink_match.source, source(31)); + }); +} + +#[test] +fn duplicate_stream_role_fails_without_replacing_the_waiter() { + let state = TempState::new("duplicate-stream-role"); + let directory = spawn_directory(&state.store()); + let logical = path("/duplicate/source"); + + future::block_on(async { + let mut first = Box::pin(directory.client.open_stream( + logical.clone(), + StreamRole::Source, + source(51), + OperationId::from_u128(51), + )); + assert!(future::poll_once(first.as_mut()).await.is_none()); + assert!(matches!( + directory + .client + .open_stream( + logical.clone(), + StreamRole::Source, + source(52), + OperationId::from_u128(52), + ) + .await, + Err(NamespaceError::DuplicateStreamRole { + role: StreamRole::Source, + .. + }) + )); + let matched = directory + .client + .open_stream( + logical, + StreamRole::Sink, + source(53), + OperationId::from_u128(53), + ) + .await + .expect("sink matches original source"); + assert_eq!(matched.source, source(51)); + assert_eq!(first.await.expect("original source survives"), matched); + }); +} #[test] fn committed_mutation_retry_has_at_most_once_effect() { let state = TempState::new("idempotent"); @@ -369,9 +589,15 @@ struct ModelBinding { revision: u64, } +#[derive(Clone, Debug)] +enum TypedModelEntry { + Blob(ModelBinding), + Stream(data_plane::namespace::StreamMatch), +} + proptest! { #![proptest_config(ProptestConfig { - cases: 16, + cases: 8, max_shrink_iters: 128, ..ProptestConfig::default() })] @@ -434,9 +660,98 @@ proptest! { let observed = future::block_on(directory.client.resolve(path.clone())) .expect("all model bindings remain resolvable"); prop_assert_eq!(observed.source, expected.source); + prop_assert_eq!(observed.length, expected.length); prop_assert_eq!(observed.revision, expected.revision); } } } + #[test] + fn typed_binding_action_strings_match_reference_model(actions in prop::collection::vec(any::(), 1..64)) { + let state = TempState::new("typed-stateful"); + let directory = spawn_directory(&state.store()); + let paths = [path("/state/a"), path("/state/b"), path("/state/c")]; + let mut model = BTreeMap::::new(); + let mut next_operation = 10_000_u128; + + for (step, action) in actions.into_iter().enumerate() { + let logical = paths[usize::from(action) % paths.len()].clone(); + match action % 4 { + 0 => { + let actor = source(action.wrapping_add(step as u8).wrapping_add(1)); + let length = u64::from(action) + 1; + let receipt = future::block_on(directory.client.register( + logical.clone(), + actor, + length, + recovery(actor), + OperationId::from_u128(next_operation), + )).expect("blob rebind"); + next_operation += 1; + model.insert(logical, TypedModelEntry::Blob(ModelBinding { + source: actor, + length, + revision: receipt.revision, + })); + } + 1 => { + let source_actor = source(action.wrapping_add(41)); + let sink_actor = source(action.wrapping_add(97)); + let mut source_open = Box::pin(directory.client.replace_with_stream( + logical.clone(), + StreamRole::Source, + source_actor, + OperationId::from_u128(next_operation), + )); + next_operation += 1; + prop_assert!(future::block_on(future::poll_once(source_open.as_mut())).is_none()); + let sink_match = future::block_on(directory.client.open_stream( + logical.clone(), + StreamRole::Sink, + sink_actor, + OperationId::from_u128(next_operation), + )).expect("sink match"); + next_operation += 1; + let source_match = future::block_on(source_open).expect("source match"); + prop_assert_eq!(&source_match, &sink_match); + model.insert(logical, TypedModelEntry::Stream(sink_match)); + } + 2 => { + if let Some(TypedModelEntry::Stream(binding)) = model.get(&logical) { + future::block_on(directory.client.close_stream( + logical.clone(), + binding.incarnation, + )).expect("close current stream"); + model.remove(&logical); + } + } + _ => { + let observed = future::block_on(directory.client.resolve(logical.clone())); + match model.get(&logical) { + Some(TypedModelEntry::Blob(expected)) => { + let observed = observed.expect("blob resolves"); + prop_assert_eq!(observed.source, expected.source); + prop_assert_eq!(observed.length, expected.length); + prop_assert_eq!(observed.revision, expected.revision); + } + Some(TypedModelEntry::Stream(_)) => { + let wrong_type = matches!( + observed, + Err(NamespaceError::WrongEntryType { + expected: EntryKind::Blob, + found: EntryKind::Stream, + .. + }) + ); + prop_assert!(wrong_type, "blob lookup must reject a stream binding"); + } + None => prop_assert!(matches!( + observed, + Err(NamespaceError::PathNotFound(found)) if found == logical + )), + } + } + } + } + } } diff --git a/crates/data-plane/tests/namespace_host_read_guarantees.rs b/crates/data-plane/tests/namespace_host_read_guarantees.rs index 918f693..b7ad500 100644 --- a/crates/data-plane/tests/namespace_host_read_guarantees.rs +++ b/crates/data-plane/tests/namespace_host_read_guarantees.rs @@ -202,12 +202,20 @@ fn host_read_resolves_file_source_and_seals_final_arena_lease() { })); let host_engine = Engine::new( host_parts, - TokioBackend::new(TokioConfig::default()).unwrap(), + TokioBackend::new(TokioConfig { + worker_threads: 1, + ..TokioConfig::default() + }) + .unwrap(), ) .unwrap(); let child_engine = Engine::new( child_parts, - TokioBackend::new(TokioConfig::default()).unwrap(), + TokioBackend::new(TokioConfig { + worker_threads: 1, + ..TokioConfig::default() + }) + .unwrap(), ) .unwrap(); @@ -264,6 +272,7 @@ fn host_read_resolves_file_source_and_seals_final_arena_lease() { source_sender: Some(Arc::clone(&sender)), source_publisher: Some(Arc::clone(&source_publisher)), route_registrar: None, + stream_transport: None, }) .unwrap(), ) @@ -296,6 +305,7 @@ fn host_read_resolves_file_source_and_seals_final_arena_lease() { source_sender: Some(Arc::clone(&sender)), source_publisher: Some(source_publisher), route_registrar: None, + stream_transport: None, }) .unwrap(), ) diff --git a/crates/data-plane/tests/stream_transport_guarantees.rs b/crates/data-plane/tests/stream_transport_guarantees.rs new file mode 100755 index 0000000..c6c93d5 --- /dev/null +++ b/crates/data-plane/tests/stream_transport_guarantees.rs @@ -0,0 +1,260 @@ +#![cfg(target_os = "linux")] + +use std::sync::Arc; + +use data_plane::arena::{ArenaConfig, ArenaManager, NodeId}; +use data_plane::byte_ring::{ByteRingSpec, RecordKind, Role, attach, install}; +use data_plane::namespace::StreamIncarnation; +use data_plane::stream_transport::{ + LocalStreamTransport, StreamSinkRequest, StreamSourceRequest, StreamTransport, + StreamTransportEvent, StreamTransportNotifier, +}; +use parking_lot::Mutex; +use proptest::prelude::*; + +#[derive(Default)] +struct Events(Mutex>); + +impl StreamTransportNotifier for Events { + fn notify(&self, event: StreamTransportEvent) { + self.0.lock().push(event); + } +} + +fn arena(node: u64) -> ArenaManager { + ArenaManager::boot(ArenaConfig { + node_id: NodeId(node), + reservation_ceiling: 1 << 20, + base_alignment: 64, + }) + .expect("arena") +} + +#[test] +fn local_transport_preserves_order_backpressure_and_eof() { + let mut source_arena = arena(1); + let source_handle = install( + &mut source_arena, + ByteRingSpec { + capacity: 32, + generation: 1, + alignment: 64, + request_id: 1, + }, + ) + .expect("source ring"); + let mut destination_arena = arena(2); + let destination_handle = install( + &mut destination_arena, + ByteRingSpec { + capacity: 16, + generation: 2, + alignment: 64, + request_id: 2, + }, + ) + .expect("destination ring"); + + let mut writer = attach(&source_arena, source_handle, Role::Producer).expect("writer"); + let source_pump = attach(&source_arena, source_handle, Role::Consumer).expect("source pump"); + let destination_pump = + attach(&destination_arena, destination_handle, Role::Producer).expect("destination pump"); + let mut reader = + attach(&destination_arena, destination_handle, Role::Consumer).expect("reader"); + + let transport = LocalStreamTransport::new(); + let incarnation = StreamIncarnation { + authority_epoch: 7, + revision: 11, + }; + let source_events = Arc::new(Events::default()); + let sink_events = Arc::new(Events::default()); + transport + .install_sink(StreamSinkRequest { + incarnation, + endpoint: destination_pump, + notifier: sink_events.clone(), + }) + .expect("install sink"); + transport + .install_source(StreamSourceRequest { + incarnation, + peer: transport.descriptor().expect("descriptor"), + endpoint: source_pump, + notifier: source_events.clone(), + }) + .expect("install source"); + assert!( + source_events + .0 + .lock() + .contains(&StreamTransportEvent::Ready) + ); + assert!(sink_events.0.lock().contains(&StreamTransportEvent::Ready)); + + writer + .send_record(RecordKind::Data, b"first") + .expect("write first"); + transport.source_progress(incarnation); + writer + .send_record(RecordKind::Data, b"next!") + .expect("write second"); + transport.source_progress(incarnation); + + assert_eq!( + reader.recv_record().expect("read first"), + Some((RecordKind::Data, b"first".to_vec())) + ); + assert_eq!(reader.recv_record().expect("second remains upstream"), None); + transport.sink_progress(incarnation); + assert_eq!( + reader.recv_record().expect("read second"), + Some((RecordKind::Data, b"next!".to_vec())) + ); + + writer.send_record(RecordKind::Eof, b"").expect("write eof"); + transport.source_progress(incarnation); + assert_eq!( + reader.recv_record().expect("read eof"), + Some((RecordKind::Eof, Vec::new())) + ); +} + +#[test] +fn local_transport_rejects_a_descriptor_from_another_backend() { + let first = LocalStreamTransport::new(); + let second = LocalStreamTransport::new(); + assert_ne!( + first.descriptor().expect("first descriptor"), + second.descriptor().expect("second descriptor") + ); +} + +proptest! { + #![proptest_config(ProptestConfig { + cases: 32, + max_shrink_iters: 256, + ..ProptestConfig::default() + })] + + #[test] + fn randomized_payloads_preserve_exact_bytes( + payload in prop::collection::vec(any::(), 0..4096), + raw_capacity in 16_u8..128, + chunk_seeds in prop::collection::vec(1_u8..=255, 1..32), + ) { + let capacity = u64::from(raw_capacity); + let mut source_arena = arena(21); + let source_handle = install( + &mut source_arena, + ByteRingSpec { + capacity, + generation: 1, + alignment: 64, + request_id: 1, + }, + ).expect("source ring"); + let mut destination_arena = arena(22); + let destination_handle = install( + &mut destination_arena, + ByteRingSpec { + capacity, + generation: 2, + alignment: 64, + request_id: 2, + }, + ).expect("destination ring"); + let mut writer = attach(&source_arena, source_handle, Role::Producer).expect("writer"); + let source_pump = attach(&source_arena, source_handle, Role::Consumer).expect("source pump"); + let destination_pump = + attach(&destination_arena, destination_handle, Role::Producer).expect("destination pump"); + let mut reader = + attach(&destination_arena, destination_handle, Role::Consumer).expect("reader"); + let transport = LocalStreamTransport::new(); + let incarnation = StreamIncarnation { + authority_epoch: 3, + revision: 9, + }; + transport.install_sink(StreamSinkRequest { + incarnation, + endpoint: destination_pump, + notifier: Arc::new(Events::default()), + }).expect("sink"); + transport.install_source(StreamSourceRequest { + incarnation, + peer: transport.descriptor().expect("descriptor"), + endpoint: source_pump, + notifier: Arc::new(Events::default()), + }).expect("source"); + + let max_chunk = capacity as usize - 5; + let mut offset = 0; + let mut seed_index = 0; + let mut observed = Vec::new(); + while offset < payload.len() { + let chunk_len = usize::from(chunk_seeds[seed_index % chunk_seeds.len()]) + .min(max_chunk) + .min(payload.len() - offset); + seed_index += 1; + loop { + match writer.send_record( + RecordKind::Data, + &payload[offset..offset + chunk_len], + ) { + Ok(()) => break, + Err(data_plane::byte_ring::FlowError::InsufficientSpace { .. }) => { + if let Some((RecordKind::Data, bytes)) = + reader.recv_record().expect("drain destination") + { + observed.extend(bytes); + transport.sink_progress(incarnation); + } else { + transport.source_progress(incarnation); + } + } + Err(error) => panic!("unexpected source error: {error:?}"), + } + } + offset += chunk_len; + transport.source_progress(incarnation); + } + + loop { + match writer.send_record(RecordKind::Eof, &[]) { + Ok(()) => break, + Err(data_plane::byte_ring::FlowError::InsufficientSpace { .. }) => { + if let Some((RecordKind::Data, bytes)) = + reader.recv_record().expect("drain for eof") + { + observed.extend(bytes); + transport.sink_progress(incarnation); + } else { + transport.source_progress(incarnation); + } + } + Err(error) => panic!("unexpected eof error: {error:?}"), + } + } + transport.source_progress(incarnation); + + let mut steps = 0; + loop { + steps += 1; + prop_assert!(steps < payload.len() + 1024, "transfer made no bounded progress"); + match reader.recv_record().expect("receive") { + Some((RecordKind::Data, bytes)) => { + observed.extend(bytes); + transport.sink_progress(incarnation); + } + Some((RecordKind::Eof, _)) => break, + Some((RecordKind::Fault, bytes)) => { + return Err(TestCaseError::fail(format!( + "unexpected fault record: {bytes:?}" + ))); + } + None => transport.source_progress(incarnation), + } + } + prop_assert_eq!(observed, payload); + } +} diff --git a/crates/iroh-driver/src/iroh_driver.rs b/crates/iroh-driver/src/iroh_driver.rs index 5e44c95..ff77306 100644 --- a/crates/iroh-driver/src/iroh_driver.rs +++ b/crates/iroh-driver/src/iroh_driver.rs @@ -33,6 +33,7 @@ use distribution::types::NodeId; use crate::edge_transport::spawn_edge_send_pump as spawn_edge_sender_task; use crate::edge_transport::{EDGE_ALPN, EdgeSendHandle, spawn_edge_recv_pump}; +use crate::stream_transport::{IrohStreamTransport, STREAM_ALPN}; use crate::telemetry_transport::{ TELEMETRY_ALPN, TelemetryQuicHeader, TelemetryQuicRead, read_events_from_stream, spawn_subscription_writer, @@ -327,6 +328,7 @@ pub struct IrohDriver { /// iroh work is scheduled through this handle; it never exposes the raw /// Tokio runtime (ENGINE_SPEC.md §7). engine: EngineHandle, + stream_transport: Arc, conns: Arc>, peer_auth: Option>>, /// Collects connections from background join tasks. @@ -450,7 +452,7 @@ impl IrohDriver { let secret_key = config.secret_key; let (endpoint_tx, endpoint_rx) = std::sync::mpsc::channel::>(); engine.spawn(async move { - let mut all_alpns = vec![ALPN.to_vec()]; + let mut all_alpns = vec![ALPN.to_vec(), STREAM_ALPN.to_vec()]; all_alpns.extend(additional_alpns); let mut builder = Endpoint::builder(iroh::endpoint::presets::Minimal) .relay_mode(effective_relay_mode) @@ -481,6 +483,7 @@ impl IrohDriver { format!("engine endpoint-bind task dropped: {e}").into() })? .map_err(|e| -> Box { e.into() })?; + let stream_transport = Arc::new(IrohStreamTransport::new(engine.clone(), endpoint.clone())); let relay_url = endpoint .addr() .relay_urls() @@ -504,6 +507,7 @@ impl IrohDriver { let peer_auth = config.peer_auth.clone(); let swim_buf = Arc::clone(&accepted_conns); let other_buf = Arc::clone(&other_accepted_conns); + let accepted_streams = Arc::clone(&stream_transport); engine.spawn(async move { while let Some(incoming) = ep.accept().await { if let Ok(conn) = incoming.await { @@ -522,6 +526,8 @@ impl IrohDriver { let negotiated_alpn = conn.alpn().to_vec(); if negotiated_alpn == ALPN { swim_buf.lock().push((node_id, conn)); + } else if negotiated_alpn == STREAM_ALPN { + accepted_streams.accept_connection(conn); } else { other_buf.lock().push((node_id, negotiated_alpn, conn)); } @@ -534,6 +540,7 @@ impl IrohDriver { keypair, endpoint, engine, + stream_transport, conns: Arc::new(Mutex::new(ConnCache { connections: HashMap::new(), next_generation: 1, @@ -561,6 +568,10 @@ impl IrohDriver { self.endpoint.clone() } + pub fn stream_transport(&self) -> Arc { + Arc::clone(&self.stream_transport) + } + /// Drain accepted connections whose negotiated ALPN exactly matches `alpn`. pub fn drain_accepted_for_alpn(&self, alpn: &[u8]) -> Vec<(NodeId, Connection)> { let mut pending = self.other_accepted_conns.lock(); diff --git a/crates/iroh-driver/src/lib.rs b/crates/iroh-driver/src/lib.rs index 1ba34fd..1bc0176 100644 --- a/crates/iroh-driver/src/lib.rs +++ b/crates/iroh-driver/src/lib.rs @@ -13,6 +13,7 @@ pub mod blob_transfer; pub mod edge_transport; pub mod endpoint_advertisement; pub mod iroh_driver; +pub mod stream_transport; pub mod telemetry_transport; pub use blob_transfer::{IrohBlobTransferReceiver, IrohBlobTransferSender}; @@ -25,6 +26,7 @@ pub use iroh_driver::{ }; pub use edge_transport::{EDGE_ALPN, EdgeSendHandle}; +pub use stream_transport::{IrohStreamTransport, STREAM_ALPN}; pub use telemetry_transport::{ PullCollectorConfig, PullCollectorHandle, TELEMETRY_ALPN, TelemetryQuicHeader, diff --git a/crates/iroh-driver/src/stream_transport.rs b/crates/iroh-driver/src/stream_transport.rs new file mode 100644 index 0000000..0b5ee9c --- /dev/null +++ b/crates/iroh-driver/src/stream_transport.rs @@ -0,0 +1,431 @@ +//! Iroh-specific implementation of the data-plane SPSC stream transport port. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use data_plane::byte_ring::{Endpoint as RingEndpoint, FlowError, RecordKind, RingProbe}; +use data_plane::namespace::StreamIncarnation; +use data_plane::stream_transport::{ + StreamPeerDescriptor, StreamSinkRequest, StreamSourceRequest, StreamTransport, + StreamTransportEvent, StreamTransportNotifier, +}; +use iroh::endpoint::{Connection, RecvStream, SendStream}; +use iroh::{Endpoint as IrohEndpoint, EndpointAddr}; +use parking_lot::Mutex; +use swactor_engine::EngineHandle; +use tokio::io::AsyncWriteExt; +use tokio::sync::mpsc; + +pub const STREAM_ALPN: &[u8] = b"swactor/data-plane-spsc/1"; +const PREAMBLE_LEN: usize = 16; +const RECORD_HEADER_LEN: usize = 5; + +#[derive(Clone)] +struct TaskControl { + wake: mpsc::Sender<()>, + progress_pending: Arc, + cancelled: Arc, +} + +impl TaskControl { + fn pair() -> (Self, mpsc::Receiver<()>) { + let (wake, receiver) = mpsc::channel(1); + ( + Self { + wake, + progress_pending: Arc::new(AtomicBool::new(false)), + cancelled: Arc::new(AtomicBool::new(false)), + }, + receiver, + ) + } + + fn progress(&self) { + if !self.progress_pending.swap(true, Ordering::AcqRel) { + let _ = self.wake.try_send(()); + } + } + + fn cancel(&self) { + self.cancelled.store(true, Ordering::Release); + self.progress(); + } + + async fn wait(&self, receiver: &mut mpsc::Receiver<()>) -> bool { + if self.cancelled.load(Ordering::Acquire) { + return false; + } + if receiver.recv().await.is_none() { + return false; + } + self.progress_pending.store(false, Ordering::Release); + !self.cancelled.load(Ordering::Acquire) + } +} + +struct PendingSink { + endpoint: RingEndpoint, + notifier: Arc, +} + +#[derive(Default)] +struct TransportState { + pending_sinks: BTreeMap, + controls: BTreeMap>, + source_probes: BTreeMap, + sink_probes: BTreeMap, +} + +struct Inner { + engine: EngineHandle, + endpoint: IrohEndpoint, + state: Mutex, +} + +/// Cloneable adapter capability installed into `HostDataPlaneConfig`. +#[derive(Clone)] +pub struct IrohStreamTransport { + inner: Arc, +} + +impl IrohStreamTransport { + pub fn new(engine: EngineHandle, endpoint: IrohEndpoint) -> Self { + Self { + inner: Arc::new(Inner { + engine, + endpoint, + state: Mutex::new(TransportState::default()), + }), + } + } + + pub(crate) fn accept_connection(&self, connection: Connection) { + let transport = self.clone(); + self.inner.engine.spawn(async move { + while let Ok(mut recv) = connection.accept_uni().await { + let mut preamble = [0_u8; PREAMBLE_LEN]; + if recv.read_exact(&mut preamble).await.is_err() { + continue; + } + let incarnation = decode_incarnation(preamble); + let pending = transport + .inner + .state + .lock() + .pending_sinks + .remove(&incarnation); + let Some(pending) = pending else { + continue; + }; + let (control, receiver) = TaskControl::pair(); + transport + .inner + .state + .lock() + .controls + .entry(incarnation) + .or_default() + .push(control.clone()); + transport.inner.engine.spawn(run_sink( + incarnation, + recv, + pending.endpoint, + pending.notifier, + control, + receiver, + )); + } + }); + } + + fn register_control(&self, incarnation: StreamIncarnation, control: TaskControl) { + self.inner + .state + .lock() + .controls + .entry(incarnation) + .or_default() + .push(control); + } + + fn progress_controls(&self, incarnation: StreamIncarnation) { + if let Some(controls) = self.inner.state.lock().controls.get(&incarnation) { + for control in controls { + control.progress(); + } + } + } +} + +impl StreamTransport for IrohStreamTransport { + fn descriptor(&self) -> Result { + serde_json::to_vec(&self.inner.endpoint.addr()) + .map(StreamPeerDescriptor) + .map_err(|error| format!("encode iroh stream endpoint: {error}")) + } + + fn install_source(&self, request: StreamSourceRequest) -> Result<(), String> { + let peer: EndpointAddr = serde_json::from_slice(&request.peer.0) + .map_err(|error| format!("decode iroh stream endpoint: {error}"))?; + let (control, receiver) = TaskControl::pair(); + self.register_control(request.incarnation, control.clone()); + let endpoint = self.inner.endpoint.clone(); + let probe = request.endpoint.probe(); + self.inner + .state + .lock() + .source_probes + .insert(request.incarnation, probe); + self.inner.engine.spawn(run_source( + request.incarnation, + endpoint, + peer, + request.endpoint, + request.notifier, + control, + receiver, + )); + Ok(()) + } + + fn install_sink(&self, request: StreamSinkRequest) -> Result<(), String> { + let mut state = self.inner.state.lock(); + if state.pending_sinks.contains_key(&request.incarnation) { + return Err("iroh stream sink is already installed".to_owned()); + } + let probe = request.endpoint.probe(); + state.sink_probes.insert(request.incarnation, probe); + state.pending_sinks.insert( + request.incarnation, + PendingSink { + endpoint: request.endpoint, + notifier: request.notifier, + }, + ); + Ok(()) + } + + fn source_progress(&self, incarnation: StreamIncarnation) { + self.progress_controls(incarnation); + } + + fn sink_progress(&self, incarnation: StreamIncarnation) { + self.progress_controls(incarnation); + } + + fn source_has_capacity(&self, incarnation: StreamIncarnation) -> bool { + self.inner + .state + .lock() + .source_probes + .get(&incarnation) + .is_some_and(RingProbe::has_capacity) + } + + fn sink_has_data(&self, incarnation: StreamIncarnation) -> bool { + self.inner + .state + .lock() + .sink_probes + .get(&incarnation) + .is_some_and(RingProbe::has_data) + } + + fn terminate(&self, incarnation: StreamIncarnation) { + let (pending, controls) = { + let mut state = self.inner.state.lock(); + let pending = state.pending_sinks.remove(&incarnation); + let controls = state.controls.remove(&incarnation).unwrap_or_default(); + state.source_probes.remove(&incarnation); + state.sink_probes.remove(&incarnation); + (pending, controls) + }; + if let Some(pending) = pending { + pending.notifier.notify(StreamTransportEvent::Quiesced); + } + for control in controls { + control.cancel(); + } + } +} + +async fn run_source( + incarnation: StreamIncarnation, + endpoint: IrohEndpoint, + peer: EndpointAddr, + mut source: RingEndpoint, + notifier: Arc, + control: TaskControl, + mut receiver: mpsc::Receiver<()>, +) { + let result = async { + let connection = endpoint + .connect(peer, STREAM_ALPN) + .await + .map_err(|error| format!("connect stream incarnation: {error}"))?; + let mut send = connection + .open_uni() + .await + .map_err(|error| format!("open stream incarnation: {error}"))?; + send.write_all(&encode_incarnation(incarnation)) + .await + .map_err(|error| format!("write stream preamble: {error}"))?; + send.flush() + .await + .map_err(|error| format!("flush stream preamble: {error}"))?; + notifier.notify(StreamTransportEvent::Ready); + + loop { + if control.cancelled.load(Ordering::Acquire) { + return Ok(()); + } + let mut moved = false; + while let Some(meta) = source + .next_record_meta() + .map_err(|error| format!("inspect source ring: {error:?}"))? + { + let view = source + .peek_record() + .map_err(|error| format!("pin source ring: {error:?}"))? + .ok_or_else(|| "source record disappeared after inspection".to_owned())?; + write_record(&mut send, meta.kind, view.spans()) + .await + .map_err(|error| format!("write stream record: {error}"))?; + view.release() + .map_err(|error| format!("release source ring: {error:?}"))?; + notifier.notify(StreamTransportEvent::CapacityAvailable); + moved = true; + if matches!(meta.kind, RecordKind::Eof | RecordKind::Fault) { + send.finish() + .map_err(|error| format!("finish stream incarnation: {error}"))?; + match send + .stopped() + .await + .map_err(|error| format!("await stream finish: {error}"))? + { + Some(code) => { + return Err(format!("peer stopped stream incarnation: {code}")); + } + None => return Ok(()), + } + } + } + if !moved && !control.wait(&mut receiver).await { + return Ok(()); + } + } + } + .await; + if let Err(reason) = result { + notifier.notify(StreamTransportEvent::Fault(reason)); + } + notifier.notify(StreamTransportEvent::Quiesced); +} + +async fn write_record( + send: &mut SendStream, + kind: RecordKind, + spans: (&[u8], &[u8]), +) -> Result<(), String> { + let len = spans.0.len() + spans.1.len(); + let len = u32::try_from(len).map_err(|_| "stream record exceeds u32 framing".to_owned())?; + let mut header = [0_u8; RECORD_HEADER_LEN]; + header[0] = kind.to_byte(); + header[1..].copy_from_slice(&len.to_le_bytes()); + send.write_all(&header) + .await + .map_err(|error| error.to_string())?; + if !spans.0.is_empty() { + send.write_all(spans.0) + .await + .map_err(|error| error.to_string())?; + } + if !spans.1.is_empty() { + send.write_all(spans.1) + .await + .map_err(|error| error.to_string())?; + } + send.flush().await.map_err(|error| error.to_string()) +} + +async fn run_sink( + _incarnation: StreamIncarnation, + mut recv: RecvStream, + mut sink: RingEndpoint, + notifier: Arc, + control: TaskControl, + mut receiver: mpsc::Receiver<()>, +) { + notifier.notify(StreamTransportEvent::Ready); + let result = async { + loop { + if control.cancelled.load(Ordering::Acquire) { + return Ok(()); + } + let mut header = [0_u8; RECORD_HEADER_LEN]; + recv.read_exact(&mut header) + .await + .map_err(|error| format!("read stream record header: {error}"))?; + let kind = RecordKind::from_byte(header[0]) + .ok_or_else(|| format!("invalid stream record kind {}", header[0]))?; + let len = u64::from(u32::from_le_bytes(header[1..].try_into().unwrap())); + let mut reservation = loop { + match sink.reserve_record(kind, len) { + Ok(reservation) => break reservation, + Err(FlowError::InsufficientSpace { .. }) => { + if !control.wait(&mut receiver).await { + return Ok(()); + } + } + Err(error) => return Err(format!("reserve sink ring: {error:?}")), + } + }; + let (first, second) = reservation.spans_mut(); + if !first.is_empty() { + recv.read_exact(first) + .await + .map_err(|error| format!("read first stream span: {error}"))?; + } + if !second.is_empty() { + recv.read_exact(second) + .await + .map_err(|error| format!("read second stream span: {error}"))?; + } + reservation + .commit() + .map_err(|error| format!("commit sink ring: {error:?}"))?; + notifier.notify(StreamTransportEvent::DataAvailable); + if matches!(kind, RecordKind::Eof | RecordKind::Fault) { + let mut trailing = [0_u8; 1]; + match recv + .read(&mut trailing) + .await + .map_err(|error| format!("read stream finish: {error}"))? + { + None | Some(0) => return Ok(()), + Some(_) => return Err("bytes followed terminal stream record".to_owned()), + } + } + } + } + .await; + if let Err(reason) = result { + notifier.notify(StreamTransportEvent::Fault(reason)); + } + notifier.notify(StreamTransportEvent::Quiesced); +} + +fn encode_incarnation(incarnation: StreamIncarnation) -> [u8; PREAMBLE_LEN] { + let mut encoded = [0_u8; PREAMBLE_LEN]; + encoded[..8].copy_from_slice(&incarnation.authority_epoch.to_le_bytes()); + encoded[8..].copy_from_slice(&incarnation.revision.to_le_bytes()); + encoded +} + +fn decode_incarnation(encoded: [u8; PREAMBLE_LEN]) -> StreamIncarnation { + StreamIncarnation { + authority_epoch: u64::from_le_bytes(encoded[..8].try_into().unwrap()), + revision: u64::from_le_bytes(encoded[8..].try_into().unwrap()), + } +} diff --git a/crates/iroh-driver/tests/stream_transport.rs b/crates/iroh-driver/tests/stream_transport.rs new file mode 100644 index 0000000..9487d51 --- /dev/null +++ b/crates/iroh-driver/tests/stream_transport.rs @@ -0,0 +1,137 @@ +pub mod common; + +use std::sync::Arc; +use std::sync::mpsc::{self, Receiver, Sender}; +use std::time::Duration; + +use data_plane::arena::{ArenaConfig, ArenaManager, NodeId}; +use data_plane::byte_ring::{ByteRingSpec, RecordKind, Role, attach, install}; +use data_plane::namespace::StreamIncarnation; +use data_plane::stream_transport::{ + StreamSinkRequest, StreamSourceRequest, StreamTransport, StreamTransportEvent, + StreamTransportNotifier, +}; + +use common::iroh::make_driver; + +struct ChannelNotifier(Sender); + +impl StreamTransportNotifier for ChannelNotifier { + fn notify(&self, event: StreamTransportEvent) { + let _ = self.0.send(event); + } +} + +fn arena(node: u64) -> ArenaManager { + ArenaManager::boot(ArenaConfig { + node_id: NodeId(node), + reservation_ceiling: 1 << 20, + base_alignment: 64, + }) + .expect("arena") +} + +fn recv_until(receiver: &Receiver, expected: StreamTransportEvent) { + loop { + let event = receiver + .recv_timeout(Duration::from_secs(10)) + .expect("transport event deadline"); + if event == expected { + return; + } + if let StreamTransportEvent::Fault(reason) = event { + panic!("unexpected stream fault: {reason}"); + } + } +} + +#[test] +fn iroh_adapter_satisfies_ordering_and_terminal_contract() { + let source_node = make_driver(); + let sink_node = make_driver(); + let source_transport = source_node.driver.stream_transport(); + let sink_transport = sink_node.driver.stream_transport(); + + let mut source_arena = arena(1); + let source_handle = install( + &mut source_arena, + ByteRingSpec { + capacity: 128, + generation: 1, + alignment: 64, + request_id: 1, + }, + ) + .expect("source ring"); + let mut sink_arena = arena(2); + let sink_handle = install( + &mut sink_arena, + ByteRingSpec { + capacity: 128, + generation: 2, + alignment: 64, + request_id: 2, + }, + ) + .expect("sink ring"); + let mut writer = attach(&source_arena, source_handle, Role::Producer).expect("writer"); + let source_pump = attach(&source_arena, source_handle, Role::Consumer).expect("source pump"); + let sink_pump = attach(&sink_arena, sink_handle, Role::Producer).expect("sink pump"); + let mut reader = attach(&sink_arena, sink_handle, Role::Consumer).expect("reader"); + + let incarnation = StreamIncarnation { + authority_epoch: 17, + revision: 23, + }; + let (source_tx, source_rx) = mpsc::channel(); + let (sink_tx, sink_rx) = mpsc::channel(); + sink_transport + .install_sink(StreamSinkRequest { + incarnation, + endpoint: sink_pump, + notifier: Arc::new(ChannelNotifier(sink_tx)), + }) + .expect("install sink"); + source_transport + .install_source(StreamSourceRequest { + incarnation, + peer: sink_transport.descriptor().expect("sink descriptor"), + endpoint: source_pump, + notifier: Arc::new(ChannelNotifier(source_tx)), + }) + .expect("install source"); + recv_until(&source_rx, StreamTransportEvent::Ready); + recv_until(&sink_rx, StreamTransportEvent::Ready); + + writer + .send_record(RecordKind::Data, b"one") + .expect("write one"); + writer + .send_record(RecordKind::Data, b"two") + .expect("write two"); + source_transport.source_progress(incarnation); + recv_until(&sink_rx, StreamTransportEvent::DataAvailable); + recv_until(&sink_rx, StreamTransportEvent::DataAvailable); + assert_eq!( + reader.recv_record().expect("read one"), + Some((RecordKind::Data, b"one".to_vec())) + ); + assert_eq!( + reader.recv_record().expect("read two"), + Some((RecordKind::Data, b"two".to_vec())) + ); + sink_transport.sink_progress(incarnation); + + writer.send_record(RecordKind::Eof, &[]).expect("write eof"); + source_transport.source_progress(incarnation); + recv_until(&sink_rx, StreamTransportEvent::DataAvailable); + assert_eq!( + reader.recv_record().expect("read eof"), + Some((RecordKind::Eof, Vec::new())) + ); + recv_until(&source_rx, StreamTransportEvent::Quiesced); + recv_until(&sink_rx, StreamTransportEvent::Quiesced); + + source_transport.terminate(incarnation); + sink_transport.terminate(incarnation); +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs index 650cbba..34b1a24 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -37,6 +37,41 @@ fn cargo_bin() -> String { .unwrap_or_else(|| "cargo".to_string()) } +fn remove_inherited_build_context(command: &mut Command) { + for (name, _) in std::env::vars_os() { + let name_text = name.to_string_lossy(); + let package_scoped = name_text.starts_with("CARGO_PKG_") + || name_text.starts_with("CARGO_FEATURE_") + || name_text.starts_with("CARGO_CFG_") + || name_text.starts_with("DEP_"); + let build_scoped = matches!( + name_text.as_ref(), + "CARGO_BIN_NAME" + | "CARGO_CRATE_NAME" + | "CARGO_MANIFEST_DIR" + | "CARGO_MANIFEST_PATH" + | "CARGO_PRIMARY_PACKAGE" + | "DEBUG" + | "HOST" + | "NUM_JOBS" + | "OPT_LEVEL" + | "OUT_DIR" + | "PROFILE" + | "PYO3_ENVIRONMENT_SIGNATURE" + | "TARGET" + ); + if package_scoped || build_scoped { + command.env_remove(name); + } + } +} + +fn cargo_command() -> Command { + let mut command = Command::new(cargo_bin()); + remove_inherited_build_context(&mut command); + command +} + fn print_usage() { println!( "\ @@ -57,9 +92,7 @@ fn run_step(step: &TestStep, python: &Path) -> bool { println!(); match swactor_process::command_status( - Command::new(cargo_bin()) - .args(step.args) - .env("PYO3_PYTHON", python), + cargo_command().args(step.args).env("PYO3_PYTHON", python), ) { Ok(status) => status.success(), Err(error) => { @@ -107,7 +140,7 @@ fn run_tests() -> ExitCode { fn nextest_available() -> bool { let available = swactor_process::command_status( - Command::new(cargo_bin()) + cargo_command() .args(["nextest", "--version"]) .stdout(Stdio::null()) .stderr(Stdio::null()), @@ -146,8 +179,10 @@ impl PythonTestTools { fn run(&self) -> bool { println!("\n=== all Python tests ==="); + let mut build_command = Command::new(&self.maturin); + remove_inherited_build_context(&mut build_command); let built = swactor_process::command_status( - Command::new(&self.maturin) + build_command .current_dir(&self.directory) .arg("develop") .env("PYO3_PYTHON", &self.python),