Account for handoff time when waiting for APS ACKs

crates/ziggurat-driver/src/zigbee_stack.rs

@@ -110,6 +110,10 @@ pub enum DeliveryError {
     TransmitFailed(TxResult),
     #[error("aps ack timeout")]
     ApsAckTimeout,
+    /// An APS-ack send's handoff never resolved within `aps_ack_handoff_timeout`: a
+    /// stalled sender, or a path that dropped its outcome.
+    #[error("handoff timeout")]
+    HandoffTimeout,
     #[error("broadcast passive-ack quorum not reached")]
     BroadcastQuorumNotReached,
     #[error("indirect transaction expired before {destination:?} polled")]
@@ -365,6 +369,9 @@ pub enum TxOutcome {
         slot: Arc<SendSlot>,
         stage: TrackStage,
     },
+    /// The handoff of an APS-ack unicast. Success resolves `HandOff` and resets the
+    /// pending ack's deadline; failure resolves both stages and drops the pending ack.
+    ApsAck { slot: Arc<SendSlot> },
     /// An extracted indirect transaction in flight to the radio. Success resolves the
     /// transaction's own completion; a failed transmit puts it back at the head of its
     /// queue for the device's next poll (802.15.4 spec 6.7.3), unless it has expired.
@@ -403,6 +410,10 @@ pub struct Broadcast {
 pub struct PendingApsAck {
     pub(crate) ack_data: ApsAckData,
     pub(crate) slot: Arc<SendSlot>,
+    /// How long to wait for the ack once the frame is handed off.
+    pub(crate) wait: Duration,
+    /// When the entry expires. Set to `aps_ack_handoff_timeout` at registration, then
+    /// reset to `wait` at handoff.
     pub(crate) deadline: CoreInstant,
 }
 

crates/ziggurat-driver/src/zigbee_stack/aps.rs

@@ -268,41 +268,39 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
             .with_discover_route(NwkRouteDiscovery::Enable)
             .with_radius(cmp::max(radius, 1));
 
-        // The end-to-end ack correlates on the swapped endpoints (our destination is the
-        // acker's source, and vice-versa).
-        let ack_data = (aps_ack == ApsAck::Request).then_some(ApsAckData {
-            src: destination,
-            destination_endpoint: Some(src_ep),
-            cluster_id: Some(cluster_id),
-            profile_id: Some(profile_id),
-            source_endpoint: Some(dst_ep),
-            counter: aps_seq,
-        });
-
         let (handle, slot) = SendHandle::new();
 
         // With an APS ack, the sender resolves only `handed_off` (next-hop acceptance);
-        // the ack arrival/timeout resolves `delivered` through the slot held here. Without
-        // one, next-hop acceptance is the whole verdict, so the sender resolves
-        // `Delivery` directly (rule 4 back-fills `handed_off`).
-        let stage = if ack_data.is_some() {
-            TrackStage::HandOff
-        } else {
-            TrackStage::Delivery
-        };
-
-        // An APS-ack send registers its pending ack (with the deadline the timeout
-        // reactor uses) before enqueueing so a fast reply is caught.
-        let registered_slot = ack_data.is_some().then(|| slot.clone());
-        if let Some(ack_data) = ack_data {
-            let deadline = self.core_now() + self.aps_ack_timeout(destination, sleepy_destination);
-            self.state.pending_aps_acks.lock().push(PendingApsAck {
-                ack_data,
+        // the ack arrival/timeout resolves `delivered` through the slot held here. The
+        // pending ack is registered before enqueueing so a fast reply is caught, but
+        // its ack deadline is armed only at handoff. Without an ack, next-hop acceptance
+        // is the whole verdict, so the sender resolves `Delivery` directly (rule 4
+        // back-fills `handed_off`).
+        let outcome = match aps_ack {
+            ApsAck::Request => {
+                self.state.pending_aps_acks.lock().push(PendingApsAck {
+                    // The end-to-end ack correlates on the swapped endpoints (our
+                    // destination is the acker's source, and vice-versa).
+                    ack_data: ApsAckData {
+                        src: destination,
+                        destination_endpoint: Some(src_ep),
+                        cluster_id: Some(cluster_id),
+                        profile_id: Some(profile_id),
+                        source_endpoint: Some(dst_ep),
+                        counter: aps_seq,
+                    },
+                    slot: slot.clone(),
+                    wait: self.aps_ack_timeout(destination, sleepy_destination),
+                    deadline: self.core_now() + self.tunables.aps_ack_handoff_timeout(),
+                });
+                self.aps_ack_wake.notify_one();
+                TxOutcome::ApsAck { slot: slot.clone() }
+            }
+            ApsAck::None => TxOutcome::Track {
                 slot: slot.clone(),
-                deadline,
-            });
-            self.aps_ack_wake.notify_one();
-        }
+                stage: TrackStage::Delivery,
+            },
+        };
 
         let accepted = self.send_unicast(Unicast {
             frame: nwk_frame,
@@ -312,24 +310,53 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
                 priority,
                 class: TrafficClass::Host,
             },
-            outcome: TxOutcome::Track { slot, stage },
+            outcome,
         });
 
         // A rejected frame gets no confirmation: unregister its pending ack and drop the
         // handle by returning the admission error.
         if let Err(err) = accepted {
-            if let Some(registered_slot) = registered_slot {
-                self.state
-                    .pending_aps_acks
-                    .lock()
-                    .retain(|entry| !Arc::ptr_eq(&entry.slot, &registered_slot));
-            }
+            self.remove_pending_aps_ack(&slot);
             return Err(err);
         }
 
         Ok(handle)
     }
 
