feat(data-plane): add routed SPSC stream transport

Implement namespace-owned stream rendezvous with incarnation isolation, cancellation, recovery, bounded local byte rings, and host/child actor protocols.

Add a replaceable stream transport boundary plus an Iroh ALPN adapter with framed records, backpressure, reconnect handling, and terminal propagation.

Wire stream sources and sinks through Myelin job deployment and Python bindings, with guarantee coverage across blob, namespace, job, byte-ring, and transport paths.

Sanitize nested Cargo build context in xtask so lint, Nextest, doctest, and maturin invocations retain stable fingerprints.
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-08-23 12:47:31 +04:00
parent df72b4ff60
commit 0411c20b94
27 changed files with 5299 additions and 252 deletions

View file

@ -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<Arc<dyn BlobTransferSender>>,
pub(crate) source_publisher: Option<Arc<dyn BlobSourcePublisher>>,
pub(crate) route_registrar: Option<Arc<dyn HostRouteRegistrar>>,
pub(crate) stream_transport: Option<Arc<dyn StreamTransport>>,
}
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<DataPlane, String> {
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

View file

@ -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<Mutex<Option<EdgeSendHandle>>>,
result_ready: Arc<Notify>,
connector: EdgeConnector,
actor_plane: Arc<ActorJobDataPlane>,
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<dyn BlobTransferReceiver>,
pub(crate) source_sender: Arc<dyn BlobTransferSender>,
pub(crate) source_publisher: Arc<dyn BlobSourcePublisher>,
pub(crate) stream_transport: Arc<dyn data_plane::stream_transport::StreamTransport>,
}
impl EmbeddedJobDataPlane {
pub(crate) fn start(
stack: &DistributionRuntimeStack,
config: EmbeddedJobDataPlaneConfig<'_>,
config: EmbeddedJobDataPlaneConfig,
) -> Result<Self, String> {
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<Mutex<Option<EdgeSendHandle>>> = 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<BTreeMap<String, String>, String> {
let assignment = serde_json::from_str::<EmbeddedDataPlaneAssignment>(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<Mutex<Vec<u8>>>,
}
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<ActorJobDataPlane>,
_result_completion: ActorCompletion<Result<(), DataPlaneError>>,
_result_bytes: Arc<Mutex<Vec<u8>>>,
}
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<dyn data_plane::stream_transport::StreamTransport> =
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<dyn StreamConsumer> = 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::<JobDone>()
@ -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,
})
}

View file

@ -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(

View file

@ -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()
})]

View file

@ -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"),
)

View file

@ -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<ActorAddress> {
Some(self.0)
}
}
struct LocalPublisher;
impl BlobSourcePublisher for LocalPublisher {
fn publish_source(&self, _source: ActorAddress) -> Result<(), String> {
Ok(())
}
}
struct BytesConsumer(Arc<Mutex<Vec<u8>>>);
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<dyn StreamTransport> = Arc::new(LocalStreamTransport::new());
let publisher: Arc<dyn BlobSourcePublisher> = 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<dyn StreamConsumer> = 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);
}

View file

