diff --git a/README.md b/README.md index d51908c..9da222a 100644 --- a/README.md +++ b/README.md @@ -121,17 +121,19 @@ cd examples/ping-pong && ./run.sh A node ships batteries-included: the actor engine (worker threads), a network driver, and the pump/fanout tasks that wire them — you don't assemble the glue. -The driver tracks its own established send/recv pumps per QUIC stream and emits -lifecycle events (edge ready, stream fault, pump stopped): +The driver's edge surface is a deliberately tiny reader/writer port; all edge +logic — lifecycle, ring bookkeeping, ingress parsing — lives in the data-plane: ```rust -// crates/iroh-driver/src/driver_pumps.rs -pub enum DriverEventOut { - DriverEdgeReady { edge_id: EdgeId }, - StreamFault { edge_id: EdgeId }, - PumpStopped { edge_id: EdgeId, ring_id: RingId }, +// crates/data-plane/src/edge_wire.rs — the whole transport contract +pub trait EdgeTransport { + type Writer: EdgeWriter; + type PeerAddr: Clone; + fn open_writer(&mut self, edge_id: EdgeId, peer: &Self::PeerAddr) + -> Result; + fn drain_events(&mut self) -> Vec; } -``` +// crates/iroh-driver: `impl EdgeTransport for IrohDriver` ### `iroh` integration — QUIC, TLS, and NAT traversal @@ -224,13 +226,6 @@ producer → egress ring ──QUIC──▶ ingress ring → consumer What travels on a wire edge is typed, so the receiver knows what arrived: ```rust -// crates/data-plane/src/actor.rs -pub enum EdgeKind { - TokenIn, - Activation, // a hidden-state tensor forwarded to the next shard - TokenOut, -} - // crates/data-plane/src/edge_lifecycle.rs pub struct ObjectSpec { pub kind: ObjectKind, // Activation @@ -239,18 +234,18 @@ pub struct ObjectSpec { } ``` -The edge actor owns the lifecycle: lease a byte range, install an ingress or +The edge runtime owns the lifecycle: lease a byte range, install an ingress or egress ring, and release it only with a `QuiescenceProof`, so a range is never recycled under a live reader. The worker consumes whole objects off its ingress ring as plain bytes: ```rust -// apps/myelin/src/node/worker_node_runtime.rs -match ingress::read_object_record(buffer, object_spec, false)? { - ingress::ObjectRecordRead::Incomplete => Ok(None), // wait for more bytes - ingress::ObjectRecordRead::Complete(record) => { // a full activation arrived +// crates/data-plane/src/edge_runtime.rs +match read_object_record(buffer, object_spec, false)? { + ObjectRecordRead::Incomplete => Ok(None), // wait for more bytes + ObjectRecordRead::Complete(record) => { // a full activation arrived let bytes: Vec = buffer.drain(..record.total_len).collect(); - Ok(Some(IngressRecordBytes { bytes, object_id: record.object_id.0, .. })) + // ... written into the edge's arena ring, then loaded on the worker } } ``` diff --git a/apps/myelin/src/node/worker_node_runtime.rs b/apps/myelin/src/node/worker_node_runtime.rs index 4c2eac0..092035d 100644 --- a/apps/myelin/src/node/worker_node_runtime.rs +++ b/apps/myelin/src/node/worker_node_runtime.rs @@ -37,15 +37,15 @@ use crate::run_plan::{GgufSource, TokenizerSource}; use crate::staging::control as stage; use data_plane::arena; use data_plane::edge_lifecycle as edge; -use data_plane::ingress; +use data_plane::object_record as ingress; use distribution::node::DistributedNodeConfig; use distribution::telemetry::{MembershipTransition, SwimProbeEvent}; use distribution::types::{MemberState, NodeId as DistNodeId}; use iroh::EndpointAddr; -use iroh_driver::driver_pumps as driver_model; +use data_plane::edge_runtime; use iroh_driver::{ - TELEMETRY_ALPN, TelemetryPublishHandle, TelemetryQuicHeader, EDGE_ALPN, EdgeSendHandle, - EdgeTransportEvent, IrohDriver, IrohDriverConfig, + EDGE_ALPN, TELEMETRY_ALPN, TelemetryPublishHandle, TelemetryQuicHeader, IrohDriver, + IrohDriverConfig, }; use iroh_driver::{EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint}; use parking_lot::Mutex; @@ -783,13 +783,6 @@ fn spawn_arena_sampler( }); } -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] -struct ObjectKey { - edge_id: u64, - object_id: u64, -} - -#[derive(Clone, Debug)] struct LoadedObject { object_id: u64, sequence: u64, @@ -797,42 +790,293 @@ struct LoadedObject { handle_id: u64, } +/// All edge logic — lifecycle, wire bookkeeping, ingress parsing, arena and +/// worker orchestration — lives in the data-plane edge runtime. This wrapper +/// only supplies myelin effects (tinygrad worker port, telemetry, node-agent +/// reporting) around it. struct WorkerEdgeRuntime { - establisher: edge::EdgeEstablisher, - driver_model: driver_model::Driver, - edge_command_cursor: usize, - edge_event_cursor: usize, - driver_event_cursor: usize, + runtime: edge_runtime::EdgeRuntime, inbound_edge: Option, outbound_edge: Option, - inbound_ring_id: Option, - outbound_ring_id: Option, - outbound_sender: Option, - next_output_object_id: u64, - object_handles: BTreeMap, - ingress_streams: BTreeMap>, +} + +/// Worker-ring effects the data-plane edge runtime drives over the tinygrad +/// worker. +struct TinygradRingPort<'a> { + worker: &'a mut TinygradWorker, + config: &'a DeploymentConfig, + telemetry: &'a mut NodeTelemetry, + inbound: Option, + outbound: Option, +} + +impl edge_runtime::WorkerPort for TinygradRingPort<'_> { + fn install_ring( + &mut self, + edge_id: edge::EdgeId, + ring_id: edge::RingId, + direction: edge::RingDirection, + layout: &arena::RingLayout, + object_spec: &edge::ObjectSpec, + ) -> Result<(), String> { + let (port, direction_name, wire_spec) = match direction { + edge::RingDirection::Ingress => ( + "input", + "ingress", + self.inbound.as_ref().map(|edge| edge.object_spec), + ), + edge::RingDirection::Egress => ( + "output", + "egress", + self.outbound.as_ref().map(|edge| edge.object_spec), + ), + }; + let wire_spec = wire_spec.unwrap_or(StageObjectSpecWire { + max_extent: object_spec.max_extent_bytes, + alignment: 4, + }); + self.worker.install_ring( + ring_id.0, + edge_id.0, + port, + direction_name, + layout.clone(), + wire_spec, + self.config, + self.telemetry, + ) + } + + fn uninstall_ring(&mut self, ring_id: edge::RingId) -> Result<(), String> { + self.worker + .uninstall_ring(ring_id.0, self.config, self.telemetry) + } + + fn load_object( + &mut self, + edge_id: edge::EdgeId, + ring_id: edge::RingId, + _record: &ingress::ObjectRecord, + spec: &ingress::ObjectSpec, + ) -> Result { + let wire_spec = StageObjectSpecWire { + max_extent: spec.max_extent, + alignment: spec.alignment.min(u64::from(u32::MAX)) as u32, + }; + let loaded = self + .worker + .ring_readable(ring_id.0, edge_id.0, wire_spec, self.config, self.telemetry)?; + Ok(edge_runtime::LoadedObject { + object_id: loaded.object_id, + sequence: loaded.sequence, + handle_generation: loaded.handle_generation, + handle_id: loaded.handle_id, + }) + } } impl WorkerEdgeRuntime { fn new(local_node_id: u64) -> Self { Self { - establisher: edge::EdgeEstablisher::new(edge::NodeId(local_node_id)), - driver_model: driver_model::Driver::new(), - edge_command_cursor: 0, - edge_event_cursor: 0, - driver_event_cursor: 0, + runtime: edge_runtime::EdgeRuntime::new(edge::NodeId(local_node_id)), inbound_edge: None, outbound_edge: None, - inbound_ring_id: None, - outbound_ring_id: None, - outbound_sender: None, - next_output_object_id: 1, - object_handles: BTreeMap::new(), - ingress_streams: BTreeMap::new(), } } + /// Run one data-plane edge-runtime tick, then map its observations to + /// telemetry and node-agent messages. The poll error (if any) propagates + /// after the observations are reported, mirroring the previous + /// composition's report-then-halt behavior. #[allow(clippy::too_many_arguments)] + fn poll_and_report( + &mut self, + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + node_actor: ActorAddress, + worker: &mut TinygradWorker, + arena_manager: &Arc>, + config: &DeploymentConfig, + telemetry: &mut NodeTelemetry, + ) -> Result<(), String> { + let result = { + let mut arena = arena_manager.lock(); + let mut port = TinygradRingPort { + worker, + config, + telemetry, + inbound: self.inbound_edge.clone(), + outbound: self.outbound_edge.clone(), + }; + self.runtime.poll(driver, &mut arena, &mut port) + }; + self.report(stack, node_actor, config, telemetry)?; + result + } + + /// Map drained runtime observations to telemetry events and node-agent + /// messages. + fn report( + &mut self, + stack: &DistributionRuntimeStack, + node_actor: ActorAddress, + config: &DeploymentConfig, + telemetry: &mut NodeTelemetry, + ) -> Result<(), String> { + let node_stage = |ds: &mut NodeTelemetry, phase: &str, status: &str, detail: Value| { + emit_node_event(ds, config, NODE_STAGE_CHANNEL, phase, status, detail) + }; + for observation in self.runtime.take_observations() { + match observation { + edge_runtime::Observation::StreamArrived { edge_id, stream_id } => { + node_stage( + telemetry, + "iroh_edge_stream_arrived", + "observed", + json!({"edge_id":edge_id.0,"stream_id":stream_id.0}), + ); + } + edge_runtime::Observation::BytesRead { + edge_id, + stream_id, + byte_count, + } => { + node_stage( + telemetry, + "iroh_edge_bytes_read", + "observed", + json!({"edge_id":edge_id.0,"stream_id":stream_id.0,"bytes":byte_count}), + ); + } + edge_runtime::Observation::IngressRingWrite { + edge_id, + ring_id, + stream_id, + object_id, + sequence, + extent, + begin_sequence, + end_of_sequence, + record_bytes, + buffered_bytes, + write_ms, + } => { + let edge_kind = self + .inbound_edge + .as_ref() + .map(|edge| format!("{:?}", edge.kind)) + .unwrap_or_else(|| "unknown".to_owned()); + node_stage( + telemetry, + "ingress_ring_write", + "ready", + json!({ + "edge_id":edge_id.0, + "edge_kind":edge_kind, + "ring_id":ring_id.0, + "stream_id":stream_id.0, + "object_id":object_id, + "sequence":sequence, + "extent":extent, + "begin_sequence":begin_sequence, + "end_of_sequence":end_of_sequence, + "record_bytes":record_bytes, + "ingress_buffer_bytes":buffered_bytes, + "ingress_ring_write_ms":write_ms, + }), + ); + } + edge_runtime::Observation::ObjectLoaded { + edge_id, + ring_id, + stream_id, + object, + load_ms, + } => { + let edge_kind = self + .inbound_edge + .as_ref() + .map(|edge| format!("{:?}", edge.kind)) + .unwrap_or_else(|| "unknown".to_owned()); + node_stage( + telemetry, + "object_loaded", + "ready", + json!({ + "edge_id":edge_id.0, + "edge_kind":edge_kind, + "ring_id":ring_id.0, + "stream_id":stream_id.0, + "object_id":object.object_id, + "sequence":object.sequence, + "handle_generation":object.handle_generation, + "handle_id":object.handle_id, + "object_load_ms":load_ms, + }), + ); + stack + .runtime + .send_to( + node_actor, + NodeAgentMsg::ObjectLoaded { + edge_id: edge_id.0, + object_id: object.object_id, + sequence: object.sequence, + handle_generation: object.handle_generation, + handle_id: object.handle_id, + }, + ) + .map_err(|e| format!("report object loaded: {e}"))?; + } + edge_runtime::Observation::ObjectFailed { edge_id, object_id } => { + let _ = stack.runtime.send_to( + node_actor, + NodeAgentMsg::ObjectFailed { + edge_id: edge_id.0, + object_id, + }, + ); + } + edge_runtime::Observation::EdgeReady { edge_id, .. } => { + if self + .inbound_edge + .as_ref() + .is_some_and(|edge| edge.edge_id == edge_id.0) + { + stack + .runtime + .send_to( + node_actor, + NodeAgentMsg::MarkInboundEdgeReady { edge_id: edge_id.0 }, + ) + .map_err(|e| format!("mark inbound ready: {e}"))?; + } + if self + .outbound_edge + .as_ref() + .is_some_and(|edge| edge.edge_id == edge_id.0) + { + stack + .runtime + .send_to( + node_actor, + NodeAgentMsg::MarkOutboundEdgeReady { edge_id: edge_id.0 }, + ) + .map_err(|e| format!("mark outbound ready: {e}"))?; + } + } + edge_runtime::Observation::EdgeFaulted { edge_id, .. } => { + stack + .runtime + .send_to(node_actor, NodeAgentMsg::EdgeFault { edge_id: edge_id.0 }) + .map_err(|e| format!("report edge fault: {e}"))?; + } + edge_runtime::Observation::EdgeStopped { .. } => {} + } + } + Ok(()) + } + fn poll_iroh( &mut self, driver: &mut IrohDriver, @@ -843,86 +1087,12 @@ impl WorkerEdgeRuntime { config: &DeploymentConfig, telemetry: &mut NodeTelemetry, ) -> Result<(), String> { - let node_stage = |ds: &mut NodeTelemetry, phase: &str, status: &str, detail: Value| { - emit_node_event(ds, config, NODE_STAGE_CHANNEL, phase, status, detail) - }; - for event in driver.drain_edge_events() { - match event { - EdgeTransportEvent::StreamArrived { - edge_id, stream_id, .. - } => { - self.driver_model.incoming_uni_stream( - driver_model::EdgeId(edge_id), - driver_model::StreamId(stream_id), - ); - node_stage( - telemetry, - "iroh_edge_stream_arrived", - "observed", - json!({"edge_id":edge_id,"stream_id":stream_id}), - ); - self.drive_edge_workflow( - stack, - node_actor, - worker, - arena_manager, - config, - telemetry, - driver, - )?; - } - EdgeTransportEvent::BytesRead { - edge_id, - stream_id, - bytes, - .. - } => { - let byte_count = bytes.len(); - node_stage( - telemetry, - "iroh_edge_bytes_read", - "observed", - json!({"edge_id":edge_id,"stream_id":stream_id,"bytes":byte_count}), - ); - self.ingest_stream_bytes( - edge_id, - stream_id, - bytes, - stack, - node_actor, - worker, - arena_manager, - config, - telemetry, - driver, - )?; - } - EdgeTransportEvent::StreamFault { - edge_id: Some(edge_id), - .. - } => { - self.driver_model.read_error(driver_model::EdgeId(edge_id)); - self.drive_edge_workflow( - stack, - node_actor, - worker, - arena_manager, - config, - telemetry, - driver, - )?; - } - EdgeTransportEvent::StreamEnded { .. } - | EdgeTransportEvent::StreamFault { edge_id: None, .. } => {} - } - } - Ok(()) + self.poll_and_report(driver, stack, node_actor, worker, arena_manager, config, telemetry) } - #[allow(clippy::too_many_arguments)] fn establish_inbound( &mut self, - edge: StageInboundEdgeWire, + edge_wire: StageInboundEdgeWire, stack: &DistributionRuntimeStack, node_actor: ActorAddress, worker: &mut TinygradWorker, @@ -931,38 +1101,36 @@ impl WorkerEdgeRuntime { telemetry: &mut NodeTelemetry, driver: &mut IrohDriver, ) -> Result<(), String> { - self.inbound_edge = Some(edge.clone()); - self.establisher - .observe(edge::EdgeEvent::ProvisionRx(edge::ProvisionRx { + let parse_spec = ingress::ObjectSpec { + max_extent: edge_wire.object_spec.max_extent, + alignment: u64::from(edge_wire.object_spec.alignment), + layout: ingress::ObjectLayout::Token, + }; + self.runtime.establish_inbound( + edge::ProvisionRx { run_id: edge::RunId(config.run_id), - edge_id: edge::EdgeId(edge.edge_id), + edge_id: edge::EdgeId(edge_wire.edge_id), local_node_id: edge::NodeId(config.logical_node_id), object_spec: edge::ObjectSpec { kind: edge::ObjectKind::Activation, dtype: edge::DType::F16, - max_extent_bytes: edge.object_spec.max_extent, + max_extent_bytes: edge_wire.object_spec.max_extent, }, ring_spec: edge::RingSpec { header_bytes: 0, - data_bytes: edge.ring_spec.data_capacity, - alignment: u64::from(edge.ring_spec.alignment), + data_bytes: edge_wire.ring_spec.data_capacity, + alignment: u64::from(edge_wire.ring_spec.alignment), }, - })); - self.drive_edge_workflow( - stack, - node_actor, - worker, - arena_manager, - config, - telemetry, - driver, - ) + }, + parse_spec, + ); + self.inbound_edge = Some(edge_wire); + self.poll_and_report(driver, stack, node_actor, worker, arena_manager, config, telemetry) } - #[allow(clippy::too_many_arguments)] fn establish_outbound( &mut self, - edge: StageOutboundEdgeWire, + edge_wire: StageOutboundEdgeWire, stack: &DistributionRuntimeStack, node_actor: ActorAddress, worker: &mut TinygradWorker, @@ -971,49 +1139,46 @@ impl WorkerEdgeRuntime { telemetry: &mut NodeTelemetry, driver: &mut IrohDriver, ) -> Result<(), String> { - if edge.consumer_endpoint.is_none() { + if edge_wire.consumer_endpoint.is_none() { stack .runtime .send_to( node_actor, NodeAgentMsg::MarkOutboundEdgeReady { - edge_id: edge.edge_id, + edge_id: edge_wire.edge_id, }, ) .map_err(|e| format!("mark outbound edge ready: {e}"))?; - self.outbound_edge = Some(edge); + self.outbound_edge = Some(edge_wire); return Ok(()); } - self.outbound_edge = Some(edge.clone()); - self.establisher - .observe(edge::EdgeEvent::ProvisionTx(edge::ProvisionTx { + let peer = edge_wire + .consumer_endpoint + .clone() + .expect("consumer endpoint presence checked above"); + self.runtime.establish_outbound( + edge::ProvisionTx { run_id: edge::RunId(config.run_id), - edge_id: edge::EdgeId(edge.edge_id), + edge_id: edge::EdgeId(edge_wire.edge_id), local_node_id: edge::NodeId(config.logical_node_id), - consumer_node_id: edge::NodeId(edge.consumer_node_id), + consumer_node_id: edge::NodeId(edge_wire.consumer_node_id), object_spec: edge::ObjectSpec { kind: edge::ObjectKind::Activation, dtype: edge::DType::F16, - max_extent_bytes: edge.object_spec.max_extent, + max_extent_bytes: edge_wire.object_spec.max_extent, }, ring_spec: edge::RingSpec { header_bytes: 0, - data_bytes: edge.ring_spec.data_capacity, - alignment: u64::from(edge.ring_spec.alignment), + data_bytes: edge_wire.ring_spec.data_capacity, + alignment: u64::from(edge_wire.ring_spec.alignment), }, - })); - self.drive_edge_workflow( - stack, - node_actor, - worker, - arena_manager, - config, - telemetry, - driver, - ) + }, + peer, + ); + self.outbound_edge = Some(edge_wire); + self.poll_and_report(driver, stack, node_actor, worker, arena_manager, config, telemetry) } - #[allow(clippy::too_many_arguments)] fn execute_step( &mut self, step_id: u64, @@ -1031,18 +1196,16 @@ impl WorkerEdgeRuntime { let node_stage = |ds: &mut NodeTelemetry, phase: &str, status: &str, detail: Value| { emit_node_event(ds, config, NODE_STAGE_CHANNEL, phase, status, detail) }; - let input_key = ObjectKey { - edge_id: input_edge_id, - object_id, - }; let loaded = self - .object_handles - .get(&input_key) + .runtime + .loaded_object(edge::EdgeId(input_edge_id), object_id) .cloned() - .ok_or_else(|| format!("object {input_key:?} has no loaded device handle"))?; + .ok_or_else(|| { + format!("object (edge {input_edge_id}, id {object_id}) has no loaded device handle") + })?; if loaded.sequence != sequence { return Err(format!( - "object {input_key:?} sequence {} does not match command sequence {sequence}", + "object (edge {input_edge_id}, id {object_id}) sequence {} does not match command sequence {sequence}", loaded.sequence )); } @@ -1051,10 +1214,11 @@ impl WorkerEdgeRuntime { .clone() .ok_or_else(|| "outbound edge missing".to_owned())?; let output_ring_id = self - .outbound_ring_id - .ok_or_else(|| "outbound ring missing".to_owned())?; - let output_object_id = self.next_output_object_id; - self.next_output_object_id = self.next_output_object_id.saturating_add(1); + .runtime + .outbound_ring_id() + .ok_or_else(|| "outbound ring missing".to_owned())? + .0; + let output_object_id = self.runtime.alloc_output_object_id()?; let final_stage = matches!( outbound.kind, crate::node_actor::StageEdgeKindWire::TokenOut @@ -1126,11 +1290,11 @@ impl WorkerEdgeRuntime { "egress_ring_read_ms":egress_read_ms, }), ); - let sender = self - .outbound_sender - .as_ref() - .ok_or_else(|| "outbound edge sender missing".to_owned())?; let edge_send_started = Instant::now(); + let sender = self + .runtime + .outbound_writer() + .ok_or_else(|| "outbound edge sender missing".to_owned())?; if let Err(e) = sender.send(record) { let _ = stack.runtime.send_to( node_actor, @@ -1173,481 +1337,6 @@ impl WorkerEdgeRuntime { ) -> Result<(), String> { worker.release_device_object(handle_id, config, telemetry) } - - #[allow(clippy::too_many_arguments)] - fn ingest_stream_bytes( - &mut self, - edge_id: u64, - stream_id: u64, - bytes: Vec, - stack: &DistributionRuntimeStack, - node_actor: ActorAddress, - worker: &mut TinygradWorker, - arena_manager: &Arc>, - config: &DeploymentConfig, - telemetry: &mut NodeTelemetry, - driver: &mut IrohDriver, - ) -> Result<(), String> { - let node_stage = |ds: &mut NodeTelemetry, phase: &str, status: &str, detail: Value| { - emit_node_event(ds, config, NODE_STAGE_CHANNEL, phase, status, detail) - }; - let Some(inbound) = self.inbound_edge.clone() else { - return Ok(()); - }; - if inbound.edge_id != edge_id { - return Ok(()); - } - let (records, buffered_bytes) = { - let buffer = self.ingress_streams.entry(stream_id).or_default(); - buffer.extend_from_slice(&bytes); - let buffered_bytes = buffer.len(); - let mut records = Vec::new(); - loop { - match take_complete_ingress_record(buffer, inbound.object_spec) { - Ok(Some(record)) => records.push(record), - Ok(None) => break, - Err(e) => { - let _ = stack.runtime.send_to( - node_actor, - NodeAgentMsg::ObjectFailed { - edge_id, - object_id: None, - }, - ); - return Err(e); - } - } - } - (records, buffered_bytes) - }; - for record in records { - let ring_id = self - .inbound_ring_id - .ok_or_else(|| "inbound ring missing".to_owned())?; - let ring_write_started = Instant::now(); - { - let arena = arena_manager.lock(); - let lease = arena - .lookup_lease(arena::RingId(ring_id)) - .ok_or_else(|| format!("inbound ring {ring_id} lease missing"))?; - arena - .write_arena(lease.layout.data_offset, &record.bytes) - .map_err(|e| format!("write ingress ring: {e}"))?; - } - let ingress_ring_write_ms = duration_ms_u64(ring_write_started.elapsed()); - node_stage( - telemetry, - "ingress_ring_write", - "ready", - json!({ - "edge_id":edge_id, - "edge_kind":format!("{:?}", inbound.kind), - "ring_id":ring_id, - "stream_id":stream_id, - "object_id":record.object_id, - "sequence":record.sequence, - "extent":record.extent, - "begin_sequence":record.begin_sequence, - "end_of_sequence":record.end_of_sequence, - "record_bytes":record.bytes.len(), - "ingress_buffer_bytes":buffered_bytes, - "ingress_ring_write_ms":ingress_ring_write_ms, - }), - ); - let object_load_started = Instant::now(); - let loaded = match worker.ring_readable( - ring_id, - edge_id, - inbound.object_spec, - config, - telemetry, - ) { - Ok(loaded) => loaded, - Err(e) => { - let _ = stack.runtime.send_to( - node_actor, - NodeAgentMsg::ObjectFailed { - edge_id, - object_id: Some(record.object_id), - }, - ); - return Err(e); - } - }; - let object_load_ms = duration_ms_u64(object_load_started.elapsed()); - let key = ObjectKey { - edge_id, - object_id: loaded.object_id, - }; - self.object_handles.insert(key, loaded.clone()); - node_stage( - telemetry, - "object_loaded", - "ready", - json!({ - "edge_id":edge_id, - "edge_kind":format!("{:?}", inbound.kind), - "ring_id":ring_id, - "stream_id":stream_id, - "object_id":loaded.object_id, - "sequence":loaded.sequence, - "handle_generation":loaded.handle_generation, - "handle_id":loaded.handle_id, - "object_load_ms":object_load_ms, - }), - ); - stack - .runtime - .send_to( - node_actor, - NodeAgentMsg::ObjectLoaded { - edge_id, - object_id: loaded.object_id, - sequence: loaded.sequence, - handle_generation: loaded.handle_generation, - handle_id: loaded.handle_id, - }, - ) - .map_err(|e| format!("report object loaded: {e}"))?; - } - self.drive_edge_workflow( - stack, - node_actor, - worker, - arena_manager, - config, - telemetry, - driver, - ) - } - - #[allow(clippy::too_many_arguments)] - fn drive_edge_workflow( - &mut self, - stack: &DistributionRuntimeStack, - node_actor: ActorAddress, - worker: &mut TinygradWorker, - arena_manager: &Arc>, - config: &DeploymentConfig, - telemetry: &mut NodeTelemetry, - driver: &mut IrohDriver, - ) -> Result<(), String> { - loop { - let progressed = - self.drain_edge_commands(worker, arena_manager, config, telemetry, driver)? - || self.drain_driver_events() - || self.drain_edge_events(stack, node_actor)?; - if !progressed { - break; - } - } - Ok(()) - } - - fn drain_edge_commands( - &mut self, - worker: &mut TinygradWorker, - arena_manager: &Arc>, - config: &DeploymentConfig, - telemetry: &mut NodeTelemetry, - driver: &mut IrohDriver, - ) -> Result { - let mut progressed = false; - while self.edge_command_cursor < self.establisher.commands().len() { - let command = self.establisher.commands()[self.edge_command_cursor].clone(); - self.edge_command_cursor += 1; - progressed = true; - match command { - edge::EdgeCommand::LeaseRing { - request_id, - ring_spec, - .. - } => { - let events = arena_manager.lock().request(arena::ArenaRequest::LeaseRing( - arena::LeaseRing { - request_id: arena::LeaseRequestId(request_id.0), - ring_spec: arena::RingSpec { - header_bytes: ring_spec.header_bytes, - data_bytes: ring_spec.data_bytes, - alignment: ring_spec.alignment, - }, - }, - )); - for event in events { - match event { - arena::ArenaEvent::RingLeased { lease } => { - self.establisher.observe(edge::EdgeEvent::RingLeased { - request_id: edge::LeaseRequestId(lease.request_id.0), - ring_id: edge::RingId(lease.ring_id.0), - layout: edge::RingLayout { - start_offset: lease.layout.start_offset, - header_offset: lease.layout.header_offset, - data_offset: lease.layout.data_offset, - end_offset: lease.layout.end_offset, - data_bytes: lease.layout.data_bytes, - alignment: lease.layout.alignment, - }, - }); - } - arena::ArenaEvent::RingLeaseRejected { request_id, reason } => { - let reason = match reason { - arena::RingLeaseRejection::CannotFitWithinCeiling => { - edge::RingLeaseRejection::CannotFit - } - arena::RingLeaseRejection::ArenaShuttingDown => { - edge::RingLeaseRejection::ArenaShuttingDown - } - }; - self.establisher - .observe(edge::EdgeEvent::RingLeaseRejected { - request_id: edge::LeaseRequestId(request_id.0), - reason, - }); - } - arena::ArenaEvent::RingLeaseQueued { .. } - | arena::ArenaEvent::RingReleased { .. } - | arena::ArenaEvent::RingReleaseRejected { .. } - | arena::ArenaEvent::CancelledFreshLeaseReleased { .. } => {} - } - } - } - edge::EdgeCommand::InstallWorkerRing { - edge_id, - ring_id, - direction, - object_spec, - .. - } => { - let lease = arena_manager - .lock() - .lookup_lease(arena::RingId(ring_id.0)) - .ok_or_else(|| format!("ring {} lease missing", ring_id.0))? - .clone(); - let (port, direction_name, wire_spec) = match direction { - edge::RingDirection::Ingress => { - self.inbound_ring_id = Some(ring_id.0); - let spec = self - .inbound_edge - .as_ref() - .map(|edge| edge.object_spec) - .unwrap_or(StageObjectSpecWire { - max_extent: object_spec.max_extent_bytes, - alignment: 4, - }); - ("input", "ingress", spec) - } - edge::RingDirection::Egress => { - self.outbound_ring_id = Some(ring_id.0); - let spec = self - .outbound_edge - .as_ref() - .map(|edge| edge.object_spec) - .unwrap_or(StageObjectSpecWire { - max_extent: object_spec.max_extent_bytes, - alignment: 4, - }); - ("output", "egress", spec) - } - }; - worker.install_ring( - ring_id.0, - edge_id.0, - port, - direction_name, - lease.layout, - wire_spec, - config, - telemetry, - )?; - self.establisher - .observe(edge::EdgeEvent::RingInstalled { edge_id, ring_id }); - } - edge::EdgeCommand::EstablishSend { edge_id, .. } => { - let outbound = self - .outbound_edge - .as_ref() - .ok_or_else(|| "outbound edge missing".to_owned())?; - let peer = outbound - .consumer_endpoint - .clone() - .ok_or_else(|| "outbound consumer endpoint missing".to_owned())?; - let record = self - .establisher - .local_record(edge_id) - .ok_or_else(|| format!("edge {} record missing", edge_id.0))?; - let ring_id = record - .ring_id - .ok_or_else(|| format!("edge {} ring missing", edge_id.0))?; - self.driver_model.establish_send( - driver_model::EdgeId(edge_id.0), - driver_model::RingId(ring_id.0), - ); - self.outbound_sender = Some(driver.spawn_edge_send_pump(peer, edge_id.0)?); - } - edge::EdgeCommand::EstablishRecv { edge_id, .. } => { - let record = self - .establisher - .local_record(edge_id) - .ok_or_else(|| format!("edge {} record missing", edge_id.0))?; - let ring_id = record - .ring_id - .ok_or_else(|| format!("edge {} ring missing", edge_id.0))?; - self.driver_model.establish_recv( - driver_model::EdgeId(edge_id.0), - driver_model::RingId(ring_id.0), - ); - } - edge::EdgeCommand::CancelQueuedLease { request_id, .. } => { - let _ = arena_manager - .lock() - .request(arena::ArenaRequest::CancelLease { - request_id: arena::LeaseRequestId(request_id.0), - }); - } - edge::EdgeCommand::StopPump { edge_id, .. } => { - self.driver_model.stop_edge(driver_model::EdgeId(edge_id.0)); - } - edge::EdgeCommand::UninstallWorkerRing { ring_id, .. } => { - worker.uninstall_ring(ring_id.0, config, telemetry)?; - self.establisher - .observe(edge::EdgeEvent::RingQuiesced { ring_id }); - } - edge::EdgeCommand::ReleaseArenaLease { ring_id, proof } => { - let proof = if proof == edge::QuiescenceProof::verified() { - arena::QuiescenceProof::verified() - } else { - arena::QuiescenceProof::missing() - }; - let _ = arena_manager - .lock() - .request(arena::ArenaRequest::ReleaseRing { - ring_id: arena::RingId(ring_id.0), - proof, - }); - } - } - } - Ok(progressed) - } - - fn drain_driver_events(&mut self) -> bool { - let mut progressed = false; - while self.driver_event_cursor < self.driver_model.events().len() { - let event = self.driver_model.events()[self.driver_event_cursor].clone(); - self.driver_event_cursor += 1; - progressed = true; - match event { - driver_model::DriverEventOut::DriverEdgeReady { edge_id } => { - self.establisher.observe(edge::EdgeEvent::DriverEdgeReady { - edge_id: edge::EdgeId(edge_id.0), - }); - } - driver_model::DriverEventOut::StreamFault { edge_id } => { - self.establisher.observe(edge::EdgeEvent::StreamFault { - edge_id: edge::EdgeId(edge_id.0), - reason: edge::StreamFaultReason::ReadError, - }); - } - driver_model::DriverEventOut::PumpStopped { edge_id, ring_id } => { - self.establisher.observe(edge::EdgeEvent::PumpStopped { - edge_id: edge::EdgeId(edge_id.0), - ring_id: edge::RingId(ring_id.0), - }); - } - } - } - progressed - } - - fn drain_edge_events( - &mut self, - stack: &DistributionRuntimeStack, - node_actor: ActorAddress, - ) -> Result { - let mut progressed = false; - while self.edge_event_cursor < self.establisher.events().len() { - let event = self.establisher.events()[self.edge_event_cursor].clone(); - self.edge_event_cursor += 1; - progressed = true; - match event { - edge::EdgeLifecycleEvent::EdgeReady { edge_id, .. } => { - if self - .inbound_edge - .as_ref() - .is_some_and(|edge| edge.edge_id == edge_id.0) - { - stack - .runtime - .send_to( - node_actor, - NodeAgentMsg::MarkInboundEdgeReady { edge_id: edge_id.0 }, - ) - .map_err(|e| format!("mark inbound ready: {e}"))?; - } - if self - .outbound_edge - .as_ref() - .is_some_and(|edge| edge.edge_id == edge_id.0) - { - stack - .runtime - .send_to( - node_actor, - NodeAgentMsg::MarkOutboundEdgeReady { edge_id: edge_id.0 }, - ) - .map_err(|e| format!("mark outbound ready: {e}"))?; - } - } - edge::EdgeLifecycleEvent::EdgeFaulted { edge_id, reason } => { - stack - .runtime - .send_to(node_actor, NodeAgentMsg::EdgeFault { edge_id: edge_id.0 }) - .map_err(|e| format!("report edge fault: {e}"))?; - return Err(format!("edge {} faulted: {reason:?}", edge_id.0)); - } - edge::EdgeLifecycleEvent::EdgeStopped { .. } => {} - } - } - Ok(progressed) - } -} - -struct IngressRecordBytes { - bytes: Vec, - object_id: u64, - sequence: u64, - extent: u64, - begin_sequence: bool, - end_of_sequence: bool, -} - -fn take_complete_ingress_record( - buffer: &mut Vec, - spec: StageObjectSpecWire, -) -> Result, String> { - let record = match ingress::read_object_record( - buffer, - ingress::ObjectSpec { - max_extent: spec.max_extent, - alignment: u64::from(spec.alignment), - layout: ingress::ObjectLayout::Token, - }, - false, - ) - .map_err(|reason| format!("invalid object record: {reason:?}"))? - { - ingress::ObjectRecordRead::Incomplete => return Ok(None), - ingress::ObjectRecordRead::Complete(record) => record, - }; - let bytes = buffer.drain(..record.total_len).collect(); - Ok(Some(IngressRecordBytes { - bytes, - object_id: record.object_id.0, - sequence: record.sequence, - extent: record.extent, - begin_sequence: record.flags.begin_sequence, - end_of_sequence: record.flags.end_of_sequence, - })) } fn value_u64(value: &Value, field: &str) -> Result { diff --git a/apps/myelin/src/orchestration/app.rs b/apps/myelin/src/orchestration/app.rs index 62a9ffa..d930d44 100644 --- a/apps/myelin/src/orchestration/app.rs +++ b/apps/myelin/src/orchestration/app.rs @@ -56,8 +56,9 @@ use distribution::node::DistributedNodeConfig; use distribution::swim::telemetry::ObservedTransition; use distribution::types::{MemberState, NodeId as DistNodeId}; use iroh::EndpointAddr; +use data_plane::edge_wire::WireEvent; use iroh_driver::{ - TELEMETRY_ALPN, EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, IrohDriver, IrohDriverConfig, + TELEMETRY_ALPN, EDGE_ALPN, EdgeSendHandle, IrohDriver, IrohDriverConfig, }; use iroh_driver::{EndpointAddrMask, MVP_IROH_ENDPOINT_ADDR_MASK_ENV, advertised_endpoint}; use parking_lot::Mutex; @@ -4259,17 +4260,17 @@ impl PipelinePromptRuntime { fn poll_driver(&mut self, driver: &mut IrohDriver) { for event in driver.drain_edge_events() { match event { - EdgeTransportEvent::BytesRead { edge_id, bytes, .. } - if edge_id == self.token_out_edge_id => + WireEvent::BytesRead { edge_id, bytes, .. } + if edge_id.0 == self.token_out_edge_id => { self.note_progress(); let _ = self.recv_tx.send(bytes); } - EdgeTransportEvent::StreamFault { + WireEvent::StreamFault { edge_id: Some(edge_id), reason, .. - } if edge_id == self.token_out_edge_id => { + } if edge_id.0 == self.token_out_edge_id => { if let Some(request_id) = self.active.as_ref().map(|active| active.request.request_id) { diff --git a/apps/myelin/src/tests/local_mock/environment.rs b/apps/myelin/src/tests/local_mock/environment.rs index b7d2512..e6a2087 100644 --- a/apps/myelin/src/tests/local_mock/environment.rs +++ b/apps/myelin/src/tests/local_mock/environment.rs @@ -3,7 +3,7 @@ use std::collections::{BTreeMap, BTreeSet}; use crate::run_fsm as fsm; use crate::run_plan as plan; use crate::tests::harness::OrchestratorHarness; -use data_plane::edge_actor; +use data_plane::object_record::ObjectIdAllocator; use myelin::observability::lifecycle as obs; use myelin::orchestration::engine_builder as engine; use myelin::staging as stage; @@ -42,7 +42,7 @@ pub struct LocalMockCluster { transport: MockTransport, resources: ResourceTracker, observed_edges: BTreeSet, - object_allocators: BTreeMap, + object_allocators: BTreeMap, scenario: LocalMockScenario, } @@ -811,9 +811,8 @@ impl LocalMockCluster { fn allocate_object_id(&mut self, edge_id: plan::EdgeId) -> u64 { self.object_allocators .entry(edge_id) - .or_insert_with(|| edge_actor::ObjectIdAllocator::new(edge_actor::EdgeId(edge_id.0))) + .or_insert_with(ObjectIdAllocator::new) .alloc() - .object_id .0 } diff --git a/apps/myelin/src/tests/local_mock/mock_node.rs b/apps/myelin/src/tests/local_mock/mock_node.rs index a0cbb6e..9544447 100644 --- a/apps/myelin/src/tests/local_mock/mock_node.rs +++ b/apps/myelin/src/tests/local_mock/mock_node.rs @@ -1,6 +1,6 @@ use crate::run_plan as plan; use crate::tests::harness::StageControllerHarness; -use data_plane::edge_actor; +use data_plane::object_record::ObjectIdAllocator; use myelin::staging as stage; use super::mock_transport::MockObject; @@ -16,7 +16,7 @@ pub struct MockNode { stage_count: u32, inbound_edge: Option, outbound_edge: Option, - outbound_object_allocator: Option, + outbound_object_allocator: Option, controller: StageControllerHarness, worker: MockWorker, event_cursor: usize, @@ -44,9 +44,7 @@ impl MockNode { pub fn provision(&mut self, from: stage::NodeId, provision: stage::ProvisionStage) { self.inbound_edge = Some(plan::EdgeId(provision.inbound.edge_id.0)); self.outbound_edge = Some(plan::EdgeId(provision.outbound.edge_id.0)); - self.outbound_object_allocator = Some(edge_actor::ObjectIdAllocator::new( - edge_actor::EdgeId(provision.outbound.edge_id.0), - )); + self.outbound_object_allocator = Some(ObjectIdAllocator::new()); self.controller .observe(stage::StageEvent::ProvisionStage { from, provision }); } @@ -54,9 +52,7 @@ impl MockNode { pub fn provision_from_wrong_orchestrator(&mut self, provision: stage::ProvisionStage) { self.inbound_edge = Some(plan::EdgeId(provision.inbound.edge_id.0)); self.outbound_edge = Some(plan::EdgeId(provision.outbound.edge_id.0)); - self.outbound_object_allocator = Some(edge_actor::ObjectIdAllocator::new( - edge_actor::EdgeId(provision.outbound.edge_id.0), - )); + self.outbound_object_allocator = Some(ObjectIdAllocator::new()); self.controller.observe(stage::StageEvent::ProvisionStage { from: stage::NodeId(provision.authorized_orchestrator.0 + 1), provision, @@ -97,7 +93,7 @@ impl MockNode { stage::StageCommand::ExecuteStep(step) => Some(step.clone()), _ => None, })?; - let output_object_id = self.outbound_object_allocator.as_mut()?.alloc().object_id.0; + let output_object_id = self.outbound_object_allocator.as_mut()?.alloc().0; let produced = self.worker.execute( &step, outbound_edge, diff --git a/crates/data-plane/Cargo.toml b/crates/data-plane/Cargo.toml index 078fa83..412d5c3 100644 --- a/crates/data-plane/Cargo.toml +++ b/crates/data-plane/Cargo.toml @@ -10,5 +10,8 @@ telemetry = { path = "../telemetry" } serde = { version = "1", features = ["derive"] } swactor = { path = "../.." } + +[dev-dependencies] +parking_lot = "0.12" [target.'cfg(target_os = "linux")'.dependencies] libc = "0.2" diff --git a/crates/data-plane/src/actor.rs b/crates/data-plane/src/actor.rs deleted file mode 100644 index 0ed7efa..0000000 --- a/crates/data-plane/src/actor.rs +++ /dev/null @@ -1,769 +0,0 @@ -//! Swactor-facing data-plane actor API. -//! -//! The actor owns lifecycle decisions for logical wire edge endpoints and their -//! bound node-local worker rings. Runtime-specific actors execute the effect -//! messages and report observations back. - -use std::collections::{BTreeMap, BTreeSet}; - -use swactor::actor::{ActorAddress, ActorInterface}; -use swactor::runtime::Ctx; - -use crate::edge_lifecycle as lifecycle; -use crate::object_record; - -pub use lifecycle::{ - DType, EdgeFaultReason, EdgeId, LeaseRequestId, NodeId, ObjectKind, ObjectSpec, - QuiescenceProof, RingDirection, RingId, RingLayout, RingLeaseRejection, RingSpec, RunId, -}; - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct PortId(pub String); - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum EdgeKind { - TokenIn, - Activation, - TokenOut, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum EdgeEndpointDirection { - Inbound, - Outbound, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct PeerEndpoint(pub String); - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum TransportBinding { - Remote { endpoint: Option }, - LocalOnly, -} - -impl TransportBinding { - fn requires_transport(&self) -> bool { - matches!(self, Self::Remote { .. }) - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum WorkerRingBinding { - Required, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct DataPlaneRunConfig { - pub run_id: RunId, - pub local_node_id: NodeId, - pub arena_actor: ActorAddress, - pub worker_actor: ActorAddress, - pub transport_actor: ActorAddress, - pub report_sink: ActorAddress, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct WireEdgeEndpoint { - pub run_id: RunId, - pub edge_id: EdgeId, - pub direction: EdgeEndpointDirection, - pub kind: EdgeKind, - pub local_node_id: NodeId, - pub peer_node_id: Option, - pub peer_endpoint: Option, - pub local_role_port: PortId, - pub object_spec: ObjectSpec, - pub ring_spec: RingSpec, - pub object_record_spec: object_record::ObjectSpec, - pub transport: TransportBinding, - pub worker_ring: WorkerRingBinding, -} - -impl WireEdgeEndpoint { - fn provision_event(&self) -> Option { - match self.direction { - EdgeEndpointDirection::Inbound => { - Some(lifecycle::EdgeEvent::ProvisionRx(lifecycle::ProvisionRx { - run_id: self.run_id, - edge_id: self.edge_id, - local_node_id: self.local_node_id, - object_spec: self.object_spec, - ring_spec: self.ring_spec, - })) - } - EdgeEndpointDirection::Outbound => { - Some(lifecycle::EdgeEvent::ProvisionTx(lifecycle::ProvisionTx { - run_id: self.run_id, - edge_id: self.edge_id, - local_node_id: self.local_node_id, - consumer_node_id: self.peer_node_id?, - object_spec: self.object_spec, - ring_spec: self.ring_spec, - })) - } - } - } -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum DataPlaneNodeMsg { - ConfigureRun(DataPlaneRunConfig), - ProvisionWireEdgeEndpoint(WireEdgeEndpoint), - Arena(ArenaObservation), - Worker(WorkerObservation), - Transport(TransportObservation), - StopEdge { edge_id: EdgeId }, - StopRun { run_id: RunId }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum ArenaObservation { - RingLeased { - request_id: LeaseRequestId, - ring_id: RingId, - layout: RingLayout, - }, - RingLeaseRejected { - request_id: LeaseRequestId, - reason: RingLeaseRejection, - }, - RingReleased { - ring_id: RingId, - }, - RingReleaseRejected { - ring_id: RingId, - }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum WorkerObservation { - WorkerReady, - RingInstalled { - edge_id: EdgeId, - ring_id: RingId, - }, - RingFaulted { - edge_id: EdgeId, - ring_id: RingId, - reason: lifecycle::RingFaultReason, - }, - RingQuiesced { - ring_id: RingId, - }, - QuiescenceProven { - ring_id: RingId, - }, - ObjectLoaded { - edge_id: EdgeId, - ring_id: RingId, - object_id: object_record::ObjectId, - sequence: u64, - extent: u64, - handle: DeviceHandle, - }, - ObjectProduced { - edge_id: EdgeId, - ring_id: RingId, - object_id: object_record::ObjectId, - sequence: u64, - extent: u64, - }, - ObjectFailed { - edge_id: EdgeId, - ring_id: RingId, - reason: object_record::ObjectFailureReason, - }, - WorkerFaulted, - WorkerStopped, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum TransportObservation { - EdgeReady { - edge_id: EdgeId, - }, - BytesAvailable { - edge_id: EdgeId, - stream_id: StreamId, - byte_count: u64, - }, - StreamClosed { - edge_id: EdgeId, - }, - StreamFaulted { - edge_id: EdgeId, - reason: lifecycle::StreamFaultReason, - }, - PumpStopped { - edge_id: EdgeId, - ring_id: RingId, - }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum DataPlaneArenaMsg { - LeaseRing { - request_id: LeaseRequestId, - edge_id: EdgeId, - direction: RingDirection, - ring_spec: RingSpec, - }, - CancelQueuedLease { - request_id: LeaseRequestId, - edge_id: EdgeId, - }, - ReleaseRing { - ring_id: RingId, - proof: QuiescenceProof, - }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum DataPlaneWorkerMsg { - InstallRing { - edge_id: EdgeId, - ring_id: RingId, - direction: RingDirection, - layout: RingLayout, - object_spec: ObjectSpec, - ring_spec: RingSpec, - role_port: PortId, - }, - UninstallRing { - edge_id: EdgeId, - ring_id: RingId, - }, - NotifyRingReadable { - ring_id: RingId, - }, - NotifyRingWritable { - ring_id: RingId, - }, - LoadObjectFromRing { - edge_id: EdgeId, - ring_id: RingId, - object_id: object_record::ObjectId, - sequence: u64, - extent: u64, - }, - ExecuteStep { - edge_id: EdgeId, - sequence: u64, - }, - ReleaseDeviceObject { - handle: DeviceHandle, - }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum DataPlaneTransportMsg { - EstablishSend { - edge_id: EdgeId, - consumer_node_id: NodeId, - ring_id: RingId, - layout: RingLayout, - }, - EstablishRecv { - edge_id: EdgeId, - ring_id: RingId, - layout: RingLayout, - }, - WriteWireObject { - edge_id: EdgeId, - object_id: object_record::ObjectId, - sequence: u64, - extent: u64, - }, - StopWirePump { - edge_id: EdgeId, - ring_id: RingId, - }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum DataPlaneReportMsg { - InboundEdgeReady { - edge_id: EdgeId, - }, - OutboundEdgeReady { - edge_id: EdgeId, - }, - ObjectLoaded { - edge_id: EdgeId, - object_id: object_record::ObjectId, - sequence: u64, - extent: u64, - handle: DeviceHandle, - }, - ObjectProduced { - edge_id: EdgeId, - object_id: object_record::ObjectId, - sequence: u64, - extent: u64, - }, - EdgeFaulted { - edge_id: EdgeId, - reason: EdgeFaultReason, - }, - EdgeStopped { - edge_id: EdgeId, - }, - LocalEdgesStopped { - run_id: RunId, - }, - WorkerDataPlaneFaulted { - reason: WorkerDataPlaneFaultReason, - }, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum WorkerDataPlaneFaultReason { - WorkerFaulted, - WorkerStopped, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct WorkerGeneration(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct DeviceHandle { - pub generation: WorkerGeneration, - pub id: u64, -} - -impl DeviceHandle { - pub const fn new(generation: WorkerGeneration, id: u64) -> Self { - Self { generation, id } - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct StreamId(pub u64); - -#[derive(Clone, Debug, PartialEq, Eq)] -struct EdgeBinding { - direction: EdgeEndpointDirection, - role_port_index: usize, - transport: TransportBinding, -} - -pub struct DataPlaneNodeActor { - run: Option, - lifecycle: lifecycle::EdgeEstablisher, - command_cursor: usize, - event_cursor: usize, - bindings: BTreeMap, - stopped_edges: BTreeSet, - role_ports: Vec, - rings: BTreeMap, - pending_reports: Vec, - stopping_run: Option, - local_edges_stopped_reported: bool, -} - -impl DataPlaneNodeActor { - pub fn new(local_node_id: NodeId) -> Self { - Self { - run: None, - lifecycle: lifecycle::EdgeEstablisher::new(local_node_id), - command_cursor: 0, - event_cursor: 0, - bindings: BTreeMap::new(), - stopped_edges: BTreeSet::new(), - role_ports: Vec::new(), - rings: BTreeMap::new(), - pending_reports: Vec::new(), - stopping_run: None, - local_edges_stopped_reported: false, - } - } - - fn configure(&mut self, config: DataPlaneRunConfig) { - self.lifecycle = lifecycle::EdgeEstablisher::new(config.local_node_id); - self.command_cursor = 0; - self.event_cursor = 0; - self.bindings.clear(); - self.role_ports.clear(); - self.stopped_edges.clear(); - self.rings.clear(); - self.pending_reports.clear(); - self.stopping_run = None; - self.local_edges_stopped_reported = false; - self.run = Some(config); - } - - fn observe(&mut self, msg: DataPlaneNodeMsg) { - match msg { - DataPlaneNodeMsg::ConfigureRun(config) => self.configure(config), - DataPlaneNodeMsg::ProvisionWireEdgeEndpoint(endpoint) => { - if let Some(event) = endpoint.provision_event() { - let role_port_index = self.role_ports.len(); - self.role_ports.push(endpoint.local_role_port.clone()); - self.bindings.insert( - endpoint.edge_id, - EdgeBinding { - direction: endpoint.direction, - role_port_index, - transport: endpoint.transport.clone(), - }, - ); - self.lifecycle.observe(event); - } - } - DataPlaneNodeMsg::Arena(observation) => match observation { - ArenaObservation::RingLeased { - request_id, - ring_id, - layout, - } => self.lifecycle.observe(lifecycle::EdgeEvent::RingLeased { - request_id, - ring_id, - layout, - }), - ArenaObservation::RingLeaseRejected { request_id, reason } => self - .lifecycle - .observe(lifecycle::EdgeEvent::RingLeaseRejected { request_id, reason }), - ArenaObservation::RingReleased { .. } - | ArenaObservation::RingReleaseRejected { .. } => {} - }, - DataPlaneNodeMsg::Worker(observation) => match observation { - WorkerObservation::WorkerReady => {} - WorkerObservation::RingInstalled { edge_id, ring_id } => { - self.rings.insert(ring_id, edge_id); - self.lifecycle - .observe(lifecycle::EdgeEvent::RingInstalled { edge_id, ring_id }); - } - WorkerObservation::RingFaulted { - edge_id, - ring_id, - reason, - } => self.lifecycle.observe(lifecycle::EdgeEvent::RingFault { - edge_id, - ring_id, - reason, - }), - WorkerObservation::RingQuiesced { ring_id } => self - .lifecycle - .observe(lifecycle::EdgeEvent::RingQuiesced { ring_id }), - WorkerObservation::QuiescenceProven { ring_id } => self - .lifecycle - .observe(lifecycle::EdgeEvent::QuiescenceProven { ring_id }), - WorkerObservation::ObjectLoaded { - edge_id, - object_id, - sequence, - extent, - handle, - .. - } => self.report_object_loaded(edge_id, object_id, sequence, extent, handle), - WorkerObservation::ObjectProduced { - edge_id, - object_id, - sequence, - extent, - .. - } => self.report_object_produced(edge_id, object_id, sequence, extent), - WorkerObservation::ObjectFailed { - edge_id, ring_id, .. - } => self.lifecycle.observe(lifecycle::EdgeEvent::RingFault { - edge_id, - ring_id, - reason: lifecycle::RingFaultReason::WorkerRejectedRing, - }), - WorkerObservation::WorkerFaulted => { - self.report_worker_fault(WorkerDataPlaneFaultReason::WorkerFaulted) - } - WorkerObservation::WorkerStopped => { - self.report_worker_fault(WorkerDataPlaneFaultReason::WorkerStopped) - } - }, - DataPlaneNodeMsg::Transport(observation) => match observation { - TransportObservation::EdgeReady { edge_id } => self - .lifecycle - .observe(lifecycle::EdgeEvent::DriverEdgeReady { edge_id }), - TransportObservation::BytesAvailable { .. } - | TransportObservation::StreamClosed { .. } => {} - TransportObservation::StreamFaulted { edge_id, reason } => self - .lifecycle - .observe(lifecycle::EdgeEvent::StreamFault { edge_id, reason }), - TransportObservation::PumpStopped { edge_id, ring_id } => self - .lifecycle - .observe(lifecycle::EdgeEvent::PumpStopped { edge_id, ring_id }), - }, - DataPlaneNodeMsg::StopEdge { edge_id } => self - .lifecycle - .observe(lifecycle::EdgeEvent::StopEdge { edge_id }), - DataPlaneNodeMsg::StopRun { run_id } => { - self.stopping_run = Some(run_id); - self.local_edges_stopped_reported = false; - let edge_ids = self.bindings.keys().copied().collect::>(); - for edge_id in edge_ids { - self.lifecycle - .observe(lifecycle::EdgeEvent::StopEdge { edge_id }); - } - self.maybe_report_local_edges_stopped(); - } - } - } - - fn ring_for_edge(&self, edge_id: EdgeId) -> Option { - self.rings - .iter() - .find_map(|(ring_id, seen_edge)| (*seen_edge == edge_id).then_some(*ring_id)) - } - - fn role_port_for(&self, edge_id: EdgeId) -> PortId { - self.bindings - .get(&edge_id) - .and_then(|binding| self.role_ports.get(binding.role_port_index)) - .cloned() - .unwrap_or_else(|| PortId(String::new())) - } - - fn report_object_loaded( - &mut self, - edge_id: EdgeId, - object_id: object_record::ObjectId, - sequence: u64, - extent: u64, - handle: DeviceHandle, - ) { - self.pending_reports.push(DataPlaneReportMsg::ObjectLoaded { - edge_id, - object_id, - sequence, - extent, - handle, - }); - } - - fn report_object_produced( - &mut self, - edge_id: EdgeId, - object_id: object_record::ObjectId, - sequence: u64, - extent: u64, - ) { - self.pending_reports - .push(DataPlaneReportMsg::ObjectProduced { - edge_id, - object_id, - sequence, - extent, - }); - } - - fn report_worker_fault(&mut self, reason: WorkerDataPlaneFaultReason) { - self.pending_reports - .push(DataPlaneReportMsg::WorkerDataPlaneFaulted { reason }); - } - - fn report_local_edges_stopped(&mut self, run_id: RunId) { - self.pending_reports - .push(DataPlaneReportMsg::LocalEdgesStopped { run_id }); - } - - fn maybe_report_local_edges_stopped(&mut self) { - let Some(run_id) = self.stopping_run else { - return; - }; - if self.local_edges_stopped_reported { - return; - } - if self - .bindings - .keys() - .all(|edge_id| self.stopped_edges.contains(edge_id)) - { - self.local_edges_stopped_reported = true; - self.report_local_edges_stopped(run_id); - } - } - - fn drain(&mut self, ctx: &Ctx) { - let Some(run) = self.run.clone() else { - return; - }; - - loop { - let mut progressed = false; - let mut immediate_ready = Vec::new(); - - while self.command_cursor < self.lifecycle.commands().len() { - let command = self.lifecycle.commands()[self.command_cursor].clone(); - self.command_cursor += 1; - progressed = true; - match command { - lifecycle::EdgeCommand::LeaseRing { - request_id, - edge_id, - direction, - ring_spec, - } => { - let _ = ctx.send( - run.arena_actor, - DataPlaneArenaMsg::LeaseRing { - request_id, - edge_id, - direction, - ring_spec, - }, - ); - } - lifecycle::EdgeCommand::CancelQueuedLease { - request_id, - edge_id, - } => { - let _ = ctx.send( - run.arena_actor, - DataPlaneArenaMsg::CancelQueuedLease { - request_id, - edge_id, - }, - ); - } - lifecycle::EdgeCommand::InstallWorkerRing { - edge_id, - ring_id, - direction, - layout, - object_spec, - ring_spec, - } => { - let _ = ctx.send( - run.worker_actor, - DataPlaneWorkerMsg::InstallRing { - edge_id, - ring_id, - direction, - layout, - object_spec, - ring_spec, - role_port: self.role_port_for(edge_id), - }, - ); - } - lifecycle::EdgeCommand::UninstallWorkerRing { edge_id, ring_id } => { - let _ = ctx.send( - run.worker_actor, - DataPlaneWorkerMsg::UninstallRing { edge_id, ring_id }, - ); - } - lifecycle::EdgeCommand::EstablishSend { - edge_id, - consumer_node_id, - layout, - } => { - let ring_id = self.ring_for_edge(edge_id).unwrap_or(RingId(0)); - if self - .bindings - .get(&edge_id) - .is_some_and(|binding| binding.transport.requires_transport()) - { - let _ = ctx.send( - run.transport_actor, - DataPlaneTransportMsg::EstablishSend { - edge_id, - consumer_node_id, - ring_id, - layout, - }, - ); - } else { - immediate_ready.push(edge_id); - } - } - lifecycle::EdgeCommand::EstablishRecv { edge_id, layout } => { - let ring_id = self.ring_for_edge(edge_id).unwrap_or(RingId(0)); - if self - .bindings - .get(&edge_id) - .is_some_and(|binding| binding.transport.requires_transport()) - { - let _ = ctx.send( - run.transport_actor, - DataPlaneTransportMsg::EstablishRecv { - edge_id, - ring_id, - layout, - }, - ); - } else { - immediate_ready.push(edge_id); - } - } - lifecycle::EdgeCommand::StopPump { edge_id, ring_id } => { - let _ = ctx.send( - run.transport_actor, - DataPlaneTransportMsg::StopWirePump { edge_id, ring_id }, - ); - } - lifecycle::EdgeCommand::ReleaseArenaLease { ring_id, proof } => { - let _ = ctx.send( - run.arena_actor, - DataPlaneArenaMsg::ReleaseRing { ring_id, proof }, - ); - } - } - } - - for edge_id in immediate_ready { - self.lifecycle - .observe(lifecycle::EdgeEvent::DriverEdgeReady { edge_id }); - } - - while self.event_cursor < self.lifecycle.events().len() { - let event = self.lifecycle.events()[self.event_cursor].clone(); - self.event_cursor += 1; - progressed = true; - match event { - lifecycle::EdgeLifecycleEvent::EdgeReady { edge_id, .. } => { - let report = - match self.bindings.get(&edge_id).map(|binding| binding.direction) { - Some(EdgeEndpointDirection::Inbound) => { - DataPlaneReportMsg::InboundEdgeReady { edge_id } - } - Some(EdgeEndpointDirection::Outbound) => { - DataPlaneReportMsg::OutboundEdgeReady { edge_id } - } - None => continue, - }; - let _ = ctx.send(run.report_sink, report); - } - lifecycle::EdgeLifecycleEvent::EdgeFaulted { edge_id, reason } => { - let _ = ctx.send( - run.report_sink, - DataPlaneReportMsg::EdgeFaulted { edge_id, reason }, - ); - } - lifecycle::EdgeLifecycleEvent::EdgeStopped { edge_id } => { - self.stopped_edges.insert(edge_id); - let _ = - ctx.send(run.report_sink, DataPlaneReportMsg::EdgeStopped { edge_id }); - self.maybe_report_local_edges_stopped(); - } - } - } - - for report in self.pending_reports.drain(..) { - progressed = true; - let _ = ctx.send(run.report_sink, report); - } - - if !progressed { - break; - } - } - } -} - -impl ActorInterface for DataPlaneNodeActor { - type Incoming = DataPlaneNodeMsg; - type Response = (); - - fn handle(&mut self, ctx: &Ctx, msg: Self::Incoming) { - self.observe(msg); - self.drain(ctx); - } -} diff --git a/crates/data-plane/src/arena.rs b/crates/data-plane/src/arena.rs index 4db48e3..ef8e3a7 100644 --- a/crates/data-plane/src/arena.rs +++ b/crates/data-plane/src/arena.rs @@ -58,14 +58,7 @@ impl From for ArenaSample { } } -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct NodeId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct LeaseRequestId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct RingId(pub u64); +pub use crate::ids::{LeaseRequestId, NodeId, RingId}; #[derive(Clone, Debug, PartialEq, Eq)] pub struct ArenaConfig { diff --git a/crates/data-plane/src/edge_actor.rs b/crates/data-plane/src/edge_actor.rs deleted file mode 100644 index 3f1b182..0000000 --- a/crates/data-plane/src/edge_actor.rs +++ /dev/null @@ -1,332 +0,0 @@ -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct EdgeId(pub u64); - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct PortId(pub String); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ObjectId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ObjectKey { - pub edge_id: EdgeId, - pub object_id: ObjectId, -} - -impl ObjectKey { - pub fn new(edge_id: EdgeId, object_id: ObjectId) -> Self { - Self { edge_id, object_id } - } -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ObjectIdAllocator { - edge_id: EdgeId, - next: u64, -} - -impl ObjectIdAllocator { - pub fn new(edge_id: EdgeId) -> Self { - Self { edge_id, next: 1 } - } - - pub fn alloc(&mut self) -> ObjectKey { - let object_key = ObjectKey::new(self.edge_id, ObjectId(self.next)); - self.next += 1; - object_key - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct OpaqueHandle(pub u64); - -impl OpaqueHandle { - pub fn new(id: u64) -> Self { - Self(id) - } -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct TxConfig { - pub edge_id: EdgeId, - pub role_port: PortId, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RxConfig { - pub edge_id: EdgeId, - pub role_port: PortId, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum ActorFaultReason { - MismatchedEdgeId, - StreamFault, - ObjectFault, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum ActorMessage { - Lifecycle { - edge_id: EdgeId, - ready: bool, - }, - ObjectIdentity { - edge_id: EdgeId, - object_id: ObjectId, - sequence: u64, - }, - OpaqueHandle { - edge_id: EdgeId, - object_id: ObjectId, - sequence: u64, - handle: OpaqueHandle, - }, - CoarseFault { - edge_id: EdgeId, - reason: ActorFaultReason, - }, -} - -impl ActorMessage { - pub fn edge_id(&self) -> EdgeId { - match self { - ActorMessage::Lifecycle { edge_id, .. } - | ActorMessage::ObjectIdentity { edge_id, .. } - | ActorMessage::OpaqueHandle { edge_id, .. } - | ActorMessage::CoarseFault { edge_id, .. } => *edge_id, - } - } -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum TxEvent { - EdgeReady { - edge_id: EdgeId, - }, - ObjectProduced { - edge_id: EdgeId, - object_id: ObjectId, - sequence: u64, - }, - StreamFault { - edge_id: EdgeId, - }, - ObjectFailed { - edge_id: EdgeId, - }, - StopEdge { - edge_id: EdgeId, - }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum RxEvent { - EdgeReady { - edge_id: EdgeId, - }, - ObjectLoaded { - edge_id: EdgeId, - object_id: ObjectId, - sequence: u64, - handle: OpaqueHandle, - }, - StreamFault { - edge_id: EdgeId, - }, - ObjectFailed { - edge_id: EdgeId, - }, - StopEdge { - edge_id: EdgeId, - }, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum ActorState { - Provisioning, - Ready, - Faulted, - Stopped, -} - -pub struct TxEdgeActor { - config: TxConfig, - state: ActorState, - messages: Vec, -} - -impl TxEdgeActor { - pub fn new(config: TxConfig) -> Self { - Self { - config, - state: ActorState::Provisioning, - messages: Vec::new(), - } - } - - pub fn observe(&mut self, event: TxEvent) { - match event { - TxEvent::EdgeReady { edge_id } => { - if !self.check_edge(edge_id) { - return; - } - if self.state == ActorState::Provisioning { - self.state = ActorState::Ready; - self.messages.push(ActorMessage::Lifecycle { - edge_id, - ready: true, - }); - } - } - TxEvent::ObjectProduced { - edge_id, - object_id, - sequence, - } => { - if self.state == ActorState::Ready && self.check_edge(edge_id) { - self.messages.push(ActorMessage::ObjectIdentity { - edge_id, - object_id, - sequence, - }); - } - } - TxEvent::StreamFault { edge_id } => { - self.fault_if_edge(edge_id, ActorFaultReason::StreamFault) - } - TxEvent::ObjectFailed { edge_id } => { - self.fault_if_edge(edge_id, ActorFaultReason::ObjectFault) - } - TxEvent::StopEdge { edge_id } => { - if edge_id == self.config.edge_id { - self.state = ActorState::Stopped; - } else { - self.mismatched(); - } - } - } - } - - pub fn messages(&self) -> &[ActorMessage] { - &self.messages - } - - fn check_edge(&mut self, edge_id: EdgeId) -> bool { - if edge_id == self.config.edge_id { - true - } else { - self.mismatched(); - false - } - } - - fn fault_if_edge(&mut self, edge_id: EdgeId, reason: ActorFaultReason) { - if self.check_edge(edge_id) && self.state != ActorState::Stopped { - self.state = ActorState::Faulted; - self.messages - .push(ActorMessage::CoarseFault { edge_id, reason }); - } - } - - fn mismatched(&mut self) { - self.state = ActorState::Faulted; - self.messages.push(ActorMessage::CoarseFault { - edge_id: self.config.edge_id, - reason: ActorFaultReason::MismatchedEdgeId, - }); - } -} - -pub struct RxEdgeActor { - config: RxConfig, - state: ActorState, - messages: Vec, -} - -impl RxEdgeActor { - pub fn new(config: RxConfig) -> Self { - Self { - config, - state: ActorState::Provisioning, - messages: Vec::new(), - } - } - - pub fn observe(&mut self, event: RxEvent) { - match event { - RxEvent::EdgeReady { edge_id } => { - if !self.check_edge(edge_id) { - return; - } - if self.state == ActorState::Provisioning { - self.state = ActorState::Ready; - self.messages.push(ActorMessage::Lifecycle { - edge_id, - ready: true, - }); - } - } - RxEvent::ObjectLoaded { - edge_id, - object_id, - sequence, - handle, - } => { - if self.state == ActorState::Ready && self.check_edge(edge_id) { - self.messages.push(ActorMessage::OpaqueHandle { - edge_id, - object_id, - sequence, - handle, - }); - } - } - RxEvent::StreamFault { edge_id } => { - self.fault_if_edge(edge_id, ActorFaultReason::StreamFault) - } - RxEvent::ObjectFailed { edge_id } => { - self.fault_if_edge(edge_id, ActorFaultReason::ObjectFault) - } - RxEvent::StopEdge { edge_id } => { - if edge_id == self.config.edge_id { - self.state = ActorState::Stopped; - } else { - self.mismatched(); - } - } - } - } - - pub fn messages(&self) -> &[ActorMessage] { - &self.messages - } - - fn check_edge(&mut self, edge_id: EdgeId) -> bool { - if edge_id == self.config.edge_id { - true - } else { - self.mismatched(); - false - } - } - - fn fault_if_edge(&mut self, edge_id: EdgeId, reason: ActorFaultReason) { - if self.check_edge(edge_id) && self.state != ActorState::Stopped { - self.state = ActorState::Faulted; - self.messages - .push(ActorMessage::CoarseFault { edge_id, reason }); - } - } - - fn mismatched(&mut self) { - self.state = ActorState::Faulted; - self.messages.push(ActorMessage::CoarseFault { - edge_id: self.config.edge_id, - reason: ActorFaultReason::MismatchedEdgeId, - }); - } -} - -pub type TxActorHarness = TxEdgeActor; -pub type RxActorHarness = RxEdgeActor; diff --git a/crates/data-plane/src/edge_lifecycle.rs b/crates/data-plane/src/edge_lifecycle.rs index 29209c4..9025dd5 100644 --- a/crates/data-plane/src/edge_lifecycle.rs +++ b/crates/data-plane/src/edge_lifecycle.rs @@ -1,22 +1,6 @@ use std::collections::BTreeMap; -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct RunId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct NodeId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct EdgeId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct LeaseRequestId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct RingId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ActorAddress(pub u64); +pub use crate::ids::{ActorAddress, EdgeId, LeaseRequestId, NodeId, RingId, RunId}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum RingDirection { diff --git a/crates/data-plane/src/edge_runtime.rs b/crates/data-plane/src/edge_runtime.rs new file mode 100644 index 0000000..2cd66e2 --- /dev/null +++ b/crates/data-plane/src/edge_runtime.rs @@ -0,0 +1,703 @@ +//! Edge runtime: the composition engine for wire edges. +//! +//! [`EdgeRuntime`] merges what used to be three separately-owned pieces: +//! +//! - the pure edge lifecycle protocol ([`crate::edge_lifecycle`]), +//! - the wire bookkeeping formerly living in iroh-driver's `driver_pumps` +//! (which edges have send/recv pumps, mapping inbound streams to edges), +//! - and the per-tick application glue (draining transport events, buffering +//! ingress streams, parsing object records into arena rings, executing +//! lifecycle commands). +//! +//! The runtime is driven synchronously from [`EdgeRuntime::poll`]. All +//! effects go through narrow ports: the byte transport is any +//! [`EdgeTransport`](crate::edge_wire::EdgeTransport) (iroh-driver provides +//! one), worker ring effects go through [`WorkerPort`] (the application +//! implements it over its GPU worker), and the arena is the in-crate +//! [`ArenaManager`]. Observations — telemetry-shaped facts and lifecycle +//! transitions — accumulate and are taken by the application after each +//! poll. + +use std::collections::BTreeMap; +use std::time::Instant; + +use crate::arena::{ + ArenaEvent, ArenaManager, ArenaRequest, LeaseRing, QuiescenceProof as ArenaQuiescenceProof, + RingLayout as ArenaRingLayout, RingSpec as ArenaRingSpec, +}; +use crate::edge_lifecycle::{ + EdgeCommand, EdgeEstablisher, EdgeEvent, EdgeFaultReason, EdgeLifecycleEvent, NodeId, + ObjectSpec, ProvisionRx, ProvisionTx, RingDirection, RingLeaseRejection, + StreamFaultReason, +}; +use crate::edge_wire::EdgeTransport; +use crate::ids::{EdgeId, LeaseRequestId, RingId, StreamId}; +use crate::object_record::{self, ObjectRecord}; + +/// Worker effects the runtime drives during edge provisioning and ingress. +pub trait WorkerPort { + /// Install a leased ring into the worker for `edge_id`; `direction` + /// decides input/output placement. `layout` is the arena lease layout. + fn install_ring( + &mut self, + edge_id: EdgeId, + ring_id: RingId, + direction: RingDirection, + layout: &ArenaRingLayout, + object_spec: &ObjectSpec, + ) -> Result<(), String>; + + /// Uninstall (quiesce) a worker ring. + fn uninstall_ring(&mut self, ring_id: RingId) -> Result<(), String>; + + /// Load one complete ingress object from `ring_id` and return its device + /// handle. `spec` is the object parse spec used on the wire. + fn load_object( + &mut self, + edge_id: EdgeId, + ring_id: RingId, + record: &ObjectRecord, + spec: &object_record::ObjectSpec, + ) -> Result; +} + +/// A worker-loaded ingress object: identity plus its device handle. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct LoadedObject { + pub object_id: u64, + pub sequence: u64, + pub handle_generation: u64, + pub handle_id: u64, +} + +/// One structured observation of runtime progress. The application maps +/// these to telemetry and node-agent messages; the runtime emits them and +/// forgets them. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Observation { + StreamArrived { + edge_id: EdgeId, + stream_id: StreamId, + }, + BytesRead { + edge_id: EdgeId, + stream_id: StreamId, + byte_count: usize, + }, + /// One complete object record was parsed from an ingress stream and + /// written into the edge's ring. + IngressRingWrite { + edge_id: EdgeId, + ring_id: RingId, + stream_id: StreamId, + object_id: u64, + sequence: u64, + extent: u64, + begin_sequence: bool, + end_of_sequence: bool, + record_bytes: usize, + buffered_bytes: usize, + write_ms: u64, + }, + /// One ingress object was loaded by the worker and now has a handle. + ObjectLoaded { + edge_id: EdgeId, + ring_id: RingId, + stream_id: StreamId, + object: LoadedObject, + load_ms: u64, + }, + /// An ingress object failed to parse or load. + ObjectFailed { + edge_id: EdgeId, + object_id: Option, + }, + EdgeReady { + edge_id: EdgeId, + direction: RingDirection, + }, + EdgeFaulted { + edge_id: EdgeId, + reason: EdgeFaultReason, + }, + EdgeStopped { + edge_id: EdgeId, + }, +} + +struct InboundEdge { + edge_id: EdgeId, + ring_id: Option, + /// Wire parse spec for objects arriving on this edge. + parse_spec: object_record::ObjectSpec, + buffers: BTreeMap>, +} + +struct OutboundEdge { + edge_id: EdgeId, + peer: P, + ring_id: Option, + writer: Option, + next_object_id: u64, +} + +/// One complete object record drained from an ingress stream buffer. +struct IngressRecord { + bytes: Vec, + record: ObjectRecord, +} + +pub struct EdgeRuntime { + establisher: EdgeEstablisher, + command_cursor: usize, + lifecycle_cursor: usize, + /// Wire bookkeeping (formerly iroh-driver `driver_pumps::Driver`). + send_rings: BTreeMap, + recv_specs: BTreeMap, + pending_streams: BTreeMap, + recv_rings: BTreeMap, + inbound: Option, + outbound: Option>, + object_handles: BTreeMap<(EdgeId, u64), LoadedObject>, + observations: Vec, +} + +impl EdgeRuntime { + pub fn new(local_node_id: NodeId) -> Self { + Self { + establisher: EdgeEstablisher::new(local_node_id), + command_cursor: 0, + lifecycle_cursor: 0, + send_rings: BTreeMap::new(), + recv_specs: BTreeMap::new(), + pending_streams: BTreeMap::new(), + recv_rings: BTreeMap::new(), + inbound: None, + outbound: None, + object_handles: BTreeMap::new(), + observations: Vec::new(), + } + } + + /// Provision the node's (single) inbound edge and begin its lifecycle. + /// + /// `parse_spec` is the object-record parse spec for this edge's wire + /// format; the application derives it from its plan. + pub fn establish_inbound( + &mut self, + provision: ProvisionRx, + parse_spec: object_record::ObjectSpec, + ) { + self.inbound = Some(InboundEdge { + edge_id: provision.edge_id, + ring_id: None, + parse_spec, + buffers: BTreeMap::new(), + }); + self.establisher.observe(EdgeEvent::ProvisionRx(provision)); + } + + /// Provision the node's (single) outbound edge toward `peer` and begin + /// its lifecycle. + pub fn establish_outbound(&mut self, provision: ProvisionTx, peer: T::PeerAddr) { + self.outbound = Some(OutboundEdge { + edge_id: provision.edge_id, + peer, + ring_id: None, + writer: None, + next_object_id: 1, + }); + self.establisher.observe(EdgeEvent::ProvisionTx(provision)); + } + + /// One tick: drain transport events, ingest inbound bytes, and run the + /// lifecycle/workflow fixpoint to quiescence. + /// + /// Returns `Err` on any fault the previous composition treated as fatal + /// (object parse/load failure, worker effect failure, writer failure, or + /// an edge fault). Observations emitted up to and including the fault + /// remain retrievable via [`Self::take_observations`]. + pub fn poll( + &mut self, + transport: &mut T, + arena: &mut ArenaManager, + worker: &mut dyn WorkerPort, + ) -> Result<(), String> { + // Run the workflow first: provisions established since the last + // tick (or earlier in this tick) lease and install their rings + // before ingress bytes are parsed against them. + self.drive(transport, arena, worker)?; + for event in transport.drain_events() { + match event { + crate::edge_wire::WireEvent::StreamArrived { edge_id, stream_id } => { + self.incoming_stream(edge_id, stream_id); + self.observations.push(Observation::StreamArrived { edge_id, stream_id }); + } + crate::edge_wire::WireEvent::BytesRead { edge_id, stream_id, bytes } => { + self.observations.push(Observation::BytesRead { + edge_id, + stream_id, + byte_count: bytes.len(), + }); + self.ingress_bytes(arena, worker, edge_id, stream_id, bytes)?; + } + crate::edge_wire::WireEvent::StreamEnded { edge_id, stream_id } => { + if let Some(inbound) = self + .inbound + .as_mut() + .filter(|edge| edge.edge_id == edge_id) + { + inbound.buffers.remove(&stream_id); + } + } + crate::edge_wire::WireEvent::StreamFault { + edge_id: Some(edge_id), + .. + } => { + self.read_error(edge_id); + } + crate::edge_wire::WireEvent::StreamFault { edge_id: None, .. } => {} + } + } + self.drive(transport, arena, worker) + } + + /// Drain all observations accumulated since the last call. + pub fn take_observations(&mut self) -> Vec { + std::mem::take(&mut self.observations) + } + + /// The device handle for one loaded ingress object. + pub fn loaded_object(&self, edge_id: EdgeId, object_id: u64) -> Option<&LoadedObject> { + self.object_handles.get(&(edge_id, object_id)) + } + + /// The outbound edge's leased ring, once installed. + pub fn outbound_ring_id(&self) -> Option { + self.outbound.as_ref().and_then(|edge| edge.ring_id) + } + + /// The outbound edge's byte writer, once the send pump is established. + pub fn outbound_writer(&self) -> Option<&T::Writer> { + self.outbound.as_ref().and_then(|edge| edge.writer.as_ref()) + } + + /// Allocate the next output object id for the outbound edge. + pub fn alloc_output_object_id(&mut self) -> Result { + self.outbound + .as_mut() + .map(|edge| { + let id = edge.next_object_id; + edge.next_object_id = edge.next_object_id.saturating_add(1); + id + }) + .ok_or_else(|| "outbound edge missing".to_owned()) + } + + /// The edge id of the provisioned inbound edge. + pub fn inbound_edge_id(&self) -> Option { + self.inbound.as_ref().map(|edge| edge.edge_id) + } + + /// The edge id of the provisioned outbound edge. + pub fn outbound_edge_id(&self) -> Option { + self.outbound.as_ref().map(|edge| edge.edge_id) + } + + // ─── Ingress ──────────────────────────────────────────────────────────── + + /// Buffer stream bytes on the inbound edge, parse complete object + /// records, write them into the edge's ring, and load them on the + /// worker. + fn ingress_bytes( + &mut self, + arena: &mut ArenaManager, + worker: &mut dyn WorkerPort, + edge_id: EdgeId, + stream_id: StreamId, + bytes: Vec, + ) -> Result<(), String> { + let (parse_spec, records, buffered_bytes) = { + let Some(inbound) = &mut self.inbound else { + return Ok(()); + }; + if inbound.edge_id != edge_id { + return Ok(()); + } + let buffer = inbound.buffers.entry(stream_id).or_default(); + buffer.extend_from_slice(&bytes); + let buffered_bytes = buffer.len(); + let mut records = Vec::new(); + loop { + match take_complete_record(buffer, inbound.parse_spec) { + Ok(Some(record)) => records.push(record), + Ok(None) => break, + Err(e) => { + self.observations.push(Observation::ObjectFailed { + edge_id, + object_id: None, + }); + return Err(e); + } + } + } + (inbound.parse_spec, records, buffered_bytes) + }; + for IngressRecord { bytes, record } in records { + let ring_id = self + .inbound + .as_ref() + .and_then(|edge| edge.ring_id) + .ok_or_else(|| "inbound ring missing".to_owned())?; + let write_started = Instant::now(); + { + let lease = arena + .lookup_lease(ring_id) + .ok_or_else(|| format!("inbound ring {} lease missing", ring_id.0))?; + let data_offset = lease.layout.data_offset; + arena + .write_arena(data_offset, &bytes) + .map_err(|e| format!("write ingress ring: {e}"))?; + } + self.observations.push(Observation::IngressRingWrite { + edge_id, + ring_id, + stream_id, + object_id: record.object_id.0, + sequence: record.sequence, + extent: record.extent, + begin_sequence: record.flags.begin_sequence, + end_of_sequence: record.flags.end_of_sequence, + record_bytes: bytes.len(), + buffered_bytes, + write_ms: elapsed_ms(write_started), + }); + let load_started = Instant::now(); + let loaded = match worker.load_object(edge_id, ring_id, &record, &parse_spec) { + Ok(loaded) => loaded, + Err(e) => { + self.observations.push(Observation::ObjectFailed { + edge_id, + object_id: Some(record.object_id.0), + }); + return Err(e); + } + }; + self.observations.push(Observation::ObjectLoaded { + edge_id, + ring_id, + stream_id, + object: loaded, + load_ms: elapsed_ms(load_started), + }); + self.object_handles + .insert((edge_id, loaded.object_id), loaded); + } + Ok(()) + } + + // ─── Workflow fixpoint ────────────────────────────────────────────────── + + /// Run command execution and lifecycle-event draining until nothing + /// progresses. + fn drive( + &mut self, + transport: &mut T, + arena: &mut ArenaManager, + worker: &mut dyn WorkerPort, + ) -> Result<(), String> { + loop { + let progressed = self.execute_commands(transport, arena, worker)? + || self.drain_lifecycle()?; + if !progressed { + break; + } + } + Ok(()) + } + + /// Execute every not-yet-executed establisher command against the arena, + /// worker, and transport. + fn execute_commands( + &mut self, + transport: &mut T, + arena: &mut ArenaManager, + worker: &mut dyn WorkerPort, + ) -> Result { + let mut progressed = false; + while self.command_cursor < self.establisher.commands().len() { + let command = self.establisher.commands()[self.command_cursor].clone(); + self.command_cursor += 1; + progressed = true; + self.execute_command(command, transport, arena, worker)?; + } + Ok(progressed) + } + + fn execute_command( + &mut self, + command: EdgeCommand, + transport: &mut T, + arena: &mut ArenaManager, + worker: &mut dyn WorkerPort, + ) -> Result<(), String> { + match command { + EdgeCommand::LeaseRing { + request_id, + ring_spec, + .. + } => { + let events = arena.request(ArenaRequest::LeaseRing(LeaseRing { + request_id: LeaseRequestId(request_id.0), + ring_spec: ArenaRingSpec { + header_bytes: ring_spec.header_bytes, + data_bytes: ring_spec.data_bytes, + alignment: ring_spec.alignment, + }, + })); + for event in events { + match event { + ArenaEvent::RingLeased { lease } => { + self.establisher.observe(EdgeEvent::RingLeased { + request_id: LeaseRequestId(lease.request_id.0), + ring_id: RingId(lease.ring_id.0), + layout: edge_layout(&lease.layout), + }); + } + ArenaEvent::RingLeaseRejected { request_id, reason } => { + let reason = match reason { + crate::arena::RingLeaseRejection::CannotFitWithinCeiling => { + RingLeaseRejection::CannotFit + } + crate::arena::RingLeaseRejection::ArenaShuttingDown => { + RingLeaseRejection::ArenaShuttingDown + } + }; + self.establisher + .observe(EdgeEvent::RingLeaseRejected { + request_id: LeaseRequestId(request_id.0), + reason, + }); + } + ArenaEvent::RingLeaseQueued { .. } + | ArenaEvent::RingReleased { .. } + | ArenaEvent::RingReleaseRejected { .. } + | ArenaEvent::CancelledFreshLeaseReleased { .. } => {} + } + } + } + EdgeCommand::InstallWorkerRing { + edge_id, + ring_id, + direction, + object_spec, + .. + } => { + let layout = arena + .lookup_lease(RingId(ring_id.0)) + .map(|lease| lease.layout.clone()) + .ok_or_else(|| format!("ring {} lease missing", ring_id.0))?; + worker.install_ring(edge_id, ring_id, direction, &layout, &object_spec)?; + match direction { + RingDirection::Ingress => { + if let Some(inbound) = self + .inbound + .as_mut() + .filter(|edge| edge.edge_id == edge_id) + { + inbound.ring_id = Some(ring_id); + } + } + RingDirection::Egress => { + if let Some(outbound) = self + .outbound + .as_mut() + .filter(|edge| edge.edge_id == edge_id) + { + outbound.ring_id = Some(ring_id); + } + } + } + self.establisher + .observe(EdgeEvent::RingInstalled { edge_id, ring_id }); + } + EdgeCommand::EstablishSend { edge_id, .. } => { + let ring_id = self.edge_ring(edge_id)?; + let peer = self + .outbound + .as_ref() + .filter(|outbound| outbound.edge_id == edge_id) + .map(|outbound| outbound.peer.clone()) + .ok_or_else(|| "outbound edge missing".to_owned())?; + let writer = transport.open_writer(edge_id, &peer)?; + if let Some(outbound) = self.outbound.as_mut().filter(|edge| edge.edge_id == edge_id) + { + outbound.writer = Some(writer); + } + self.establish_send_wire(edge_id, ring_id); + } + EdgeCommand::EstablishRecv { edge_id, .. } => { + let ring_id = self.edge_ring(edge_id)?; + self.establish_recv_wire(edge_id, ring_id); + } + EdgeCommand::CancelQueuedLease { request_id, .. } => { + let _ = arena.request(ArenaRequest::CancelLease { + request_id: LeaseRequestId(request_id.0), + }); + } + EdgeCommand::StopPump { edge_id, .. } => { + self.stop_wire(edge_id); + } + EdgeCommand::UninstallWorkerRing { ring_id, .. } => { + worker.uninstall_ring(ring_id)?; + self.establisher + .observe(EdgeEvent::RingQuiesced { ring_id }); + } + EdgeCommand::ReleaseArenaLease { ring_id, proof } => { + let proof = if proof == crate::edge_lifecycle::QuiescenceProof::verified() { + ArenaQuiescenceProof::verified() + } else { + ArenaQuiescenceProof::missing() + }; + let _ = arena.request(ArenaRequest::ReleaseRing { + ring_id: RingId(ring_id.0), + proof, + }); + } + } + Ok(()) + } + + fn edge_ring(&self, edge_id: EdgeId) -> Result { + self.establisher + .local_record(edge_id) + .and_then(|record| record.ring_id) + .ok_or_else(|| format!("edge {} ring missing", edge_id.0)) + } + + /// Drain establisher lifecycle events into observations. An edge fault + /// is reported and then treated as fatal (mirrors the previous + /// composition, which halted the tick after reporting). + fn drain_lifecycle(&mut self) -> Result { + let mut progressed = false; + while self.lifecycle_cursor < self.establisher.events().len() { + let event = self.establisher.events()[self.lifecycle_cursor].clone(); + self.lifecycle_cursor += 1; + progressed = true; + match event { + EdgeLifecycleEvent::EdgeReady { edge_id, .. } => { + let direction = self + .establisher + .local_record(edge_id) + .map(|record| record.direction) + .unwrap_or(RingDirection::Ingress); + self.observations + .push(Observation::EdgeReady { edge_id, direction }); + } + EdgeLifecycleEvent::EdgeFaulted { edge_id, reason } => { + self.observations + .push(Observation::EdgeFaulted { edge_id, reason }); + return Err(format!("edge {} faulted: {reason:?}", edge_id.0)); + } + EdgeLifecycleEvent::EdgeStopped { edge_id } => { + self.observations + .push(Observation::EdgeStopped { edge_id }); + } + } + } + Ok(progressed) + } + + // ─── Wire bookkeeping (formerly iroh-driver `driver_pumps::Driver`) ───── + + fn establish_send_wire(&mut self, edge_id: EdgeId, ring_id: RingId) { + self.send_rings.insert(edge_id, ring_id); + self.establisher + .observe(EdgeEvent::DriverEdgeReady { edge_id }); + } + + fn establish_recv_wire(&mut self, edge_id: EdgeId, ring_id: RingId) { + self.recv_specs.insert(edge_id, ring_id); + if let Some(stream_id) = self.pending_streams.remove(&edge_id) { + self.spawn_recv(edge_id, stream_id); + } + } + + fn incoming_stream(&mut self, edge_id: EdgeId, stream_id: StreamId) { + if self.recv_specs.contains_key(&edge_id) { + self.spawn_recv(edge_id, stream_id); + } else { + self.pending_streams.insert(edge_id, stream_id); + } + } + + fn spawn_recv(&mut self, edge_id: EdgeId, stream_id: StreamId) { + let Some(ring_id) = self.recv_specs.get(&edge_id).copied() else { + self.pending_streams.insert(edge_id, stream_id); + return; + }; + self.recv_rings.insert(edge_id, ring_id); + self.establisher + .observe(EdgeEvent::DriverEdgeReady { edge_id }); + } + + fn read_error(&mut self, edge_id: EdgeId) { + self.establisher.observe(EdgeEvent::StreamFault { + edge_id, + reason: StreamFaultReason::ReadError, + }); + } + + fn stop_wire(&mut self, edge_id: EdgeId) { + let ring_id = self + .send_rings + .get(&edge_id) + .copied() + .or_else(|| self.recv_rings.get(&edge_id).copied()) + .or_else(|| self.recv_specs.get(&edge_id).copied()) + .unwrap_or(RingId(0)); + + self.send_rings.remove(&edge_id); + self.recv_rings.remove(&edge_id); + self.recv_specs.remove(&edge_id); + self.pending_streams.remove(&edge_id); + self.establisher + .observe(EdgeEvent::PumpStopped { edge_id, ring_id }); + } +} + +/// Parse one complete object record off the front of `buffer`, draining its +/// bytes; `Ok(None)` means more bytes are needed. +fn take_complete_record( + buffer: &mut Vec, + spec: object_record::ObjectSpec, +) -> Result, String> { + match object_record::read_object_record(buffer, spec, false) { + Ok(object_record::ObjectRecordRead::Incomplete) => Ok(None), + Ok(object_record::ObjectRecordRead::Complete(record)) => { + let bytes = buffer.drain(..record.total_len).collect(); + Ok(Some(IngressRecord { bytes, record })) + } + Err(reason) => Err(format!("invalid ingress record: {reason:?}")), + } +} + +fn edge_layout(layout: &ArenaRingLayout) -> crate::edge_lifecycle::RingLayout { + crate::edge_lifecycle::RingLayout { + start_offset: layout.start_offset, + header_offset: layout.header_offset, + data_offset: layout.data_offset, + end_offset: layout.end_offset, + data_bytes: layout.data_bytes, + alignment: layout.alignment, + } +} + +fn elapsed_ms(started: Instant) -> u64 { + started + .elapsed() + .as_millis() + .try_into() + .unwrap_or(u64::MAX) +} + diff --git a/crates/data-plane/src/edge_wire.rs b/crates/data-plane/src/edge_wire.rs new file mode 100644 index 0000000..07bb755 --- /dev/null +++ b/crates/data-plane/src/edge_wire.rs @@ -0,0 +1,63 @@ +//! Transport-port contracts for ring-backed edge byte streams. +//! +//! This is the *entire* contract a byte transport must satisfy for the +//! data-plane to drive it: open a writer for one edge, and report inbound +//! stream events. The iroh-driver crate provides the concrete implementation +//! over its QUIC endpoint; everything semantic — edges, rings, lifecycle, +//! object parsing — lives in this crate and consumes this port. + +use crate::ids::{EdgeId, StreamId}; + +/// One observed edge-byte transport event, in transport vocabulary only. +/// +/// Streams are unidirectional byte streams tagged with an edge id (the wire +/// preamble). The transport knows nothing about edges beyond passing the tag +/// through. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum WireEvent { + /// A new inbound stream arrived, tagged for `edge_id`. + StreamArrived { edge_id: EdgeId, stream_id: StreamId }, + /// Bytes were read from an inbound stream. + BytesRead { + edge_id: EdgeId, + stream_id: StreamId, + bytes: Vec, + }, + /// An inbound stream ended cleanly. + StreamEnded { edge_id: EdgeId, stream_id: StreamId }, + /// A stream or connection-level fault. `None` ids mean the fault could + /// not be attributed to a specific edge or stream. + StreamFault { + edge_id: Option, + stream_id: Option, + reason: WireFault, + }, +} + +/// Why an edge byte stream faulted. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WireFault { + ReadError, + WriteError, + ProtocolError, +} + +/// Cloneable-ish handle for writing opaque bytes to one edge's peer. +pub trait EdgeWriter { + fn send(&self, bytes: Vec) -> Result<(), String>; +} + +/// Byte transport port the data-plane edge runtime drives. +/// +/// `PeerAddr` is the transport's own peer-address notion (e.g. an iroh +/// `EndpointAddr`); the data-plane treats it as opaque. +pub trait EdgeTransport { + type Writer: EdgeWriter; + type PeerAddr: Clone; + + /// Open (or continue) the writer pumping bytes to `edge_id`'s peer. + fn open_writer(&mut self, edge_id: EdgeId, peer: &Self::PeerAddr) -> Result; + + /// Drain all transport events observed since the last call. + fn drain_events(&mut self) -> Vec; +} diff --git a/crates/data-plane/src/egress.rs b/crates/data-plane/src/egress.rs deleted file mode 100644 index 8558d86..0000000 --- a/crates/data-plane/src/egress.rs +++ /dev/null @@ -1,377 +0,0 @@ -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct WorkerGeneration(pub u64); -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct RingId(pub u64); -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct EdgeId(pub u64); -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct PortId(pub String); - -pub use crate::object_record::{ - FLAG_BEGIN_SEQUENCE, FLAG_END_OF_SEQUENCE, HEADER_LEN, KNOWN_FLAGS_MASK, OBJECT_MAGIC, - OBJECT_MAGIC_BYTES, OBJECT_VERSION, ObjectFlags, ObjectHeader, ObjectId, ObjectLayout, - ObjectSpec, -}; -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ObjectKey { - pub edge_id: EdgeId, - pub object_id: ObjectId, -} - -impl ObjectKey { - pub fn new(edge_id: EdgeId, object_id: ObjectId) -> Self { - Self { edge_id, object_id } - } -} -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct StepId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct DeviceHandle { - pub generation: WorkerGeneration, - pub id: u64, -} - -impl DeviceHandle { - pub fn new(generation: WorkerGeneration, id: u64) -> Self { - Self { generation, id } - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum RingDirection { - Ingress, - Egress, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct InstallRing { - pub ring_id: RingId, - pub edge_id: EdgeId, - pub port_id: PortId, - pub direction: RingDirection, - pub object_spec: ObjectSpec, - pub generation: WorkerGeneration, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct OutputBinding { - pub ring_id: RingId, - pub object_id: ObjectId, - pub sequence: u64, - pub extent: u64, - pub flags: ObjectFlags, - pub device_source: DeviceHandle, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum WorkerEgressEvent { - InstallRing(InstallRing), - ExecuteStep { - step_id: StepId, - outputs: Vec, - }, - HeaderReady { - object_id: ObjectId, - }, - EgressRingFull { - ring_id: RingId, - }, - RingWritable { - ring_id: RingId, - }, - DeviceToHostCopyCompleted { - object_id: ObjectId, - byte_count: u64, - }, - DeviceCopyFailed { - object_id: ObjectId, - }, - RoleStateUpdated { - step_id: StepId, - }, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum StepFailureReason { - InvalidOutputRing, - OutputExtentViolation, - DeviceCopyFailed, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum WorkerEgressOut { - ObjectProduced { - ring_id: RingId, - edge_id: EdgeId, - port_id: PortId, - object_id: ObjectId, - sequence: u64, - extent: u64, - }, - StepCompleted { - step_id: StepId, - }, - StepFailed { - step_id: StepId, - reason: StepFailureReason, - }, - RingFault { - ring_id: RingId, - }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum WakeHint { - RingReadable { ring_id: RingId }, - RingWritable { ring_id: RingId }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -struct PendingStep { - step_id: StepId, - outputs: Vec, - role_state_updated: bool, -} - -pub struct EgressProducerHarness { - generation: WorkerGeneration, - rings: std::collections::BTreeMap, - pending_outputs: Vec, - steps: Vec, - committed: std::collections::BTreeMap>, - payload_committed: std::collections::BTreeMap, - full_rings: std::collections::BTreeSet, - produced: std::collections::BTreeSet, - wake_hints: Vec, - events: Vec, -} - -impl EgressProducerHarness { - pub fn new(generation: WorkerGeneration) -> Self { - Self { - generation, - rings: std::collections::BTreeMap::new(), - pending_outputs: Vec::new(), - steps: Vec::new(), - committed: std::collections::BTreeMap::new(), - payload_committed: std::collections::BTreeMap::new(), - full_rings: std::collections::BTreeSet::new(), - produced: std::collections::BTreeSet::new(), - wake_hints: Vec::new(), - events: Vec::new(), - } - } - - pub fn observe(&mut self, event: WorkerEgressEvent) { - match event { - WorkerEgressEvent::InstallRing(install) => { - if install.direction == RingDirection::Egress - && install.generation == self.generation - { - self.rings.insert(install.ring_id, install); - } - } - WorkerEgressEvent::ExecuteStep { step_id, outputs } => { - self.execute_step(step_id, outputs) - } - WorkerEgressEvent::HeaderReady { object_id } => self.header_ready(object_id), - WorkerEgressEvent::EgressRingFull { ring_id } => { - self.full_rings.insert(ring_id); - } - WorkerEgressEvent::RingWritable { ring_id } => { - self.full_rings.remove(&ring_id); - self.wake_hints.push(WakeHint::RingWritable { ring_id }); - } - WorkerEgressEvent::DeviceToHostCopyCompleted { - object_id, - byte_count, - } => self.copy_completed(object_id, byte_count), - WorkerEgressEvent::DeviceCopyFailed { object_id } => self.copy_failed(object_id), - WorkerEgressEvent::RoleStateUpdated { step_id } => { - if let Some(step) = self.steps.iter_mut().find(|step| step.step_id == step_id) { - step.role_state_updated = true; - } - self.maybe_step_completed(step_id); - } - } - } - - pub fn pending_outputs(&self) -> &[OutputBinding] { - &self.pending_outputs - } - - pub fn committed_bytes(&self, ring_id: RingId) -> &[u8] { - self.committed - .get(&ring_id) - .map(Vec::as_slice) - .unwrap_or(&[]) - } - - pub fn committed_payload_bytes(&self, ring_id: RingId) -> u64 { - self.payload_committed.get(&ring_id).copied().unwrap_or(0) - } - - pub fn wake_hints(&self) -> &[WakeHint] { - &self.wake_hints - } - - pub fn events(&self) -> &[WorkerEgressOut] { - &self.events - } - - pub fn complete_output(&mut self, object_id: ObjectId) { - self.header_ready(object_id); - let Some(output) = self - .pending_outputs - .iter() - .find(|output| output.object_id == object_id) - .copied() - else { - return; - }; - self.copy_completed(object_id, output.extent); - } - - fn execute_step(&mut self, step_id: StepId, outputs: Vec) { - for output in &outputs { - let Some(ring) = self.rings.get(&output.ring_id) else { - self.events.push(WorkerEgressOut::StepFailed { - step_id, - reason: StepFailureReason::InvalidOutputRing, - }); - return; - }; - if output.extent > ring.object_spec.max_extent - || (ring.object_spec.alignment != 0 - && output.extent % ring.object_spec.alignment != 0) - { - self.events.push(WorkerEgressOut::StepFailed { - step_id, - reason: StepFailureReason::OutputExtentViolation, - }); - return; - } - } - self.pending_outputs.extend(outputs.iter().copied()); - self.steps.push(PendingStep { - step_id, - outputs, - role_state_updated: false, - }); - } - - fn header_ready(&mut self, object_id: ObjectId) { - let Some(output) = self - .pending_outputs - .iter() - .find(|output| output.object_id == object_id) - .copied() - else { - return; - }; - let bytes = encode_header(output); - self.committed - .entry(output.ring_id) - .or_default() - .extend(bytes); - self.wake_hints.push(WakeHint::RingReadable { - ring_id: output.ring_id, - }); - } - - fn copy_completed(&mut self, object_id: ObjectId, byte_count: u64) { - let Some(output) = self - .pending_outputs - .iter() - .find(|output| output.object_id == object_id) - .copied() - else { - return; - }; - if self.full_rings.contains(&output.ring_id) || byte_count != output.extent { - return; - } - self.committed - .entry(output.ring_id) - .or_default() - .extend(std::iter::repeat(0).take(byte_count as usize)); - *self.payload_committed.entry(output.ring_id).or_insert(0) += byte_count; - let Some(object_key) = self.output_key(output) else { - return; - }; - if self.produced.insert(object_key) { - let ring = self.rings.get(&output.ring_id).unwrap(); - self.events.push(WorkerEgressOut::ObjectProduced { - ring_id: output.ring_id, - edge_id: ring.edge_id, - port_id: ring.port_id.clone(), - object_id, - sequence: output.sequence, - extent: output.extent, - }); - } - let step_ids = self - .steps - .iter() - .filter(|step| { - step.outputs - .iter() - .any(|output| output.object_id == object_id) - }) - .map(|step| step.step_id) - .collect::>(); - for step_id in step_ids { - self.maybe_step_completed(step_id); - } - } - - fn copy_failed(&mut self, object_id: ObjectId) { - let step_id = self - .steps - .iter() - .find(|step| { - step.outputs - .iter() - .any(|output| output.object_id == object_id) - }) - .map(|step| step.step_id) - .unwrap_or(StepId(0)); - self.events.push(WorkerEgressOut::StepFailed { - step_id, - reason: StepFailureReason::DeviceCopyFailed, - }); - } - - fn output_key(&self, output: OutputBinding) -> Option { - self.rings - .get(&output.ring_id) - .map(|ring| ObjectKey::new(ring.edge_id, output.object_id)) - } - - fn maybe_step_completed(&mut self, step_id: StepId) { - let Some(step) = self.steps.iter().find(|step| step.step_id == step_id) else { - return; - }; - if !step.role_state_updated { - return; - } - if step.outputs.iter().all(|output| { - self.output_key(*output) - .is_some_and(|key| self.produced.contains(&key)) - }) && !self.events.iter().any(|event| matches!(event, WorkerEgressOut::StepCompleted { step_id: seen } if *seen == step_id)) - { - self.events.push(WorkerEgressOut::StepCompleted { step_id }); - } - } -} - -fn encode_header(output: OutputBinding) -> Vec { - ObjectHeader { - object_id: output.object_id, - sequence: output.sequence, - extent: output.extent, - flags: output.flags, - } - .encode() -} diff --git a/crates/data-plane/src/ids.rs b/crates/data-plane/src/ids.rs new file mode 100644 index 0000000..e743356 --- /dev/null +++ b/crates/data-plane/src/ids.rs @@ -0,0 +1,29 @@ +//! Single definitions of the data-plane identifier newtypes. +//! +//! Every module — and the iroh-driver edge transport — re-exports these, so +//! edge, ring, arena, and wire code share one vocabulary instead of defining +//! structurally identical `u64` newtypes and re-wrapping them at every crate +//! boundary. + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RunId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct NodeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct EdgeId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct RingId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct StreamId(pub u64); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct LeaseRequestId(pub u64); + +/// Data-plane-local actor address for edge endpoints. Distinct from the +/// swactor runtime's actor address; the application maps between them. +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ActorAddress(pub u64); diff --git a/crates/data-plane/src/ingress.rs b/crates/data-plane/src/ingress.rs deleted file mode 100644 index 709cefa..0000000 --- a/crates/data-plane/src/ingress.rs +++ /dev/null @@ -1,320 +0,0 @@ -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct WorkerGeneration(pub u64); -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct RingId(pub u64); -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct EdgeId(pub u64); -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct PortId(pub String); - -pub use crate::object_record::{ - FLAG_BEGIN_SEQUENCE, FLAG_END_OF_SEQUENCE, HEADER_LEN, KNOWN_FLAGS_MASK, OBJECT_MAGIC, - OBJECT_MAGIC_BYTES, OBJECT_VERSION, ObjectFailureReason, ObjectFlags, ObjectHeader, ObjectId, - ObjectLayout, ObjectRecord, ObjectRecordBuilder, ObjectRecordRead, ObjectSpec, - read_object_record, -}; - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct ObjectKey { - pub edge_id: EdgeId, - pub object_id: ObjectId, -} - -impl ObjectKey { - pub fn new(edge_id: EdgeId, object_id: ObjectId) -> Self { - Self { edge_id, object_id } - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct DeviceHandle { - pub generation: WorkerGeneration, - pub id: u64, -} - -impl DeviceHandle { - pub fn new(generation: WorkerGeneration, id: u64) -> Self { - Self { generation, id } - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum RingDirection { - Ingress, - Egress, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct InstallRing { - pub ring_id: RingId, - pub edge_id: EdgeId, - pub port_id: PortId, - pub direction: RingDirection, - pub object_spec: ObjectSpec, - pub generation: WorkerGeneration, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum WorkerIngressEvent { - InstallRing(InstallRing), - RingReadable { - ring_id: RingId, - }, - Eof { - ring_id: RingId, - }, - DeviceCopyCompleted { - object_id: ObjectId, - byte_count: u64, - }, - DeviceHandleCreated { - object_id: ObjectId, - handle: DeviceHandle, - }, - RingFault { - ring_id: RingId, - }, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum WorkerIngressOut { - ObjectLoaded { - ring_id: RingId, - edge_id: EdgeId, - port_id: PortId, - object_id: ObjectId, - sequence: u64, - extent: u64, - handle: DeviceHandle, - }, - ObjectFailed { - ring_id: RingId, - edge_id: EdgeId, - port_id: PortId, - object_id: Option, - sequence: Option, - reason: ObjectFailureReason, - }, - RingFault { - ring_id: RingId, - }, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct DeviceCopyLog { - pub object_id: ObjectId, - pub byte_count: u64, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -struct PendingObject { - record: ObjectRecord, - copy_done: bool, - handle: Option, -} - -pub struct IngressParserHarness { - generation: WorkerGeneration, - install: Option, - buffers: std::collections::BTreeMap>, - consume: std::collections::BTreeMap, - cursor_reload: std::collections::BTreeMap, - faulted_rings: std::collections::BTreeSet, - expected_sequence: u64, - pending: std::collections::BTreeMap, - copy_log: Vec, - events: Vec, -} - -impl IngressParserHarness { - pub fn new(generation: WorkerGeneration) -> Self { - Self { - generation, - install: None, - buffers: std::collections::BTreeMap::new(), - consume: std::collections::BTreeMap::new(), - cursor_reload: std::collections::BTreeMap::new(), - faulted_rings: std::collections::BTreeSet::new(), - expected_sequence: 0, - pending: std::collections::BTreeMap::new(), - copy_log: Vec::new(), - events: Vec::new(), - } - } - - pub fn observe(&mut self, event: WorkerIngressEvent) { - match event { - WorkerIngressEvent::InstallRing(install) => { - if install.direction == RingDirection::Ingress - && install.generation == self.generation - { - self.consume.entry(install.ring_id).or_insert(0); - self.install = Some(install); - } - } - WorkerIngressEvent::RingReadable { ring_id } => self.parse_ring(ring_id, false), - WorkerIngressEvent::Eof { ring_id } => self.parse_ring(ring_id, true), - WorkerIngressEvent::DeviceCopyCompleted { - object_id, - byte_count, - } => { - self.copy_log.push(DeviceCopyLog { - object_id, - byte_count, - }); - if let Some(key) = self.object_key(object_id) { - if let Some(pending) = self.pending.get_mut(&key) { - pending.copy_done = true; - if byte_count == pending.record.extent { - if let Some(install) = &self.install { - *self.consume.entry(install.ring_id).or_insert(0) += - pending.record.total_len as u64; - } - } - } - self.maybe_loaded(key); - } - } - WorkerIngressEvent::DeviceHandleCreated { object_id, handle } => { - if let Some(key) = self.object_key(object_id) { - if let Some(pending) = self.pending.get_mut(&key) { - pending.handle = Some(handle); - } - self.maybe_loaded(key); - } - } - WorkerIngressEvent::RingFault { ring_id } => { - self.faulted_rings.insert(ring_id); - self.events.push(WorkerIngressOut::RingFault { ring_id }); - } - } - } - - pub fn write_committed_bytes(&mut self, ring_id: RingId, bytes: Vec) { - self.buffers.entry(ring_id).or_default().extend(bytes); - } - - pub fn write_uncommitted_bytes(&mut self, _ring_id: RingId, _bytes: Vec) {} - - pub fn consume_cursor(&self, ring_id: RingId) -> u64 { - self.consume.get(&ring_id).copied().unwrap_or(0) - } - - pub fn cursor_reload_count(&self, ring_id: RingId) -> u64 { - self.cursor_reload.get(&ring_id).copied().unwrap_or(0) - } - - pub fn device_copy_log(&self) -> &[DeviceCopyLog] { - &self.copy_log - } - - pub fn events(&self) -> &[WorkerIngressOut] { - &self.events - } - - fn parse_ring(&mut self, ring_id: RingId, eof: bool) { - let Some(install) = &self.install else { - return; - }; - if install.ring_id != ring_id || self.faulted_rings.contains(&ring_id) { - return; - } - *self.cursor_reload.entry(ring_id).or_insert(0) += 1; - let buffer = self.buffers.get(&ring_id).cloned().unwrap_or_default(); - if buffer.is_empty() { - return; - } - match read_object_record(&buffer, install.object_spec, eof) { - Ok(ObjectRecordRead::Complete(record)) => { - if record.sequence != self.expected_sequence { - self.events.push(WorkerIngressOut::ObjectFailed { - ring_id, - edge_id: install.edge_id, - port_id: install.port_id.clone(), - object_id: Some(record.object_id), - sequence: Some(record.sequence), - reason: ObjectFailureReason::SequenceViolation, - }); - return; - } - self.expected_sequence += 1; - self.pending.insert( - ObjectKey::new(install.edge_id, record.object_id), - PendingObject { - record: record.clone(), - copy_done: false, - handle: None, - }, - ); - self.copy_log.push(DeviceCopyLog { - object_id: record.object_id, - byte_count: record.extent, - }); - } - Ok(ObjectRecordRead::Incomplete) => {} - Err(reason) => { - let (object_id, sequence) = object_failure_metadata(&buffer, reason); - self.events.push(WorkerIngressOut::ObjectFailed { - ring_id, - edge_id: install.edge_id, - port_id: install.port_id.clone(), - object_id, - sequence, - reason, - }); - } - } - } - - fn object_key(&self, object_id: ObjectId) -> Option { - self.install - .as_ref() - .map(|install| ObjectKey::new(install.edge_id, object_id)) - } - - fn maybe_loaded(&mut self, key: ObjectKey) { - let Some(pending) = self.pending.get(&key).cloned() else { - return; - }; - let Some(handle) = pending.handle else { - return; - }; - if !pending.copy_done || handle.generation != self.generation { - return; - } - let install = self.install.as_ref().unwrap(); - if !self.events.iter().any(|event| matches!(event, WorkerIngressOut::ObjectLoaded { edge_id, object_id, .. } if *edge_id == key.edge_id && *object_id == key.object_id)) { - self.events.push(WorkerIngressOut::ObjectLoaded { - ring_id: install.ring_id, - edge_id: install.edge_id, - port_id: install.port_id.clone(), - object_id: key.object_id, - sequence: pending.record.sequence, - extent: pending.record.extent, - handle, - }); - } - } -} - -fn object_failure_metadata( - bytes: &[u8], - reason: ObjectFailureReason, -) -> (Option, Option) { - if bytes.len() < HEADER_LEN - || matches!( - reason, - ObjectFailureReason::UnsupportedMagic - | ObjectFailureReason::UnsupportedVersion - | ObjectFailureReason::MalformedHeaderLength - ) - { - return (None, None); - } - ( - Some(ObjectId(u64::from_le_bytes( - bytes[8..16].try_into().unwrap(), - ))), - Some(u64::from_le_bytes(bytes[16..24].try_into().unwrap())), - ) -} diff --git a/crates/data-plane/src/lib.rs b/crates/data-plane/src/lib.rs index a7bf73a..2825dcf 100644 --- a/crates/data-plane/src/lib.rs +++ b/crates/data-plane/src/lib.rs @@ -1,11 +1,14 @@ //! Reusable actor-oriented data-plane contracts for wire edges, local IPC rings, //! arena-backed byte movement, and GPU worker object movement. +//! +//! Composition lives in [`edge_runtime`]: [`edge_runtime::EdgeRuntime`] drives +//! the edge lifecycle ([`edge_lifecycle`]) over a byte transport +//! ([`edge_wire`]), the arena ([`arena`]), and an application worker port. -pub mod actor; pub mod arena; -pub mod edge_actor; pub mod edge_lifecycle; -pub mod egress; -pub mod ingress; +pub mod edge_runtime; +pub mod edge_wire; +pub mod ids; pub mod object_record; pub mod ring; diff --git a/crates/data-plane/src/object_record.rs b/crates/data-plane/src/object_record.rs index 430c515..79d6366 100644 --- a/crates/data-plane/src/object_record.rs +++ b/crates/data-plane/src/object_record.rs @@ -268,3 +268,28 @@ pub fn read_object_record( total_len, })) } + +/// Monotonic per-producer source of object ids for one edge's outbound +/// objects. Ids start at 1. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ObjectIdAllocator { + next: u64, +} + +impl ObjectIdAllocator { + pub fn new() -> Self { + Self { next: 1 } + } + + pub fn alloc(&mut self) -> ObjectId { + let object_id = ObjectId(self.next); + self.next = self.next.saturating_add(1); + object_id + } +} + +impl Default for ObjectIdAllocator { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/data-plane/src/ring.rs b/crates/data-plane/src/ring.rs index c486355..20291d7 100644 --- a/crates/data-plane/src/ring.rs +++ b/crates/data-plane/src/ring.rs @@ -1,10 +1,6 @@ //! Reusable bounded ring cursor, wake, and process-local view contracts. -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct NodeId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct RingId(pub u64); +pub use crate::ids::{NodeId, RingId}; #[derive(Clone, Debug, PartialEq, Eq)] pub struct EndpointId(pub String); diff --git a/crates/data-plane/tests/data_plane_actor_guarantees.rs b/crates/data-plane/tests/data_plane_actor_guarantees.rs deleted file mode 100644 index 6e553e0..0000000 --- a/crates/data-plane/tests/data_plane_actor_guarantees.rs +++ /dev/null @@ -1,235 +0,0 @@ -use data_plane::actor as dp; -use data_plane::object_record; -use swactor::config::RuntimeConfig; -use swactor::runtime::{RuntimeParts, SingleThreadRuntime}; - -fn object_spec() -> dp::ObjectSpec { - dp::ObjectSpec { - kind: dp::ObjectKind::Activation, - dtype: dp::DType::F16, - max_extent_bytes: 4096, - } -} - -fn record_spec() -> object_record::ObjectSpec { - object_record::ObjectSpec { - max_extent: 4096, - alignment: 4, - layout: object_record::ObjectLayout::Token, - } -} - -fn ring_spec() -> dp::RingSpec { - dp::RingSpec { - header_bytes: 128, - data_bytes: 4096, - alignment: 64, - } -} - -fn layout() -> dp::RingLayout { - dp::RingLayout { - start_offset: 0, - header_offset: 0, - data_offset: 128, - end_offset: 4224, - data_bytes: 4096, - alignment: 64, - } -} - -fn inbound_endpoint() -> dp::WireEdgeEndpoint { - dp::WireEdgeEndpoint { - run_id: dp::RunId(55), - edge_id: dp::EdgeId(7001), - direction: dp::EdgeEndpointDirection::Inbound, - kind: dp::EdgeKind::Activation, - local_node_id: dp::NodeId(10), - peer_node_id: Some(dp::NodeId(9)), - peer_endpoint: Some(dp::PeerEndpoint("node-9".to_owned())), - local_role_port: dp::PortId("input".to_owned()), - object_spec: object_spec(), - ring_spec: ring_spec(), - object_record_spec: record_spec(), - transport: dp::TransportBinding::Remote { - endpoint: Some(dp::PeerEndpoint("node-9".to_owned())), - }, - worker_ring: dp::WorkerRingBinding::Required, - } -} - -#[test] -fn inbound_wire_edge_establishes_through_arena_worker_transport_then_reports_ready() { - let parts = RuntimeParts::new(RuntimeConfig::default()); - let runtime = parts.runtime().clone(); - let mut host = SingleThreadRuntime::new(parts); - let arena = runtime - .new_inbox::() - .expect("arena inbox"); - let worker = runtime - .new_inbox::() - .expect("worker inbox"); - let transport = runtime - .new_inbox::() - .expect("transport inbox"); - let reports = runtime - .new_inbox::() - .expect("report inbox"); - let actor = runtime - .spawn(dp::DataPlaneNodeActor::new(dp::NodeId(10))) - .expect("spawn data-plane actor"); - - runtime - .send_to( - actor, - dp::DataPlaneNodeMsg::ConfigureRun(dp::DataPlaneRunConfig { - run_id: dp::RunId(55), - local_node_id: dp::NodeId(10), - arena_actor: *arena.addr(), - worker_actor: *worker.addr(), - transport_actor: *transport.addr(), - report_sink: *reports.addr(), - }), - ) - .expect("configure run"); - runtime - .send_to( - actor, - dp::DataPlaneNodeMsg::ProvisionWireEdgeEndpoint(inbound_endpoint()), - ) - .expect("provision inbound"); - host.tick(); - - assert_eq!( - arena.try_recv(), - Some(dp::DataPlaneArenaMsg::LeaseRing { - request_id: dp::LeaseRequestId(1), - edge_id: dp::EdgeId(7001), - direction: dp::RingDirection::Ingress, - ring_spec: ring_spec(), - }) - ); - - runtime - .send_to( - actor, - dp::DataPlaneNodeMsg::Arena(dp::ArenaObservation::RingLeased { - request_id: dp::LeaseRequestId(1), - ring_id: dp::RingId(8001), - layout: layout(), - }), - ) - .expect("ring leased"); - host.tick(); - - assert_eq!( - worker.try_recv(), - Some(dp::DataPlaneWorkerMsg::InstallRing { - edge_id: dp::EdgeId(7001), - ring_id: dp::RingId(8001), - direction: dp::RingDirection::Ingress, - layout: layout(), - object_spec: object_spec(), - ring_spec: ring_spec(), - role_port: dp::PortId("input".to_owned()), - }) - ); - - runtime - .send_to( - actor, - dp::DataPlaneNodeMsg::Worker(dp::WorkerObservation::RingInstalled { - edge_id: dp::EdgeId(7001), - ring_id: dp::RingId(8001), - }), - ) - .expect("worker installed"); - host.tick(); - - assert_eq!( - transport.try_recv(), - Some(dp::DataPlaneTransportMsg::EstablishRecv { - edge_id: dp::EdgeId(7001), - ring_id: dp::RingId(8001), - layout: layout(), - }) - ); - - runtime - .send_to( - actor, - dp::DataPlaneNodeMsg::Transport(dp::TransportObservation::EdgeReady { - edge_id: dp::EdgeId(7001), - }), - ) - .expect("transport ready"); - host.tick(); - - assert_eq!( - reports.try_recv(), - Some(dp::DataPlaneReportMsg::InboundEdgeReady { - edge_id: dp::EdgeId(7001) - }) - ); -} - -#[test] -fn object_loaded_observation_is_reported_as_coarse_data_plane_outcome() { - let parts = RuntimeParts::new(RuntimeConfig::default()); - let runtime = parts.runtime().clone(); - let mut host = SingleThreadRuntime::new(parts); - let arena = runtime - .new_inbox::() - .expect("arena inbox"); - let worker = runtime - .new_inbox::() - .expect("worker inbox"); - let transport = runtime - .new_inbox::() - .expect("transport inbox"); - let reports = runtime - .new_inbox::() - .expect("report inbox"); - let actor = runtime - .spawn(dp::DataPlaneNodeActor::new(dp::NodeId(10))) - .expect("spawn data-plane actor"); - - runtime - .send_to( - actor, - dp::DataPlaneNodeMsg::ConfigureRun(dp::DataPlaneRunConfig { - run_id: dp::RunId(55), - local_node_id: dp::NodeId(10), - arena_actor: *arena.addr(), - worker_actor: *worker.addr(), - transport_actor: *transport.addr(), - report_sink: *reports.addr(), - }), - ) - .expect("configure run"); - runtime - .send_to( - actor, - dp::DataPlaneNodeMsg::Worker(dp::WorkerObservation::ObjectLoaded { - edge_id: dp::EdgeId(7001), - ring_id: dp::RingId(8001), - object_id: object_record::ObjectId(9000), - sequence: 7, - extent: 16, - handle: dp::DeviceHandle::new(dp::WorkerGeneration(3), 42), - }), - ) - .expect("object loaded"); - host.tick(); - - assert_eq!( - reports.try_recv(), - Some(dp::DataPlaneReportMsg::ObjectLoaded { - edge_id: dp::EdgeId(7001), - object_id: object_record::ObjectId(9000), - sequence: 7, - extent: 16, - handle: dp::DeviceHandle::new(dp::WorkerGeneration(3), 42), - }) - ); -} diff --git a/crates/data-plane/tests/edge_runtime_guarantees.rs b/crates/data-plane/tests/edge_runtime_guarantees.rs new file mode 100644 index 0000000..f15db28 --- /dev/null +++ b/crates/data-plane/tests/edge_runtime_guarantees.rs @@ -0,0 +1,350 @@ +//! Black-box contract tests for the data-plane edge runtime. +//! +//! These tests know only the public `EdgeRuntime` surface driven over a mock +//! byte transport and a mock worker port, with the real arena manager. They +//! assert the composition guarantees: +//! +//! - provisioning an edge reaches `EdgeReady` through the real lifecycle +//! (lease → worker ring → transport establishment), +//! - inbound streams arriving before recv establishment are held pending, +//! - complete ingress records are parsed, written to the edge's ring, and +//! surfaced as loaded objects, +//! - outbound edges open the transport writer and allocate object ids, +//! - ingress faults are fatal and reported as observations. +use parking_lot::Mutex; +use std::sync::Arc; + +use data_plane::arena::{ArenaConfig, ArenaManager}; +use data_plane::edge_lifecycle::{ + DType, NodeId, ObjectKind, ObjectSpec as EdgeObjectSpec, ProvisionRx, ProvisionTx, RingSpec, +}; +use data_plane::edge_runtime::{EdgeRuntime, LoadedObject, Observation, WorkerPort}; +use data_plane::edge_wire::{EdgeTransport, EdgeWriter, WireEvent}; +use data_plane::ids::{EdgeId, RingId, StreamId}; +use data_plane::object_record::{ + ObjectFlags, ObjectId, ObjectLayout, ObjectRecord, ObjectRecordBuilder, + ObjectSpec as ParseSpec, +}; + +// ─── Mock transport ───────────────────────────────────────────────────────── + +#[derive(Clone, Default)] +struct MockWriter { + sent: Arc>>>, +} + +impl EdgeWriter for MockWriter { + fn send(&self, bytes: Vec) -> Result<(), String> { + self.sent.lock().push(bytes); + Ok(()) + } +} + +struct MockTransport { + events: Vec, + opened: Vec, + writer: MockWriter, +} + +impl EdgeTransport for MockTransport { + type Writer = MockWriter; + type PeerAddr = (); + + fn open_writer(&mut self, edge_id: EdgeId, _peer: &()) -> Result { + self.opened.push(edge_id); + Ok(self.writer.clone()) + } + + fn drain_events(&mut self) -> Vec { + std::mem::take(&mut self.events) + } +} + +// ─── Mock worker ──────────────────────────────────────────────────────────── + +#[derive(Default)] +struct MockWorker { + installed: Vec<(EdgeId, RingId)>, + uninstalled: Vec, + fail_load: bool, +} + +impl WorkerPort for MockWorker { + fn install_ring( + &mut self, + edge_id: EdgeId, + ring_id: RingId, + _direction: data_plane::edge_lifecycle::RingDirection, + _layout: &data_plane::arena::RingLayout, + _object_spec: &EdgeObjectSpec, + ) -> Result<(), String> { + self.installed.push((edge_id, ring_id)); + Ok(()) + } + + fn uninstall_ring(&mut self, ring_id: RingId) -> Result<(), String> { + self.uninstalled.push(ring_id); + Ok(()) + } + + fn load_object( + &mut self, + _edge_id: EdgeId, + _ring_id: RingId, + record: &ObjectRecord, + _spec: &ParseSpec, + ) -> Result { + if self.fail_load { + return Err("mock load failure".to_owned()); + } + Ok(LoadedObject { + object_id: record.object_id.0, + sequence: record.sequence, + handle_generation: 3, + handle_id: 4242, + }) + } +} + +// ─── Fixtures ─────────────────────────────────────────────────────────────── + +fn boot_arena() -> ArenaManager { + ArenaManager::boot(ArenaConfig { + node_id: NodeId(10), + reservation_ceiling: 1 << 20, + base_alignment: 64, + }) + .expect("boot arena") +} + +fn parse_spec() -> ParseSpec { + ParseSpec { + max_extent: 4096, + alignment: 16, + layout: ObjectLayout::Token, + } +} + +fn provision_rx(edge_id: u64) -> ProvisionRx { + ProvisionRx { + run_id: data_plane::ids::RunId(1), + edge_id: EdgeId(edge_id), + local_node_id: NodeId(10), + object_spec: EdgeObjectSpec { + kind: ObjectKind::Activation, + dtype: DType::F16, + max_extent_bytes: 4096, + }, + ring_spec: RingSpec { + header_bytes: 0, + data_bytes: 8192, + alignment: 64, + }, + } +} + +fn provision_tx(edge_id: u64) -> ProvisionTx { + ProvisionTx { + run_id: data_plane::ids::RunId(1), + edge_id: EdgeId(edge_id), + local_node_id: NodeId(10), + consumer_node_id: NodeId(11), + object_spec: EdgeObjectSpec { + kind: ObjectKind::Activation, + dtype: DType::F16, + max_extent_bytes: 4096, + }, + ring_spec: RingSpec { + header_bytes: 0, + data_bytes: 8192, + alignment: 64, + }, + } +} + +fn record_bytes(object_id: u64, sequence: u64) -> Vec { + ObjectRecordBuilder::new(parse_spec()) + .object_id(ObjectId(object_id)) + .sequence(sequence) + .payload(vec![7_u8; 64]) + .flags(ObjectFlags::default()) + .encode() +} + +fn arena_runtime(events: Vec) -> ( + EdgeRuntime, + MockTransport, + ArenaManager, + MockWorker, +) { + let runtime = EdgeRuntime::new(NodeId(10)); + let transport = MockTransport { + events, + opened: Vec::new(), + writer: MockWriter::default(), + }; + (runtime, transport, boot_arena(), MockWorker::default()) +} + +fn find_observation<'a>( + observations: &'a [Observation], + predicate: impl Fn(&Observation) -> bool, +) -> Option<&'a Observation> { + observations.iter().find(|obs| predicate(obs)) +} + +// ─── Contracts ────────────────────────────────────────────────────────────── + +// A fully provisioned inbound edge must reach Ready through the real +// lifecycle: ring leased, worker ring installed, then — because the inbound +// stream already arrived — transport readiness observed and EdgeReady fired. +#[test] +fn inbound_edge_reaches_ready_and_delivers_objects() { + let (mut runtime, mut transport, mut arena, mut worker) = arena_runtime(vec![ + WireEvent::StreamArrived { + edge_id: EdgeId(7001), + stream_id: StreamId(1), + }, + WireEvent::BytesRead { + edge_id: EdgeId(7001), + stream_id: StreamId(1), + bytes: record_bytes(9001, 1), + }, + ]); + runtime.establish_inbound(provision_rx(7001), parse_spec()); + + runtime.poll(&mut transport, &mut arena, &mut worker).expect("poll"); + let observations = runtime.take_observations(); + + assert_eq!(worker.installed.len(), 1, "worker ring must be installed"); + let Observation::EdgeReady { direction, .. } = find_observation(&observations, |obs| { + matches!(obs, Observation::EdgeReady { .. }) + }) + .expect("edge ready observation") else { + unreachable!() + }; + assert_eq!(*direction, data_plane::edge_lifecycle::RingDirection::Ingress); + + let Observation::ObjectLoaded { object, .. } = find_observation(&observations, |obs| { + matches!(obs, Observation::ObjectLoaded { .. }) + }) + .expect("object loaded observation") else { + unreachable!() + }; + assert_eq!(object.object_id, 9001); + assert_eq!(object.sequence, 1); + + // The loaded object is retrievable by identity for compute admission. + let loaded = runtime + .loaded_object(EdgeId(7001), 9001) + .expect("loaded object handle"); + assert_eq!(loaded.handle_id, 4242); + + // The record bytes were written into the leased ingress ring. + let ring_id = worker.installed[0].1; + let lease = arena.lookup_lease(ring_id).expect("ingress lease"); + let written = arena + .read_arena(lease.layout.data_offset, 40 + 64) + .expect("read ring"); + assert_eq!(written, record_bytes(9001, 1)); +} + +// An outbound edge must open the transport writer, become ready, and hand +// out monotonically increasing output object ids. +#[test] +fn outbound_edge_opens_writer_and_allocates_object_ids() { + let (mut runtime, mut transport, mut arena, mut worker) = arena_runtime(Vec::new()); + runtime.establish_outbound(provision_tx(7002), ()); + + runtime.poll(&mut transport, &mut arena, &mut worker).expect("poll"); + let observations = runtime.take_observations(); + + assert!(matches!( + find_observation(&observations, |obs| matches!( + obs, + Observation::EdgeReady { .. } + )), + Some(_) + )); + assert_eq!(transport.opened, vec![EdgeId(7002)]); + assert_eq!(runtime.outbound_ring_id().map(|ring| ring.0), Some(1)); + assert!(runtime.outbound_writer().is_some()); + assert_eq!(runtime.alloc_output_object_id(), Ok(1)); + assert_eq!(runtime.alloc_output_object_id(), Ok(2)); +} + +// Streams that arrive before recv establishment are held pending; the edge +// still becomes Ready once establishment completes. +#[test] +fn early_stream_waits_for_recv_establishment() { + let (mut runtime, mut transport, mut arena, mut worker) = arena_runtime(vec![WireEvent::StreamArrived { + edge_id: EdgeId(7003), + stream_id: StreamId(9), + }]); + runtime.establish_inbound(provision_rx(7003), parse_spec()); + + runtime.poll(&mut transport, &mut arena, &mut worker).expect("poll"); + let observations = runtime.take_observations(); + assert!(matches!( + find_observation(&observations, |obs| matches!( + obs, + Observation::EdgeReady { .. } + )), + Some(_) + )); +} + +// A malformed ingress record must fault: ObjectFailed observation with no +// object id, and a fatal poll error. +#[test] +fn malformed_ingress_record_is_fatal_and_reported() { + let garbage = vec![0xDE; 64]; + let (mut runtime, mut transport, mut arena, mut worker) = arena_runtime(vec![WireEvent::BytesRead { + edge_id: EdgeId(7004), + stream_id: StreamId(1), + bytes: garbage, + }]); + runtime.establish_inbound(provision_rx(7004), parse_spec()); + + let result = runtime.poll(&mut transport, &mut arena, &mut worker); + assert!(result.is_err(), "malformed record must be fatal"); + let observations = runtime.take_observations(); + assert!(matches!( + find_observation(&observations, |obs| { + matches!( + obs, + Observation::ObjectFailed { object_id: None, .. } + ) + }), + Some(_) + )); +} + +// A worker load failure must surface ObjectFailed with the object id and +// remain fatal. +#[test] +fn worker_load_failure_reports_object_and_is_fatal() { + let (mut runtime, mut transport, mut arena, mut worker) = arena_runtime(vec![WireEvent::BytesRead { + edge_id: EdgeId(7005), + stream_id: StreamId(1), + bytes: record_bytes(9002, 1), + }]); + worker.fail_load = true; + runtime.establish_inbound(provision_rx(7005), parse_spec()); + + let result = runtime.poll(&mut transport, &mut arena, &mut worker); + assert!(result.is_err(), "load failure must be fatal"); + let observations = runtime.take_observations(); + assert!(matches!( + find_observation(&observations, |obs| { + matches!( + obs, + Observation::ObjectFailed { + object_id: Some(9002), + .. + } + ) + }), + Some(_) + )); +} diff --git a/crates/data-plane/tests/egress_guarantees.rs b/crates/data-plane/tests/egress_guarantees.rs deleted file mode 100644 index f30f96f..0000000 --- a/crates/data-plane/tests/egress_guarantees.rs +++ /dev/null @@ -1,316 +0,0 @@ -//! Black-box contract tests for data-plane GPU worker egress production. -//! -//! These tests intentionally know only the public data-plane egress surface: -//! -//! - `InstallRing`, `ExecuteStep` output bindings, device-copy outcomes, -//! backpressure, and shutdown events in -//! - committed ring bytes, cursor publication, `ObjectProduced`, -//! `StepCompleted`, and step failures out -//! -//! They assert the reusable data-plane egress producer contract. - -use data_plane::egress; - -// A valid egress ring config supplies the edge object spec and current worker -// generation. The producer remains free to choose copy scheduling internally. -fn egress_ring() -> egress::InstallRing { - egress::InstallRing { - ring_id: egress::RingId(8002), - edge_id: egress::EdgeId(7002), - port_id: egress::PortId("out".into()), - direction: egress::RingDirection::Egress, - object_spec: egress::ObjectSpec { - max_extent: 16, - alignment: 4, - layout: egress::ObjectLayout::Token, - }, - generation: egress::WorkerGeneration(1), - } -} - -// The harness exposes egress ring writes and worker events, not private output -// queues, device kernels, or role internals. -fn new_producer() -> egress::EgressProducerHarness { - egress::EgressProducerHarness::new(egress::WorkerGeneration(1)) -} - -// This helper installs the output ring through the public worker command path. -fn installed_producer() -> egress::EgressProducerHarness { - let mut harness = new_producer(); - harness.observe(egress::WorkerEgressEvent::InstallRing(egress_ring())); - harness -} - -// A valid output binding carries object identity, sequence, extent, flags, and -// target ring. The worker must not invent these graph-visible facts. -fn output_binding(sequence: u64, extent: u64) -> egress::OutputBinding { - egress::OutputBinding { - ring_id: egress::RingId(8002), - object_id: egress::ObjectId(9000 + sequence), - sequence, - extent, - flags: egress::ObjectFlags::default(), - device_source: egress::DeviceHandle::new(egress::WorkerGeneration(1), 40 + sequence), - } -} - -// This proves egress production starts only after InstallRing, only for -// ExecuteStep output bindings naming that ring, and never invents object ids or -// sequence numbers. -#[test] -fn output_admission_requires_installed_ring_and_explicit_binding() { - // Execute before InstallRing must not write output. - let mut not_installed = new_producer(); - not_installed.observe(egress::WorkerEgressEvent::ExecuteStep { - step_id: egress::StepId(77), - outputs: vec![output_binding(0, 8)], - }); - assert_eq!(not_installed.committed_bytes(egress::RingId(8002)).len(), 0); - - // Install the ring and execute with a binding that names it. - let mut harness = installed_producer(); - harness.observe(egress::WorkerEgressEvent::ExecuteStep { - step_id: egress::StepId(77), - outputs: vec![output_binding(0, 8)], - }); - - // The pending output identity must match the binding exactly. - let pending = harness.pending_outputs(); - assert_eq!(pending[0].object_id, egress::ObjectId(9000)); - assert_eq!(pending[0].sequence, 0); - assert_eq!(pending[0].extent, 8); -} - -// This proves the worker creates a valid ObjectHeader from ObjectSpec, writes -// header bytes before payload bytes, advances commit only after valid header -// bytes, and emits readable wake after committed header bytes. -#[test] -fn header_is_written_and_committed_before_payload() { - // Start one egress output. - let mut output = output_binding(0, 8); - output.flags = egress::ObjectFlags { - end_of_sequence: true, - begin_sequence: false, - }; - let mut harness = installed_producer(); - harness.observe(egress::WorkerEgressEvent::ExecuteStep { - step_id: egress::StepId(77), - outputs: vec![output], - }); - - // Complete header production but not payload copy. - harness.observe(egress::WorkerEgressEvent::HeaderReady { - object_id: egress::ObjectId(9000), - }); - - // The committed prefix must decode as a header for the configured spec. - let committed = harness.committed_bytes(egress::RingId(8002)); - let header = egress::ObjectHeader::decode(committed).expect("header must decode"); - assert_eq!(header.object_id, egress::ObjectId(9000)); - assert_eq!(header.sequence, 0); - assert_eq!(header.extent, 8); - assert_eq!( - header.flags, - egress::ObjectFlags { - end_of_sequence: true, - begin_sequence: false, - } - ); - - // Payload bytes are not committed before the payload copy is valid. - assert_eq!(harness.committed_payload_bytes(egress::RingId(8002)), 0); - assert!(harness.wake_hints().iter().any(|wake| { - matches!( - wake, - egress::WakeHint::RingReadable { - ring_id: egress::RingId(8002) - } - ) - })); -} - -// This proves payload production copies exactly extent bytes from device to the -// egress ring, advances commit only after host bytes are valid, and blocks on -// egress backpressure without dropping ownership. -#[test] -fn payload_copy_is_exact_extent_and_respects_backpressure() { - // Start one output with extent 8. - let mut harness = installed_producer(); - harness.observe(egress::WorkerEgressEvent::ExecuteStep { - step_id: egress::StepId(77), - outputs: vec![output_binding(0, 8)], - }); - harness.observe(egress::WorkerEgressEvent::HeaderReady { - object_id: egress::ObjectId(9000), - }); - - // Backpressure prevents committing payload bytes. - harness.observe(egress::WorkerEgressEvent::EgressRingFull { - ring_id: egress::RingId(8002), - }); - harness.observe(egress::WorkerEgressEvent::DeviceToHostCopyCompleted { - object_id: egress::ObjectId(9000), - byte_count: 4, - }); - assert_eq!(harness.committed_payload_bytes(egress::RingId(8002)), 0); - - // Once writable, the full exact extent can commit. - harness.observe(egress::WorkerEgressEvent::RingWritable { - ring_id: egress::RingId(8002), - }); - harness.observe(egress::WorkerEgressEvent::DeviceToHostCopyCompleted { - object_id: egress::ObjectId(9000), - byte_count: 8, - }); - assert_eq!(harness.committed_payload_bytes(egress::RingId(8002)), 8); -} - -// This proves ObjectProduced is emitted after the full output object is -// committed, and StepCompleted is emitted only after all declared outputs are -// produced and role state updates are complete. -#[test] -fn object_produced_precedes_step_completed_after_all_outputs() { - // Execute a step with two outputs. - let mut harness = installed_producer(); - harness.observe(egress::WorkerEgressEvent::InstallRing( - egress::InstallRing { - ring_id: egress::RingId(8003), - edge_id: egress::EdgeId(7003), - port_id: egress::PortId("out2".into()), - ..egress_ring() - }, - )); - harness.observe(egress::WorkerEgressEvent::ExecuteStep { - step_id: egress::StepId(77), - outputs: vec![ - output_binding(0, 8), - egress::OutputBinding { - ring_id: egress::RingId(8003), - object_id: egress::ObjectId(9100), - sequence: 0, - extent: 8, - flags: egress::ObjectFlags::default(), - device_source: egress::DeviceHandle::new(egress::WorkerGeneration(1), 55), - }, - ], - }); - - // Produce only the first output and prove StepCompleted is still absent. - harness.complete_output(egress::ObjectId(9000)); - assert!( - !harness - .events() - .iter() - .any(|event| { matches!(event, egress::WorkerEgressOut::StepCompleted { .. }) }) - ); - - // Produce the second output and complete role state update. - harness.complete_output(egress::ObjectId(9100)); - harness.observe(egress::WorkerEgressEvent::RoleStateUpdated { - step_id: egress::StepId(77), - }); - - // Both object-produced events precede StepCompleted. - let first_object_pos = harness - .events() - .iter() - .position(|event| { - matches!( - event, - egress::WorkerEgressOut::ObjectProduced { - object_id: egress::ObjectId(9000), - .. - } - ) - }) - .expect("first object produced"); - let second_object_pos = harness - .events() - .iter() - .position(|event| { - matches!( - event, - egress::WorkerEgressOut::ObjectProduced { - object_id: egress::ObjectId(9100), - .. - } - ) - }) - .expect("second object produced"); - let completed_pos = harness - .events() - .iter() - .position(|event| { - matches!( - event, - egress::WorkerEgressOut::StepCompleted { - step_id: egress::StepId(77), - .. - } - ) - }) - .expect("step completed"); - assert!(first_object_pos < completed_pos); - assert!(second_object_pos < completed_pos); -} - -// This proves invalid output ring, extent violation, device copy failure, and -// shutdown reject or abort egress production with visible step/ring faults. -#[test] -fn egress_faults_are_visible_and_suppress_success_events() { - // Invalid output ring fails the step. - let mut invalid_ring = installed_producer(); - invalid_ring.observe(egress::WorkerEgressEvent::ExecuteStep { - step_id: egress::StepId(77), - outputs: vec![egress::OutputBinding { - ring_id: egress::RingId(9999), - ..output_binding(0, 8) - }], - }); - assert!(invalid_ring.events().iter().any(|event| { - matches!( - event, - egress::WorkerEgressOut::StepFailed { - reason: egress::StepFailureReason::InvalidOutputRing, - .. - } - ) - })); - - // Extent violation fails the step. - let mut bad_extent = installed_producer(); - bad_extent.observe(egress::WorkerEgressEvent::ExecuteStep { - step_id: egress::StepId(78), - outputs: vec![output_binding(0, 32)], - }); - assert!(bad_extent.events().iter().any(|event| { - matches!( - event, - egress::WorkerEgressOut::StepFailed { - reason: egress::StepFailureReason::OutputExtentViolation, - .. - } - ) - })); - - // Copy failure faults the ring or fails the step, but must not emit - // ObjectProduced. - let mut copy_failed = installed_producer(); - copy_failed.observe(egress::WorkerEgressEvent::ExecuteStep { - step_id: egress::StepId(79), - outputs: vec![output_binding(0, 8)], - }); - copy_failed.observe(egress::WorkerEgressEvent::DeviceCopyFailed { - object_id: egress::ObjectId(9000), - }); - assert!(copy_failed.events().iter().any(|event| { - matches!(event, egress::WorkerEgressOut::StepFailed { .. }) - || matches!(event, egress::WorkerEgressOut::RingFault { .. }) - })); - assert!( - !copy_failed - .events() - .iter() - .any(|event| { matches!(event, egress::WorkerEgressOut::ObjectProduced { .. }) }) - ); -} diff --git a/crates/iroh-driver/Cargo.toml b/crates/iroh-driver/Cargo.toml index 24eb439..a670427 100644 --- a/crates/iroh-driver/Cargo.toml +++ b/crates/iroh-driver/Cargo.toml @@ -12,6 +12,7 @@ path = "src/lib.rs" swactor = { path = "../..", features = ["serde", "transport"] } swactor-engine = { path = "../engine" } swactor-transport = { path = "../transport" } +data-plane = { path = "../data-plane" } distribution = { path = "../distribution" } telemetry = { path = "../telemetry" } crossbeam-channel = "0.5" diff --git a/crates/iroh-driver/IROH_DRIVER_SPEC.md b/crates/iroh-driver/IROH_DRIVER_SPEC.md index 9ff3ece..40176e2 100644 --- a/crates/iroh-driver/IROH_DRIVER_SPEC.md +++ b/crates/iroh-driver/IROH_DRIVER_SPEC.md @@ -524,13 +524,21 @@ relay URLs ## 5. Code Architecture -The crate has two public modules: +The crate has four public modules: ```text iroh_driver telemetry_transport +edge_transport +endpoint_advertisement ``` +Edge semantics (lifecycle, ring bookkeeping, object-record parsing) live in +the `data-plane` crate's edge runtime. This crate's edge surface is the thin +`data_plane::edge_wire::EdgeTransport` port implementation in +`edge_transport`: open one writer per outbound edge, and surface inbound +edge-stream events. `IrohDriver` implements the port directly. + ### 5.1 Driver Core `IrohDriver` owns: diff --git a/crates/iroh-driver/src/driver_pumps.rs b/crates/iroh-driver/src/driver_pumps.rs deleted file mode 100644 index 90ee322..0000000 --- a/crates/iroh-driver/src/driver_pumps.rs +++ /dev/null @@ -1,102 +0,0 @@ -//! Driver edge/ring/stream bookkeeping. -//! -//! The driver is the node's swactor-to-iroh boundary. It tracks which edges -//! have an established send or recv pump, maps inbound uni-streams to their -//! recv rings, and emits driver lifecycle events (edge ready, stream fault, -//! pump stopped). This module is the pure state machine; the iroh stream pumps -//! themselves live in [`crate::edge_transport`]. - -use std::collections::BTreeMap; - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct EdgeId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct RingId(pub u64); - -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct StreamId(pub u64); - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum DriverEventOut { - DriverEdgeReady { edge_id: EdgeId }, - StreamFault { edge_id: EdgeId }, - PumpStopped { edge_id: EdgeId, ring_id: RingId }, -} - -#[derive(Debug)] -pub struct Driver { - sends: BTreeMap, - recv_specs: BTreeMap, - pending_streams: BTreeMap, - recvs: BTreeMap, - events: Vec, -} - -impl Driver { - pub fn new() -> Self { - Self { - sends: BTreeMap::new(), - recv_specs: BTreeMap::new(), - pending_streams: BTreeMap::new(), - recvs: BTreeMap::new(), - events: Vec::new(), - } - } - - pub fn establish_send(&mut self, edge_id: EdgeId, ring_id: RingId) { - self.sends.insert(edge_id, ring_id); - self.events - .push(DriverEventOut::DriverEdgeReady { edge_id }); - } - - pub fn establish_recv(&mut self, edge_id: EdgeId, ring_id: RingId) { - self.recv_specs.insert(edge_id, ring_id); - if let Some(stream_id) = self.pending_streams.remove(&edge_id) { - self.spawn_recv(edge_id, stream_id); - } - } - - pub fn incoming_uni_stream(&mut self, edge_id: EdgeId, stream_id: StreamId) { - if self.recv_specs.contains_key(&edge_id) { - self.spawn_recv(edge_id, stream_id); - } else { - self.pending_streams.insert(edge_id, stream_id); - } - } - - fn spawn_recv(&mut self, edge_id: EdgeId, stream_id: StreamId) { - let Some(ring_id) = self.recv_specs.get(&edge_id).copied() else { - self.pending_streams.insert(edge_id, stream_id); - return; - }; - self.recvs.insert(edge_id, ring_id); - self.events - .push(DriverEventOut::DriverEdgeReady { edge_id }); - } - - pub fn read_error(&mut self, edge_id: EdgeId) { - self.events.push(DriverEventOut::StreamFault { edge_id }); - } - - pub fn stop_edge(&mut self, edge_id: EdgeId) { - let ring_id = self - .sends - .get(&edge_id) - .copied() - .or_else(|| self.recvs.get(&edge_id).copied()) - .or_else(|| self.recv_specs.get(&edge_id).copied()) - .unwrap_or(RingId(0)); - - self.sends.remove(&edge_id); - self.recvs.remove(&edge_id); - self.recv_specs.remove(&edge_id); - self.pending_streams.remove(&edge_id); - self.events - .push(DriverEventOut::PumpStopped { edge_id, ring_id }); - } - - pub fn events(&self) -> &[DriverEventOut] { - &self.events - } -} diff --git a/crates/iroh-driver/src/edge_transport.rs b/crates/iroh-driver/src/edge_transport.rs index cb20521..691e6d3 100644 --- a/crates/iroh-driver/src/edge_transport.rs +++ b/crates/iroh-driver/src/edge_transport.rs @@ -1,13 +1,15 @@ //! Driver-owned byte transport for ring-backed MVP edge protocols. //! -//! This module deliberately owns only transport framing: one edge id preamble per -//! unidirectional stream, followed by opaque byte chunks. Object-record parsing, -//! ring ownership, and stage semantics stay in the MVP/dataplane crates. +//! This module deliberately owns only transport framing: one edge id preamble +//! per unidirectional stream, followed by opaque byte chunks. Object-record +//! parsing, ring ownership, and edge semantics stay in the data-plane crate, +//! which drives this transport through the `data_plane::edge_wire` port +//! (`IrohDriver` implements `EdgeTransport`). use std::sync::Arc; -use std::time::Duration; -use distribution::types::NodeId; +use data_plane::edge_wire::{EdgeWriter, WireEvent, WireFault}; +use data_plane::ids::{EdgeId, StreamId}; use iroh::endpoint::Connection; use iroh::{Endpoint, EndpointAddr}; use parking_lot::Mutex; @@ -17,39 +19,8 @@ use tokio::sync::mpsc as tokio_mpsc; pub const EDGE_ALPN: &[u8] = b"mvp/pipeline-edge/0"; -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum EdgeTransportEvent { - StreamArrived { - peer: NodeId, - edge_id: u64, - stream_id: u64, - }, - BytesRead { - peer: NodeId, - edge_id: u64, - stream_id: u64, - bytes: Vec, - }, - StreamEnded { - peer: NodeId, - edge_id: u64, - stream_id: u64, - }, - StreamFault { - peer: NodeId, - edge_id: Option, - stream_id: Option, - reason: EdgeTransportFault, - }, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum EdgeTransportFault { - ReadError, - WriteError, - ProtocolError, -} - +/// Cloneable handle for pushing opaque record bytes onto one edge's send +/// pump. Implements the data-plane [`EdgeWriter`] port. #[derive(Clone)] pub struct EdgeSendHandle { tx: tokio_mpsc::UnboundedSender>, @@ -63,6 +34,12 @@ impl EdgeSendHandle { } } +impl EdgeWriter for EdgeSendHandle { + fn send(&self, bytes: Vec) -> Result<(), String> { + self.send(bytes) + } +} + pub(crate) fn spawn_edge_send_pump( engine: EngineHandle, endpoint: Endpoint, @@ -100,16 +77,17 @@ pub(crate) fn spawn_edge_send_pump( let mut attempts = 0_u8; loop { attempts = attempts.saturating_add(1); - let write_result = engine_handle.timeout(Duration::from_secs(30), async { - send.write_all(&record) - .await - .map_err(|e| format!("write edge record {edge_id}: {e}"))?; - send.flush() - .await - .map_err(|e| format!("flush edge record {edge_id}: {e}")) - }) - .await - .map_err(|_| format!("write edge record {edge_id}: timed out"))?; + let write_result = engine_handle + .timeout(std::time::Duration::from_secs(30), async { + send.write_all(&record) + .await + .map_err(|e| format!("write edge record {edge_id}: {e}"))?; + send.flush() + .await + .map_err(|e| format!("flush edge record {edge_id}: {e}")) + }) + .await + .map_err(|_| format!("write edge record {edge_id}: timed out"))?; match write_result { Ok(()) => break, @@ -139,28 +117,25 @@ pub(crate) fn spawn_edge_send_pump( pub(crate) fn spawn_edge_recv_pump( engine: EngineHandle, conn: Connection, - peer: NodeId, - events: Arc>>, + events: Arc>>, stream_group: u64, ) { engine.spawn(async move { let mut next_uni_stream_id = stream_group << 32; while let Ok(mut recv) = conn.accept_uni().await { next_uni_stream_id = next_uni_stream_id.saturating_add(1); - let current_stream_id = next_uni_stream_id; + let current_stream_id = StreamId(next_uni_stream_id); let mut preamble = [0u8; 8]; if recv.read_exact(&mut preamble).await.is_err() { - events.lock().push(EdgeTransportEvent::StreamFault { - peer, + events.lock().push(WireEvent::StreamFault { edge_id: None, stream_id: Some(current_stream_id), - reason: EdgeTransportFault::ProtocolError, + reason: WireFault::ProtocolError, }); continue; } - let edge_id = u64::from_le_bytes(preamble); - events.lock().push(EdgeTransportEvent::StreamArrived { - peer, + let edge_id = EdgeId(u64::from_le_bytes(preamble)); + events.lock().push(WireEvent::StreamArrived { edge_id, stream_id: current_stream_id, }); @@ -168,27 +143,24 @@ pub(crate) fn spawn_edge_recv_pump( loop { match recv.read(&mut chunk).await { Ok(Some(0)) | Ok(None) => { - events.lock().push(EdgeTransportEvent::StreamEnded { - peer, + events.lock().push(WireEvent::StreamEnded { edge_id, stream_id: current_stream_id, }); break; } Ok(Some(n)) => { - events.lock().push(EdgeTransportEvent::BytesRead { - peer, + events.lock().push(WireEvent::BytesRead { edge_id, stream_id: current_stream_id, bytes: chunk[..n].to_vec(), }); } Err(_) => { - events.lock().push(EdgeTransportEvent::StreamFault { - peer, + events.lock().push(WireEvent::StreamFault { edge_id: Some(edge_id), stream_id: Some(current_stream_id), - reason: EdgeTransportFault::ReadError, + reason: WireFault::ReadError, }); break; } diff --git a/crates/iroh-driver/src/iroh_driver.rs b/crates/iroh-driver/src/iroh_driver.rs index 2d2c045..348d919 100644 --- a/crates/iroh-driver/src/iroh_driver.rs +++ b/crates/iroh-driver/src/iroh_driver.rs @@ -31,10 +31,9 @@ use distribution::swim::actor::SwimIn; use distribution::transport_bridge::{OutFrame, Outbox, RelayMirror, RouteView, peer_addr}; use distribution::types::NodeId; -use crate::edge_transport::{ - EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, spawn_edge_recv_pump, - spawn_edge_send_pump as spawn_edge_sender_task, -}; +use crate::edge_transport::{EDGE_ALPN, EdgeSendHandle, spawn_edge_recv_pump}; +use crate::edge_transport::spawn_edge_send_pump as spawn_edge_sender_task; +use data_plane::edge_wire::WireEvent; use crate::telemetry_transport::{ TELEMETRY_ALPN, TelemetryQuicHeader, TelemetryQuicRead, read_events_from_stream, spawn_subscription_writer, @@ -291,7 +290,7 @@ pub struct IrohDriver { /// Completed telemetry QUIC reads from driver-owned TELEMETRY_ALPN adapters. telemetry_reads: Arc>>, /// Logical edge events emitted by driver-owned EDGE_ALPN byte pumps. - edge_events: Arc>>, + edge_events: Arc>>, next_edge_stream_group: Arc, /// Frames read by per-connection reader tasks, drained by the engine-hosted /// adapter pump ([`Self::install_actor_bridge_pump`]). This decouples network @@ -445,7 +444,7 @@ impl IrohDriver { let other_accepted_conns: Arc, Connection)>>> = Arc::new(Mutex::new(Vec::new())); let telemetry_reads: Arc>> = Arc::new(Mutex::new(Vec::new())); - let edge_events: Arc>> = Arc::new(Mutex::new(Vec::new())); + let edge_events: Arc>> = Arc::new(Mutex::new(Vec::new())); { let ep = endpoint.clone(); let peer_auth = config.peer_auth.clone(); @@ -577,14 +576,14 @@ impl IrohDriver { } /// Drain logical edge transport events emitted by driver-owned byte pumps. - pub fn drain_edge_events(&self) -> Vec { + pub fn drain_edge_events(&self) -> Vec { self.edge_events.lock().drain(..).collect() } /// Clone of the shared edge-event queue, so callers outside the driver /// (e.g. a job worker) can drain `EDGE_ALPN` byte events from their own /// thread/task without going through `&self`. - pub fn edge_events_handle(&self) -> Arc>> { + pub fn edge_events_handle(&self) -> Arc>> { Arc::clone(&self.edge_events) } @@ -1181,7 +1180,7 @@ struct AdapterPump { accepted_conns: Arc>>, other_accepted_conns: Arc, Connection)>>>, telemetry_reads: Arc>>, - edge_events: Arc>>, + edge_events: Arc>>, next_edge_stream_group: Arc, dialing: Arc>>, peer_auth: Option>>, @@ -1518,7 +1517,7 @@ impl AdapterPump { *pending = keep; drained }; - for (node, conn) in drained { + for (_node, conn) in drained { let stream_group = self .next_edge_stream_group .fetch_add(1, Ordering::Relaxed) @@ -1526,7 +1525,6 @@ impl AdapterPump { spawn_edge_recv_pump( self.engine.clone(), conn, - node, Arc::clone(&self.edge_events), stream_group, ); @@ -1624,3 +1622,25 @@ async fn read_message( Ok((dest, tag, payload)) } + +// ─── Data-plane edge transport port ──────────────────────────────────────── + +/// The driver as a data-plane byte transport: open one writer per outbound +/// edge and expose inbound edge-stream events. All edge semantics live in +/// the data-plane crate's edge runtime; this impl is deliberately thin. +impl data_plane::edge_wire::EdgeTransport for IrohDriver { + type Writer = EdgeSendHandle; + type PeerAddr = EndpointAddr; + + fn open_writer( + &mut self, + edge_id: data_plane::ids::EdgeId, + peer: &EndpointAddr, + ) -> Result { + self.spawn_edge_send_pump(peer.clone(), edge_id.0) + } + + fn drain_events(&mut self) -> Vec { + self.drain_edge_events() + } +} diff --git a/crates/iroh-driver/src/lib.rs b/crates/iroh-driver/src/lib.rs index 24e66f3..a91e948 100644 --- a/crates/iroh-driver/src/lib.rs +++ b/crates/iroh-driver/src/lib.rs @@ -9,7 +9,6 @@ // work goes through `EngineHandle`. #![deny(clippy::disallowed_methods)] -pub mod driver_pumps; pub mod edge_transport; pub mod endpoint_advertisement; pub mod iroh_driver; @@ -23,7 +22,7 @@ pub use iroh_driver::{ conn_type_of, discover_lan_ips, }; -pub use edge_transport::{EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, EdgeTransportFault}; +pub use edge_transport::{EDGE_ALPN, EdgeSendHandle}; pub use telemetry_transport::{ TELEMETRY_ALPN, TelemetryQuicHeader, TelemetryQuicRead, TelemetryQuicWriteStats,