+    /// An APS-ack unicast's handoff resolved. Success resets its pending ack's deadline
+    /// to `wait`; failure drops the pending ack, since the slot already holds the real
+    /// cause.
+    pub(super) fn resolve_aps_ack_handoff(
+        &self,
+        slot: Arc<SendSlot>,
+        result: Result<(), DeliveryError>,
+    ) {
+        if result.is_ok() {
+            let now = self.core_now();
+            if let Some(entry) = self
+                .state
+                .pending_aps_acks
+                .lock()
+                .iter_mut()
+                .find(|entry| Arc::ptr_eq(&entry.slot, &slot))
+            {
+                entry.deadline = now + entry.wait;
+                self.aps_ack_wake.notify_one();
+            }
+        } else {
+            self.remove_pending_aps_ack(&slot);
+        }
+
+        slot.resolve(TrackStage::HandOff, result);
+    }
+
+    fn remove_pending_aps_ack(&self, slot: &Arc<SendSlot>) {
+        self.state
+            .pending_aps_acks
+            .lock()
+            .retain(|entry| !Arc::ptr_eq(&entry.slot, slot));
+    }
+
     /// Send a broadcast APS data frame to a broadcast sink, returning a [`SendHandle`].
     /// `Err` rejects at admission; on `Ok`, `handed_off` resolves when the first copy
     /// reaches the air and `delivered` on the passive-ack quorum result. A broadcast is
@@ -471,33 +498,35 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
     }
 
     /// Expire pending APS acks: an entry past its deadline resolves `delivered` with a
-    /// timeout, and a cancelled entry with `Cancelled`. A stale entry whose send already
-    /// failed at handoff still expires here; its slot is already resolved, so the late
-    /// write is ignored (write-once, rule 1).
+    /// timeout (`HandoffTimeout` if its handoff never resolved), and a cancelled entry
+    /// with `Cancelled`.
     fn expire_aps_acks(&self) {
         let now = self.core_now();
 
-        let due: Vec<(Arc<SendSlot>, bool)> = self
+        let due: Vec<PendingApsAck> = self
             .state
             .pending_aps_acks
             .lock()
             .extract_if(.., |entry| {
                 entry.deadline <= now || entry.slot.is_cancelled()
             })
-            .map(|entry| {
-                let cancelled = entry.slot.is_cancelled();
-                (entry.slot, cancelled)
-            })
             .collect();
 
-        for (slot, cancelled) in due {
-            let result = if cancelled {
+        for entry in due {
+            let result = if entry.slot.is_cancelled() {
                 Err(DeliveryError::Cancelled)
+            } else if !entry.slot.is_handed_off() {
+                tracing::warn!(
+                    "APS ack to {:?} (counter {}) expired before its handoff resolved",
+                    entry.ack_data.src,
+                    entry.ack_data.counter
+                );
+                Err(DeliveryError::HandoffTimeout)
             } else {
                 Err(DeliveryError::ApsAckTimeout)
             };
 
-            slot.resolve(TrackStage::Delivery, result);
+            entry.slot.resolve(TrackStage::Delivery, result);
         }
     }
 }

crates/ziggurat-driver/src/zigbee_stack/nwk.rs

@@ -1164,6 +1164,7 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
                 }
             }
             TxOutcome::Track { slot, stage } => slot.resolve(stage, result),