@ -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<String>) -> PyErr {
PyErr::new::<BootstrapError, _>(message.into())
@ -56,6 +57,11 @@ fn data_plane_error(error: DataPlaneError) -> PyErr {
PyErr::new::<DataPathError, _>(format!("data path not found: {path}"))
}
DataPlaneError::Blob(reason) => PyErr::new::<BlobError, _>(format!("{reason:?}")),
DataPlaneError::WrongEntryType { .. }
| DataPlaneError::PathReplaced(_)
| DataPlaneError::PeerLost
| DataPlaneError::StreamFault(_)
| DataPlaneError::StreamClosed => PyErr::new::<StreamError, _>(error.to_string()),
DataPlaneError::Attachment(reason) => {
PyErr::new::<SessionError, _>(format!("attachment failed: {reason:?}"))
}
@ -229,7 +235,6 @@ fn build_child_routing(
pub struct PyDataPlane {
inner: Arc<DataPlane>,
_routing: Arc<JobRouting>,
legacy_output: Option<String>,
}
#[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<Bound<'py, PyAny>> {
let path = DataPath::parse(path)
.map_err(|error| PyErr::new::<DataPathError, _>(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<PyStreamContext> {
DataPath::parse(path).map_err(|error| PyErr::new::<DataPathError, _>(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::<DataPathError, _>(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<tokio::sync::Mutex<StreamReader>>,
}
#[pymethods]
impl PyStreamReader {
fn read<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
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<tokio::sync::Mutex<Option<BufWriter<UnixStream>>>>,
stream: Arc<tokio::sync::Mutex<Option<StreamWriter>>>,
}
#[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<tokio::sync::Mutex<Option<BufWriter<UnixStream>>>>,
data_plane: Arc<DataPlane>,
path: DataPath,
stream: Arc<tokio::sync::Mutex<Option<StreamWriter>>>,
entered: Arc<AtomicBool>,
}
#[pymethods]
impl PyStreamContext {
fn __aenter__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
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<Bound<'py, PyAny>> {
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<ParkingMutex<Vec<u8>>>,
}
#[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<DataPlane>,
}
#[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<Bound<'py, PyAny>> {
let path = DataPath::parse(path)
.map_err(|error| PyErr::new::<DataPathError, _>(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<PyTestStreamCapture> {
let path = DataPath::parse(path)
.map_err(|error| PyErr::new::<DataPathError, _>(error.to_string()))?;
let bytes = Arc::new(ParkingMutex::new(Vec::new()));
let consumer: Arc<dyn data_plane::data_plane::StreamConsumer> =
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<Result<(), DataPlaneError>>,
bytes: Arc<ParkingMutex<Vec<u8>>>,
}
#[cfg(all(debug_assertions, target_os = "linux"))]
#[pymethods]
impl PyTestStreamCapture {
fn result<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
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<PyTestDataPlaneHost> {
))
.map_err(|error| PyRuntimeError::new_err(error.to_string()))?;
let namespace = data_plane::namespace::NamespaceClient::new(runtime.clone(), namespace_proxy);
let stream_transport: Arc<dyn data_plane::stream_transport::StreamTransport> =
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<PyTestDataPlaneHost> {
.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<PyTestDataPlaneHost> {
.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<PyTestDataPlaneHost> {
_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::<DataPathError>())?;
module.add("BlobError", module.py().get_type::<BlobError>())?;
module.add("SessionError", module.py().get_type::<SessionError>())?;
module.add("StreamError", module.py().get_type::<StreamError>())?;
module.add_class::<PyDataPlane>()?;
module.add_class::<PyBlob>()?;
module.add_class::<PyBlobView>()?;
module.add_class::<PyStreamWriter>()?;
module.add_class::<PyStreamReader>()?;
module.add_class::<PyContext>()?;
module.add_function(wrap_pyfunction!(run, module)?)?;
#[cfg(all(debug_assertions, target_os = "linux"))]
{
module.add_class::<PyTestStreamCapture>()?;
module.add_class::<PyTestDataPlaneHost>()?;
module.add_function(wrap_pyfunction!(_test_data_plane_host, module)?)?;
}

View file

@ -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])

View file

@ -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"

View file

@ -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<Reservation>,
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<Self> {
pub fn from_byte(byte: u8) -> Option<Self> {
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<u8>,
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<RingHandl
offset: lease.layout.start_offset,
capacity: spec.capacity,
generation: spec.generation,
lease_id: lease.ring_id.0,
})
}
@ -339,6 +525,38 @@ pub fn attach(
Ok(endpoint)
}
/// Attach through a child/host process-local mapping of the same arena.
pub fn attach_mapped(
arena: &crate::mapped_arena::MappedArena,
handle: RingHandle,
role: Role,
) -> Result<Endpoint, AttachError> {
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<u8> {
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<WritableRecord<'_>, 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<Option<RecordMeta>, 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<Option<PinnedRecord<'_>>, 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<Option<(RecordKind, Vec<u8>)>, FlowError> {

File diff suppressed because it is too large Load diff

View file

@ -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<Arc<dyn BlobTransferSender>>,
pub source_publisher: Option<Arc<dyn BlobSourcePublisher>>,
pub route_registrar: Option<Arc<dyn HostRouteRegistrar>>,
pub stream_transport: Option<Arc<dyn StreamTransport>>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@ -75,9 +85,12 @@ pub struct HostDataPlaneSessionActor {
source_sender: Option<Arc<dyn BlobTransferSender>>,
source_publisher: Option<Arc<dyn BlobSourcePublisher>>,
route_registrar: Option<Arc<dyn HostRouteRegistrar>>,
stream_arena: Option<Arc<MappedArena>>,
stream_transport: Option<Arc<dyn StreamTransport>>,
child_session: Option<ActorAddress>,
allocator: Option<ActorAddress>,
active_bindings: HashSet<ActorAddress>,
stream_bindings: HashMap<ActorAddress, ActorAddress>,
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<RingHandle, DataPlaneError> {
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<u8>,
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<MappedArena>,
transport: Arc<dyn StreamTransport>,
local_descriptor: StreamPeerDescriptor,
ring: Option<RingHandle>,
matched: Option<StreamMatch>,
peer_descriptor: Option<StreamPeerDescriptor>,
transport_installed: bool,
transport_ready: bool,
opened: bool,
transport_quiesced: bool,
terminal: Option<DataPlaneError>,
data_waiters: Vec<ActorAddress>,
capacity_waiters: Vec<ActorAddress>,
close_waiters: Vec<ActorAddress>,
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<dyn StreamTransportNotifier> = 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<ActorAddress>,
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 {

View file

@ -25,3 +25,4 @@ pub mod path;
pub mod protocol;
pub mod ring;
pub mod source;
pub mod stream_transport;

View file

@ -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<u8>,
pub sink_descriptor: Vec<u8>,
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<u8>,
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<MutationReceipt, NamespaceError>,
},
StreamOpened {
request_id: DirectoryRequestId,
authority_epoch: u64,
result: Result<StreamMatch, NamespaceError>,
},
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<u8>,
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<u8>,
reply_to: ActorAddress,
revision: u64,
}
struct StreamOpenRequest {
request_id: DirectoryRequestId,
path: DataPath,
role: StreamRole,
endpoint: ActorAddress,
replace: bool,
descriptor: Vec<u8>,
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<DataPath, RuntimeSource>,
streams: BTreeMap<DataPath, RuntimeStream>,
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<ActorAddress>,
) -> Result<MutationReceipt, NamespaceError> {
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<StreamMatch, NamespaceError>,
) {
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<BlobBinding, NamespaceError> {
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<StreamMatch, NamespaceError> {
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<StreamMatch, NamespaceError> {
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<StreamMatch, NamespaceError> {
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<StreamMatch, NamespaceError> {
let inbox = self
.runtime
.new_inbox::<NamespaceClientIn>()
.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<StreamMatch, NamespaceError> {
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<StreamMatch, NamespaceError> {
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::<NamespaceClientIn>()
.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) {

View file

@ -72,6 +72,9 @@ pub enum MutationRequest {
length: u64,
recovery: SourceRecovery,
},
BindStream {
path: DataPath,
},
Unregister {
path: DataPath,
},

View file

@ -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<StreamMatch, NamespaceError>),
Allocated(Result<RingHandle, DataPlaneError>),
PeerOffer {
incarnation: StreamIncarnation,
descriptor: StreamPeerDescriptor,
},
Transport(StreamTransportEvent),
DataAvailable,
CapacityAvailable,
WaitData {
reply_to: ActorAddress,
},
WaitCapacity {
reply_to: ActorAddress,
},
Close {
clean: bool,
reply_to: Option<ActorAddress>,
},
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::<HostSessionIn, _>(JsonCodec::default());
registry.register::<ChildSessionIn, _>(JsonCodec::default());
registry.register::<HostStreamIn, _>(JsonCodec::default());
crate::namespace::register_namespace_codecs(registry);
crate::blob_transfer::register_blob_transfer_codecs(registry);
crate::source::register_blob_source_codecs(registry);

View file

@ -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<u8>);
#[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<dyn StreamTransportNotifier>,
}
pub struct StreamSinkRequest {
pub incarnation: StreamIncarnation,
pub endpoint: Endpoint,
pub notifier: Arc<dyn StreamTransportNotifier>,
}
/// 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<StreamPeerDescriptor, String>;
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<dyn StreamTransportNotifier>,
}
struct LocalSink {
endpoint: Endpoint,
notifier: Arc<dyn StreamTransportNotifier>,
}
#[derive(Default)]
struct LocalTransfer {
source: Option<LocalSource>,
sink: Option<LocalSink>,
ready: bool,
}
#[derive(Default)]
struct LocalState {
transfers: BTreeMap<StreamIncarnation, LocalTransfer>,
}
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<LocalState>,
}
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<StreamPeerDescriptor, String> {
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:?}")
}

View file

@ -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<data_plane::stream_transport::StreamPeerDescriptor, String> {
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<parking_lot::Mutex<Vec<u8>>>);
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<dyn data_plane::stream_transport::StreamTransport>,
) -> 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<u8> = (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<dyn data_plane::data_plane::StreamConsumer> =
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");
}

View file

@ -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<u8> = (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<u8> = 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<u8> = view
.spans()
.0
.iter()
.chain(view.spans().1)
.copied()
.collect();
assert_eq!(observed, (0..17).collect::<Vec<_>>());
}

View file

@ -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::<u8>(), 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::<DataPath, TypedModelEntry>::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
)),
}
}
}
}
}
}

View file

@ -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(),
)

View file

@ -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<Vec<StreamTransportEvent>>);
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::<u8>(), 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);
}
}

