Skip to main content

helix_driver_host/
tick_ingress.rs

1use std::sync::Arc;
2use std::time::{Duration, Instant};
3
4use helix_core::Tick;
5use tokio::sync::mpsc;
6
7use crate::metrics::{
8    AsyncMetricSink, LabelKey, MetricEvent, MetricId, MetricLabels, NoopMetricSink,
9};
10use crate::trace::TraceCarrier;
11
12/// 将 Tick、入队等待起点和可选 trace carrier 作为一个不可拆分的入站元素。
13struct StampedTick {
14    tick: Tick,
15    enqueued_at: Instant,
16    carrier: Option<TraceCarrier>,
17}
18
19/// 生产 tick 入站发送端:时间戳、Tick 和可选 carrier 同在一个有界 mpsc 元素内。
20///
21/// `reserve` 成功后才取时间戳,因此等待值严格表示“实际入队 → engine 消费”;
22/// 不增加旁路队列,也不改变原 channel 容量和背压点。
23#[derive(Clone)]
24pub struct TickIngressSender {
25    tx: mpsc::Sender<StampedTick>,
26    metrics: Arc<dyn AsyncMetricSink>,
27}
28
29pub struct TickIngressReceiver {
30    rx: mpsc::Receiver<StampedTick>,
31}
32
33pub fn tick_ingress_channel(capacity: usize) -> (TickIngressSender, TickIngressReceiver) {
34    tick_ingress_channel_observed(capacity, Arc::new(NoopMetricSink))
35}
36
37/// 创建带 command admission 指标的生产 channel,容量与原背压点完全一致。
38pub fn tick_ingress_channel_observed(
39    capacity: usize,
40    metrics: Arc<dyn AsyncMetricSink>,
41) -> (TickIngressSender, TickIngressReceiver) {
42    let (tx, rx) = mpsc::channel(capacity.max(1));
43    (
44        TickIngressSender { tx, metrics },
45        TickIngressReceiver { rx },
46    )
47}
48
49impl TickIngressSender {
50    /// 兼容无 trace carrier 的业务入站,统一复用带上下文的发送路径。
51    pub async fn send(&self, tick: Tick) -> Result<(), Tick> {
52        self.send_with_trace(tick, None).await
53    }
54
55    /// 将可选 WS trace carrier 与 Tick 原子放入同一个有界元素,避免旁路队列错配。
56    pub async fn send_with_trace(
57        &self,
58        tick: Tick,
59        carrier: Option<TraceCarrier>,
60    ) -> Result<(), Tick> {
61        let started = self.metrics.is_enabled().then(Instant::now);
62        let kind = tick_kind(&tick);
63        let is_command = matches!(tick, Tick::Command(_));
64        let permit = match self.tx.reserve().await {
65            Ok(permit) => permit,
66            Err(_) => {
67                self.record_rejected(kind, is_command, "closed");
68                return Err(tick);
69            }
70        };
71        permit.send(StampedTick {
72            tick,
73            enqueued_at: Instant::now(),
74            carrier,
75        });
76        self.record_accepted(kind, is_command, started);
77        Ok(())
78    }
79
80    /// 兼容无 trace carrier 的非阻塞入站,保持原有队列语义。
81    pub fn try_send(&self, tick: Tick) -> Result<(), mpsc::error::TrySendError<Tick>> {
82        self.try_send_with_trace(tick, None)
83    }
84
85    /// 非阻塞版本的带 carrier 入站;队满时 carrier 随 Tick 一起返回,绝不泄漏。
86    pub(crate) fn try_send_with_trace(
87        &self,
88        tick: Tick,
89        carrier: Option<TraceCarrier>,
90    ) -> Result<(), mpsc::error::TrySendError<Tick>> {
91        let started = self.metrics.is_enabled().then(Instant::now);
92        let kind = tick_kind(&tick);
93        let is_command = matches!(tick, Tick::Command(_));
94        let permit = match self.tx.try_reserve() {
95            Ok(permit) => permit,
96            Err(mpsc::error::TrySendError::Full(())) => {
97                self.record_rejected(kind, is_command, "full");
98                return Err(mpsc::error::TrySendError::Full(tick));
99            }
100            Err(mpsc::error::TrySendError::Closed(())) => {
101                self.record_rejected(kind, is_command, "closed");
102                return Err(mpsc::error::TrySendError::Closed(tick));
103            }
104        };
105        permit.send(StampedTick {
106            tick,
107            enqueued_at: Instant::now(),
108            carrier,
109        });
110        self.record_accepted(kind, is_command, started);
111        Ok(())
112    }
113
114    /// 记录真正成功 reserve+send 的 admission,不把 engine 消费误当接纳。
115    fn record_accepted(&self, kind: &'static str, is_command: bool, started: Option<Instant>) {
116        if !self.metrics.is_enabled() {
117            return;
118        }
119        let labels = MetricLabels::one(LabelKey::Stage, "ingress")
120            .with(LabelKey::TickKind, kind)
121            .with(LabelKey::Status, "accepted");
122        if is_command {
123            let _ = self.metrics.try_record(MetricEvent::counter(
124                MetricId::CommandAcceptedTotal,
125                1.0,
126                labels,
127            ));
128            if let Some(started) = started {
129                let _ = self.metrics.try_record(MetricEvent::histogram(
130                    MetricId::CommandAdmissionDurationSeconds,
131                    started.elapsed().as_secs_f64(),
132                    labels,
133                ));
134            }
135        }
136    }
137
138    /// 记录 full/closed 拒绝和队列满分母,保持 try_send 非阻塞。
139    fn record_rejected(&self, kind: &'static str, is_command: bool, reason: &'static str) {
140        if !self.metrics.is_enabled() {
141            return;
142        }
143        let labels = MetricLabels::one(LabelKey::Stage, "ingress")
144            .with(LabelKey::TickKind, kind)
145            .with(LabelKey::Status, reason);
146        if reason == "full" {
147            let _ = self.metrics.try_record(MetricEvent::counter(
148                MetricId::TickQueueFullTotal,
149                1.0,
150                labels,
151            ));
152        }
153        if is_command {
154            let _ = self.metrics.try_record(MetricEvent::counter(
155                MetricId::CommandRejectedTotal,
156                1.0,
157                labels,
158            ));
159        }
160    }
161}
162
163/// 将 Tick 映射为固定低基数 admission 标签。
164fn tick_kind(tick: &Tick) -> &'static str {
165    match tick {
166        Tick::Inbound(_) => "inbound",
167        Tick::PortReply { .. } => "port_reply",
168        Tick::PortProgress { .. } => "port_progress",
169        Tick::Timer(_) => "timer",
170        Tick::Command(_) => "command",
171        Tick::Connected(_) => "connected",
172        Tick::Disconnected(_) => "disconnected",
173    }
174}
175
176impl TickIngressReceiver {
177    /// 接收一个带入队等待时间和可选 trace carrier 的 stamped Tick。
178    pub async fn recv(&mut self) -> Option<(Tick, Duration, Option<TraceCarrier>)> {
179        self.rx
180            .recv()
181            .await
182            .map(|item| (item.tick, item.enqueued_at.elapsed(), item.carrier))
183    }
184
185    pub(crate) fn len(&self) -> usize {
186        self.rx.len()
187    }
188
189    /// 非阻塞接收 stamped Tick,并原子取出同一元素上的 carrier。
190    pub fn try_recv(
191        &mut self,
192    ) -> Result<(Tick, Duration, Option<TraceCarrier>), mpsc::error::TryRecvError> {
193        self.rx
194            .try_recv()
195            .map(|item| (item.tick, item.enqueued_at.elapsed(), item.carrier))
196    }
197
198    pub(crate) fn max_capacity(&self) -> usize {
199        self.rx.max_capacity()
200    }
201}
202
203pub(crate) enum EngineTickSender {
204    Raw(mpsc::Sender<Tick>),
205    Stamped(TickIngressSender),
206}
207
208impl Clone for EngineTickSender {
209    fn clone(&self) -> Self {
210        match self {
211            Self::Raw(tx) => Self::Raw(tx.clone()),
212            Self::Stamped(tx) => Self::Stamped(tx.clone()),
213        }
214    }
215}
216
217impl EngineTickSender {
218    pub(crate) async fn send(&self, tick: Tick) -> Result<(), Tick> {
219        match self {
220            Self::Raw(tx) => tx.send(tick).await.map_err(|error| error.0),
221            Self::Stamped(tx) => tx.send(tick).await,
222        }
223    }
224}
225
226pub(crate) enum EngineTickReceiver {
227    Raw(mpsc::Receiver<Tick>),
228    Stamped(TickIngressReceiver),
229}
230
231impl EngineTickReceiver {
232    pub(crate) async fn recv(&mut self) -> Option<(Tick, Option<Duration>, Option<TraceCarrier>)> {
233        match self {
234            Self::Raw(rx) => rx.recv().await.map(|tick| (tick, None, None)),
235            Self::Stamped(rx) => rx
236                .recv()
237                .await
238                .map(|(tick, wait, carrier)| (tick, Some(wait), carrier)),
239        }
240    }
241
242    pub(crate) fn len(&self) -> usize {
243        match self {
244            Self::Raw(rx) => rx.len(),
245            Self::Stamped(rx) => rx.len(),
246        }
247    }
248
249    pub(crate) fn try_recv(
250        &mut self,
251    ) -> Result<(Tick, Option<Duration>, Option<TraceCarrier>), mpsc::error::TryRecvError> {
252        match self {
253            Self::Raw(rx) => rx.try_recv().map(|tick| (tick, None, None)),
254            Self::Stamped(rx) => rx
255                .try_recv()
256                .map(|(tick, wait, carrier)| (tick, Some(wait), carrier)),
257        }
258    }
259
260    pub(crate) fn max_capacity(&self) -> usize {
261        match self {
262            Self::Raw(rx) => rx.max_capacity(),
263            Self::Stamped(rx) => rx.max_capacity(),
264        }
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271    use crate::metrics::RecordOutcome;
272    use crate::trace::TraceCarrier;
273    use helix_core::effect::TimerId;
274    use helix_core::tick::AppCommand;
275    use std::sync::Mutex;
276
277    /// WS 入站 carrier 必须和 Tick 在同一有界队列元素中保持顺序。
278    #[tokio::test]
279    async fn stamped_ingress_preserves_ws_trace_carrier() {
280        let (tx, mut rx) = tick_ingress_channel(1);
281        let carrier = TraceCarrier::from_json_str(
282            r#"{"traceparent":"00-00000000000000000000000000000001-0000000000000002-01"}"#,
283        )
284        .expect("carrier");
285
286        tx.send_with_trace(
287            Tick::Inbound(helix_core::tick::InboundBytes::from_static(b"{}")),
288            Some(carrier.clone()),
289        )
290        .await
291        .expect("ingress send");
292
293        let (_, _, received) = rx.recv().await.expect("stamped tick");
294        assert_eq!(received, Some(carrier));
295    }
296
297    #[derive(Default)]
298    struct RecordingMetricSink(Mutex<Vec<MetricEvent>>);
299
300    impl AsyncMetricSink for RecordingMetricSink {
301        /// 保存 admission 指标事件,不启动真实 exporter。
302        fn try_record(&self, event: MetricEvent) -> RecordOutcome {
303            self.0.lock().unwrap().push(event);
304            RecordOutcome::Accepted
305        }
306    }
307
308    #[tokio::test]
309    async fn timestamp_lives_in_same_bounded_queue_element() {
310        let (tx, mut rx) = tick_ingress_channel(1);
311        tx.send(Tick::Timer(TimerId::from_raw(1)))
312            .await
313            .expect("first tick enters queue");
314        assert!(matches!(
315            tx.try_send(Tick::Timer(TimerId::from_raw(2))),
316            Err(mpsc::error::TrySendError::Full(_))
317        ));
318        tokio::time::sleep(Duration::from_millis(5)).await;
319        let (_, wait, _) = rx.recv().await.expect("stamped tick");
320        assert!(wait >= Duration::from_millis(5));
321        assert_eq!(rx.max_capacity(), 1, "stamping 不得引入第二缓冲区");
322    }
323
324    /// command 的成功接纳与队满拒绝必须在真正 channel admission 点计数。
325    #[tokio::test]
326    async fn observed_channel_records_command_admission_and_rejection() {
327        let metrics = Arc::new(RecordingMetricSink::default());
328        let (tx, _rx) = tick_ingress_channel_observed(1, metrics.clone());
329        tx.try_send(Tick::Command(AppCommand::new("first", bytes::Bytes::new())))
330            .expect("第一条 command 应成功入队");
331        assert!(matches!(
332            tx.try_send(Tick::Command(AppCommand::new(
333                "second",
334                bytes::Bytes::new()
335            ))),
336            Err(mpsc::error::TrySendError::Full(_))
337        ));
338
339        let events = metrics.0.lock().unwrap();
340        for expected in [
341            MetricId::CommandAcceptedTotal,
342            MetricId::CommandAdmissionDurationSeconds,
343            MetricId::CommandRejectedTotal,
344            MetricId::TickQueueFullTotal,
345        ] {
346            assert!(
347                events.iter().any(|event| event.id == expected),
348                "admission 缺少 {:?}",
349                expected
350            );
351        }
352    }
353}