+            TxOutcome::ApsAck { slot } => self.resolve_aps_ack_handoff(slot, result),
             TxOutcome::IndirectDelivery {
                 destination,
                 transaction,
@@ -1377,7 +1378,10 @@ impl<P: RadioPhy, R: Runtime> ZigbeeStack<P, R> {
     /// slot, and only a tracked send can be cancelled; every reactor checks this before
     /// acting on the entry it holds, dropping it and resolving `Err(Cancelled)` if set.
     fn outcome_cancelled(outcome: &TxOutcome) -> bool {
-        matches!(outcome, TxOutcome::Track { slot, .. } if slot.is_cancelled())
+        matches!(
+            outcome,
+            TxOutcome::Track { slot, .. } | TxOutcome::ApsAck { slot } if slot.is_cancelled()
+        )
     }
 
     /// Nudge every reactor that parks a cancellable send, so a just-set cancel flag is

crates/ziggurat-driver/src/zigbee_stack/send_handle.rs

@@ -70,6 +70,11 @@ impl SendSlot {
         self.cancelled.load(Ordering::Relaxed)
     }
 
+    /// Whether the send's handoff has resolved, successfully or not.
+    pub(crate) fn is_handed_off(&self) -> bool {
+        self.progress.lock().handed_off.is_some()
+    }
+
     /// Resolve one of the send's stages, applying the write rules:
     ///
     /// 1. **First write wins** — both fields are write-once; a later write is ignored.

crates/ziggurat-protocol/src/bridge.rs

@@ -429,6 +429,7 @@ pub(crate) const fn send_status(result: &Result<(), DeliveryError>) -> SendStatu
             DeliveryError::CcaFailure => SendStatus::CcaFailure,
             DeliveryError::TransmitFailed(_) => SendStatus::TransmitFailed,
             DeliveryError::ApsAckTimeout => SendStatus::ApsAckTimeout,
+            DeliveryError::HandoffTimeout => SendStatus::HandoffTimeout,
             DeliveryError::BroadcastQuorumNotReached => SendStatus::BroadcastQuorumNotReached,
             DeliveryError::IndirectExpired { .. } => SendStatus::IndirectExpired,
             DeliveryError::BudgetExhausted => SendStatus::FrameBudgetExhausted,

crates/ziggurat-protocol/src/wire.rs

@@ -696,6 +696,7 @@ pub enum SendStatus {
     FrameBudgetExhausted = 10,
     Cancelled = 11,
     RadioError = 12,
+    HandoffTimeout = 13,
 }
 
 #[abstract_bits]

crates/ziggurat-zigbee/src/constants.rs

@@ -356,6 +356,10 @@ tunables! {
     /// back.
     aps_ack_timeout_indirect: Duration = Duration::from_millis(10000),
 
+    /// How long an APS-ack send may wait for handoff before it fails. Should be longer
+    /// than route discovery plus indirect persistence.
+    aps_ack_handoff_timeout: Duration = Duration::from_secs(30),
+
     /// `macMaxCSMABackoffs`: how many times the radio backs off on a busy channel
     /// before declaring a transmit failed.
     mac_max_csma_backoffs: u8 = 2,