View file

@ -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<IrohStreamTransport>,
conns: Arc<Mutex<ConnCache>>,
peer_auth: Option<Arc<Mutex<PeerAllowList>>>,
/// 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::<Result<Endpoint, String>>();
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<dyn std::error::Error> { 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<IrohStreamTransport> {
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();

View file

@ -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,

View file

@ -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<AtomicBool>,
cancelled: Arc<AtomicBool>,
}
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<dyn StreamTransportNotifier>,
}
#[derive(Default)]
struct TransportState {
pending_sinks: BTreeMap<StreamIncarnation, PendingSink>,
controls: BTreeMap<StreamIncarnation, Vec<TaskControl>>,
source_probes: BTreeMap<StreamIncarnation, RingProbe>,
sink_probes: BTreeMap<StreamIncarnation, RingProbe>,
}
struct Inner {
engine: EngineHandle,
endpoint: IrohEndpoint,
state: Mutex<TransportState>,
}
/// Cloneable adapter capability installed into `HostDataPlaneConfig`.
#[derive(Clone)]
pub struct IrohStreamTransport {
inner: Arc<Inner>,
}
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<StreamPeerDescriptor, String> {
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<dyn StreamTransportNotifier>,
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<dyn StreamTransportNotifier>,
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()),
}
}

View file

@ -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<StreamTransportEvent>);
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<StreamTransportEvent>, 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);
}

View file

@ -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),