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
12struct StampedTick {
14 tick: Tick,
15 enqueued_at: Instant,
16 carrier: Option<TraceCarrier>,
17}
18
19#[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
37pub 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 pub async fn send(&self, tick: Tick) -> Result<(), Tick> {
52 self.send_with_trace(tick, None).await
53 }
54
55 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 pub fn try_send(&self, tick: Tick) -> Result<(), mpsc::error::TrySendError<Tick>> {
82 self.try_send_with_trace(tick, None)
83 }
84
85 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 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 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
163fn 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 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 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 #[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 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 #[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}