laminar_core/shuffle/transport/grpc/
topology.rs

1//! Install a logical generation using existing assignment publication and stream cancellation.
2
3use super::*;
4
5pub(super) fn parse_topology_fence(
6    version: u64,
7    digest: &[u8],
8) -> Result<Option<ShuffleTopologyFence>, tonic::Status> {
9    if version == 0 && digest.is_empty() {
10        return Ok(None);
11    }
12    let digest = parse_certificate_digest(digest, "topology manifest digest")?;
13    ShuffleTopologyFence::new(version, digest)
14        .map(Some)
15        .map_err(|error| tonic::Status::invalid_argument(error.to_string()))
16}
17
18impl ShuffleSender {
19    /// Logical generation retained by this endpoint, including while assignment intake is closed.
20    #[must_use]
21    pub fn topology_fence(&self) -> Option<ShuffleTopologyFence> {
22        self.assignment.read().as_ref().and_then(|a| a.topology)
23    }
24    /// Cheap graph guard. Zero denotes the legacy fabric, never an inferred topology.
25    #[must_use]
26    pub fn topology_version(&self) -> u64 {
27        self.topology_version.load(Ordering::Acquire)
28    }
29
30    /// Fence both directions of one process's fabric to an exact committed catalog generation.
31    /// The trusted caller must hold intake/execution and prove predecessor retirement and Commit.
32    /// This supplies no actor readiness, output permit or cluster Release.
33    ///
34    /// Existing assignment locks cancel prior streams, handshakes and blocked sends before starting
35    /// a new sequence domain. Identical retries preserve sequences. Process/assignment/recovery
36    /// identities and unrepaired delivery-loss fences are preserved.
37    ///
38    /// # Errors
39    /// Rejects stale/conflicting topology, incompatible endpoints, inactive assignments, expired
40    /// process leases and unrepaired delivery loss before changing either direction.
41    pub fn install_topology_fence_pair(
42        &self,
43        receiver: &ShuffleReceiver,
44        expected: Option<ShuffleTopologyFence>,
45        target: ShuffleTopologyFence,
46    ) -> io::Result<bool> {
47        self.install_topology_fence_pair_inner(receiver, expected, target, None)
48    }
49
50    /// Install held transport for an exact authorized coordinated recovery Start.
51    /// The trusted caller must prove the stopped roster, selected cut and current generation.
52    /// Loss covered by that generation's prepared cutoff may remain pending during installation;
53    /// this does not forgive it or authorize input/output. Only recovery completion promotes the
54    /// repair floor. Later loss and an exhausted incident counter still reject installation.
55    ///
56    /// # Errors
57    /// Rejects a zero/stale generation or any ordinary process, assignment and topology mismatch.
58    pub fn install_topology_fence_pair_for_recovery(
59        &self,
60        receiver: &ShuffleReceiver,
61        expected: Option<ShuffleTopologyFence>,
62        target: ShuffleTopologyFence,
63        generation: u64,
64    ) -> io::Result<bool> {
65        if generation == 0 {
66            return Err(io::Error::new(
67                io::ErrorKind::InvalidInput,
68                "topology recovery installation requires a nonzero generation",
69            ));
70        }
71        self.install_topology_fence_pair_inner(receiver, expected, target, Some(generation))
72    }
73
74    fn install_topology_fence_pair_inner(
75        &self,
76        receiver: &ShuffleReceiver,
77        expected: Option<ShuffleTopologyFence>,
78        target: ShuffleTopologyFence,
79        recovery: Option<u64>,
80    ) -> io::Result<bool> {
81        let mut sender_assignment = self.assignment.write();
82        let mut receiver_assignment = receiver.assignment.write();
83        self.process_lease.require_live_io()?;
84        receiver.process_lease.require_live_io()?;
85        let sender = sender_assignment.as_ref().ok_or_else(scope_cancelled_io)?;
86        let inbound = receiver_assignment
87            .as_ref()
88            .ok_or_else(scope_cancelled_io)?;
89        let mut delivery_peers = receiver.delivery.peers.lock();
90        if self.local_id != receiver.local_id
91            || self.sender_incarnation != receiver.receiver_incarnation
92            || sender.fence != inbound.fence
93            || sender.digest != inbound.digest
94            || sender.owners != inbound.owners
95            || sender.topology != inbound.topology
96            || self.assignment_version.load(Ordering::Acquire) != sender.fence.assignment_version
97            || receiver.assignment_version.load(Ordering::Acquire)
98                != inbound.fence.assignment_version
99            || self.recovery_gen.load(Ordering::Acquire)
100                != receiver.recovery_gen.load(Ordering::Acquire)
101            || self.scope_cancel.read().is_cancelled()
102            || receiver.scope_cancel.read().is_cancelled()
103            || recovery.is_some_and(|generation| {
104                receiver.recovery_gen.load(Ordering::Acquire) != generation
105            })
106            || (receiver.has_unrecovered_delivery_loss()
107                && recovery
108                    .is_none_or(|generation| !receiver.delivery.recovery_covers_loss(generation)))
109        {
110            return Err(io::Error::new(
111                io::ErrorKind::PermissionDenied,
112                "topology installation requires the exact live local shuffle fabric",
113            ));
114        }
115        if sender.topology == Some(target) {
116            return Ok(false);
117        }
118        if sender.topology != expected
119            || expected.is_some_and(|old| target.version() <= old.version())
120        {
121            return Err(io::Error::new(
122                io::ErrorKind::InvalidData,
123                "shuffle topology conflicts with its retained generation",
124            ));
125        }
126        let next = |old: &InstalledAssignment| {
127            Arc::new(InstalledAssignment {
128                fence: old.fence.clone(),
129                digest: old.digest,
130                owners: Arc::clone(&old.owners),
131                topology: Some(target),
132            })
133        };
134        let next_sender = next(sender);
135        let next_receiver = next(inbound);
136        rotate_scope_token(&self.scope_cancel, &self.process_lease, false);
137        rotate_scope_token(&receiver.scope_cancel, &receiver.process_lease, false);
138        self.pool.lock().clear();
139        self.seqs.lock().clear();
140        self.connect_locks.lock().clear();
141        receiver.pending_handshakes.clear();
142        receiver
143            .delivery
144            .reset_topology(target, &mut delivery_peers);
145        receiver.holdover.clear_staged_frontiers();
146        let barriers = receiver.holdover.take_staged_barriers();
147        receiver.holdover.release_items(barriers.len());
148        drop(barriers);
149        *sender_assignment = Some(next_sender);
150        *receiver_assignment = Some(next_receiver);
151        self.topology_version
152            .store(target.version(), Ordering::Release);
153        receiver
154            .topology_version
155            .store(target.version(), Ordering::Release);
156        rotate_scope_token(&self.scope_cancel, &self.process_lease, true);
157        rotate_scope_token(&receiver.scope_cancel, &receiver.process_lease, true);
158        receiver.work_ready.notify_waiters();
159        receiver.assignment_resumed.notify_waiters();
160        Ok(true)
161    }
162
163    /// Send under an immutable graph topology binding and exact routing assignment.
164    /// Existing unbound send methods target only the legacy fabric.
165    /// # Errors
166    /// Rejects stale topology before sequence allocation, as well as ordinary send failures.
167    pub async fn send_to_for_topology(
168        &self,
169        peer: ShufflePeerId,
170        expected_assignment_version: u64,
171        topology: Option<ShuffleTopologyFence>,
172        message: &ShuffleMessage,
173    ) -> io::Result<()> {
174        self.send_to_inner(
175            peer,
176            message,
177            Some(expected_assignment_version),
178            None,
179            topology,
180            None,
181        )
182        .await
183    }
184
185    /// Send under the exact routing assignment, recovery generation, and graph topology.
186    /// A retained pre-recovery plan cannot acquire a stream in a newer recovery generation.
187    ///
188    /// # Errors
189    /// Rejects a stale generation before connection or sequence allocation, plus ordinary send errors.
190    pub async fn send_to_for_generation(
191        &self,
192        peer: ShufflePeerId,
193        assignment_version: u64,
194        recovery_generation: u64,
195        topology: Option<ShuffleTopologyFence>,
196        message: &ShuffleMessage,
197    ) -> io::Result<()> {
198        self.send_to_inner(
199            peer,
200            message,
201            Some(assignment_version),
202            None,
203            topology,
204            Some(recovery_generation),
205        )
206        .await
207    }
208
209    pub(super) fn current_send_scope(
210        &self,
211        assignment: Option<u64>,
212        topology: Option<ShuffleTopologyFence>,
213        recovery: Option<u64>,
214    ) -> io::Result<ScopeLease> {
215        let scope = self.current_scope(assignment, topology)?;
216        if recovery.is_some_and(|expected| expected != scope.recovery_gen) {
217            return Err(io::Error::new(
218                io::ErrorKind::ConnectionAborted,
219                "shuffle send plan belongs to a stale recovery generation",
220            ));
221        }
222        Ok(scope)
223    }
224}
225
226pub(super) fn outbound_admission_bytes(
227    message: &ShuffleMessage,
228    assignment_fence: Option<&CheckpointAssignmentFence>,
229) -> io::Result<usize> {
230    let bytes = outbound_workspace_bytes(message)?;
231    if matches!(message, ShuffleMessage::Barrier(_)) && assignment_fence.is_none() {
232        return Err(io::Error::new(
233            io::ErrorKind::InvalidInput,
234            "shuffle checkpoint barriers require an admitted assignment certificate",
235        ));
236    }
237    Ok(bytes)
238}
239
240impl ShuffleReceiver {
241    /// Logical generation retained by this endpoint, including while assignment intake is closed.
242    #[must_use]
243    pub fn topology_fence(&self) -> Option<ShuffleTopologyFence> {
244        self.assignment.read().as_ref().and_then(|a| a.topology)
245    }
246    /// Cheap graph guard. Zero denotes the legacy fabric, never an inferred topology.
247    #[must_use]
248    pub fn topology_version(&self) -> u64 {
249        self.topology_version.load(Ordering::Acquire)
250    }
251}
252
253impl ShuffleSender {
254    pub(super) fn current_scope(
255        &self,
256        expected: Option<u64>,
257        topology: Option<ShuffleTopologyFence>,
258    ) -> io::Result<ScopeLease> {
259        let assignment = self.assignment.read();
260        let installed = assignment.as_ref().ok_or_else(|| {
261            io::Error::new(
262                io::ErrorKind::NotConnected,
263                "shuffle assignment certificate is not installed",
264            )
265        })?;
266        let version = self.assignment_version.load(Ordering::Acquire);
267        let recovery_gen = self.recovery_gen.load(Ordering::Acquire);
268        let cancel = self.scope_cancel.read().clone();
269        self.process_lease.require_live_io()?;
270        if version == 0
271            || version != installed.fence.assignment_version
272            || cancel.is_cancelled()
273            || installed.topology != topology
274        {
275            return Err(scope_cancelled_io());
276        }
277        if expected.is_some_and(|expected| expected == 0 || expected != version) {
278            return Err(io::Error::new(
279                io::ErrorKind::ConnectionAborted,
280                format!(
281                    "shuffle assignment scope mismatch: routed at {}, sender at {version}",
282                    expected.unwrap_or_default()
283                ),
284            ));
285        }
286        Ok(ScopeLease {
287            assignment: Arc::clone(installed),
288            recovery_gen,
289            cancel,
290        })
291    }
292
293    pub(super) fn validate_expected_assignment(&self, expected: Option<u64>) -> io::Result<()> {
294        let Some(expected) = expected else {
295            return Ok(());
296        };
297        let current = self.assignment_version.load(Ordering::Acquire);
298        if expected != 0 && expected == current {
299            return Ok(());
300        }
301        Err(io::Error::new(
302            io::ErrorKind::ConnectionAborted,
303            format!("shuffle assignment scope mismatch: routed at {expected}, sender at {current}"),
304        ))
305    }
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    #[test]
313    fn topology_transport_rejects_partial_malformed_and_zero_wire_identities() {
314        assert_eq!(parse_topology_fence(0, &[]).unwrap(), None);
315        for (version, digest) in [
316            (0, vec![1; 32]),
317            (1, vec![]),
318            (1, vec![0; 32]),
319            (1, vec![1; 31]),
320            (1, vec![1; 33]),
321        ] {
322            assert!(parse_topology_fence(version, &digest).is_err());
323        }
324        let target = ShuffleTopologyFence::new(2, [42; 32]).unwrap();
325        assert_eq!(parse_topology_fence(2, &[42; 32]).unwrap(), Some(target));
326        assert!(ShuffleTopologyFence::new(0, [42; 32]).is_err());
327    }
328
329    #[test]
330    fn topology_transport_late_old_admission_cannot_fault_or_advance_target_delivery() {
331        let delivery = DeliveryTracker::default();
332        let original = StreamFence {
333            sender_node_id: 1,
334            sender_incarnation: Uuid::from_u128(1),
335            receiver_incarnation: Uuid::from_u128(2),
336            stream_id: Uuid::from_u128(3),
337            assignment_version: 1,
338            assignment_certificate_digest: [1; 32],
339            recovery_gen: 0,
340            topology: None,
341        };
342        delivery.observe_hello(original).unwrap();
343        let pending = delivery.prepare_data(&original, 0).unwrap().unwrap();
344        let barrier = delivery.prepare_barrier(&original, 1).unwrap();
345        let topology = ShuffleTopologyFence::new(2, [42; 32]).unwrap();
346        delivery.reset_topology(topology, &mut delivery.peers.lock());
347        assert_eq!(
348            delivery.commit_data(pending).unwrap_err().code(),
349            tonic::Code::Cancelled
350        );
351        assert_eq!(
352            delivery.commit_barrier(barrier).unwrap_err().code(),
353            tonic::Code::Cancelled
354        );
355        assert_eq!(
356            delivery.validate_stream(&original).unwrap_err().code(),
357            tonic::Code::Cancelled
358        );
359        assert_eq!(
360            delivery.prepare_barrier(&original, 1).err().unwrap().code(),
361            tonic::Code::Cancelled
362        );
363        let current = StreamFence {
364            topology: Some(topology),
365            stream_id: Uuid::from_u128(4),
366            ..original
367        };
368        delivery.observe_hello(current).unwrap();
369        let pending = delivery.prepare_data(&current, 0).unwrap().unwrap();
370        delivery.commit_data(pending).unwrap();
371        assert_eq!(delivery.delivery_loss_incidents.load(Ordering::Acquire), 0);
372        assert_eq!(delivery.peers.lock()[&1].expected, 1);
373    }
374}
laminar_core/shuffle/transport/grpc/topology.rs

laminar_core/shuffle/transport/grpc/
topology.rs
