1use crate::control::queries;
4use crate::control::{ControlMessage, ControlSenders, ControlSocket, commands};
5use crate::node::{
6 EndpointDataBatchRx, EndpointEventSender, Node, NodeError, endpoint_data_batch_channel,
7};
8use crate::transport::PacketRx;
9use crate::upper::tun::TunOutboundRx;
10use std::time::{Duration, Instant};
11use tokio::sync::mpsc::Receiver;
12use tracing::{debug, info, warn};
13
14mod budget;
15mod drain;
16mod packet_mover2;
17
18#[cfg(test)]
19mod tests;
20
21use budget::*;
22use drain::*;
23
24impl Node {
25 pub async fn run_rx_loop(&mut self) -> Result<(), NodeError> {
45 let mut packet_rx = self.packet_rx.take().ok_or(NodeError::NotStarted)?;
46
47 let (mut tun_outbound_rx, _tun_guard) = match self.tun_outbound_rx.take() {
51 Some(rx) => (rx, None),
52 None => {
53 let (tx, rx) = crate::upper::tun::tun_outbound_channel(1);
54 (rx, Some(tx))
55 }
56 };
57
58 let (mut dns_identity_rx, _dns_guard) = match self.dns_identity_rx.take() {
61 Some(rx) => (rx, None),
62 None => {
63 let (tx, rx) = tokio::sync::mpsc::channel(1);
64 (rx, Some(tx))
65 }
66 };
67
68 let (mut endpoint_control_rx, _endpoint_control_guard) =
71 match self.endpoint_control_rx.take() {
72 Some(rx) => (rx, None),
73 None => {
74 let (tx, rx) = tokio::sync::mpsc::channel(1);
75 (rx, Some(tx))
76 }
77 };
78 let (mut endpoint_data_rx, _endpoint_data_guard) = match self.endpoint_data_rx.take() {
79 Some(rx) => (rx, None),
80 None => {
81 let (tx, rx) = endpoint_data_batch_channel(1);
82 (rx, Some(tx))
83 }
84 };
85 let (mut packet_mover2_fast_ingress_rx, _packet_mover2_fast_ingress_guard) =
86 match self.packet_mover2_fast_ingress_rx.take() {
87 Some(rx) => (rx, None),
88 None => {
89 let (tx, rx) = tokio::sync::mpsc::channel(1);
90 (rx, Some(tx))
91 }
92 };
93
94 let mut tick =
95 tokio::time::interval(Duration::from_secs(self.config.node.tick_interval_secs));
96 tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
97 let mut maintenance_state = RxLoopMaintenanceState::default();
98
99 let (control_query_tx, mut control_query_rx) =
103 tokio::sync::mpsc::channel::<ControlMessage>(32);
104 let (control_command_tx, mut control_command_rx) =
105 tokio::sync::mpsc::channel::<ControlMessage>(32);
106
107 if self.config.node.control.enabled {
108 let config = self.config.node.control.clone();
109 let senders = ControlSenders::new(control_query_tx.clone(), control_command_tx.clone());
110 tokio::spawn(async move {
111 match ControlSocket::bind(&config) {
112 Ok(socket) => {
113 socket.accept_loop_split(senders).await;
114 }
115 Err(e) => {
116 warn!(error = %e, "Failed to bind control socket");
117 }
118 }
119 });
120 }
121 drop(control_query_tx);
123 drop(control_command_tx);
124
125 let packet_mover2_tun_tx = self.tun_tx.clone().unwrap_or_else(|| {
126 let (tx, rx) = crate::upper::tun::write_channel();
127 drop(rx);
128 tx
129 });
130 let packet_mover2_endpoint_tx = self.endpoint_events.sender().unwrap_or_else(|| {
131 let (tx, rx) = EndpointEventSender::channel(1);
132 drop(rx);
133 tx
134 });
135 let packet_mover2_completion_notify = self.packet_mover2.completion_notify();
136
137 info!("RX event loop started");
138 crate::perf_profile::maybe_spawn_reporter();
140 tick.tick().await;
144
145 loop {
146 tokio::select! {
147 biased;
148 _ = tick.tick() => {
154 let drained = self.drain_rx_loop_data_queues(
155 &mut packet_rx,
156 &mut packet_mover2_fast_ingress_rx,
157 &mut tun_outbound_rx,
158 &mut endpoint_data_rx,
159 &packet_mover2_tun_tx,
160 &packet_mover2_endpoint_tx,
161 ENDPOINT_DRAIN_BUDGET,
162 ).await;
163 if drained.has_drained() {
164 maintenance_state.record_data_activity(Instant::now());
165 debug!(
166 drained = drained.total(),
167 drained_packets = drained.packets,
168 drained_tun = drained.tun,
169 drained_endpoint = drained.endpoint,
170 "Drained queued packets before rx-loop maintenance"
171 );
172 }
173 let maintenance_plan = maintenance_state.plan_maintenance(
174 drained,
175 Instant::now(),
176 RX_LOOP_RECENT_DATA_ACTIVITY_WINDOW,
177 RX_LOOP_SLOW_MAINTENANCE_IDLE_TIMEOUT,
178 RX_LOOP_SLOW_MAINTENANCE_BUSY_TIMEOUT,
179 );
180
181 let slow_timed_out = self.run_rx_loop_maintenance_tick(
182 maintenance_plan,
183 ).await;
184 maintenance_state.record_maintenance_result(
185 maintenance_plan.data_pressure(),
186 slow_timed_out,
187 );
188
189 let post_drained = self.drain_rx_loop_data_queues(
190 &mut packet_rx,
191 &mut packet_mover2_fast_ingress_rx,
192 &mut tun_outbound_rx,
193 &mut endpoint_data_rx,
194 &packet_mover2_tun_tx,
195 &packet_mover2_endpoint_tx,
196 PACKET_DRAIN_BUDGET,
197 ).await;
198 if post_drained.has_drained() {
199 maintenance_state.record_data_activity(Instant::now());
200 debug!(
201 drained = post_drained.total(),
202 drained_packets = post_drained.packets,
203 drained_tun = post_drained.tun,
204 drained_endpoint = post_drained.endpoint,
205 "Drained queued packets after rx-loop maintenance"
206 );
207 }
208 }
209 Some(message) = control_query_rx.recv() => {
210 self.drain_control_queries(
211 &mut control_query_rx,
212 Some(message),
213 ENDPOINT_DRAIN_BUDGET,
214 ).await;
215 }
216 Some(command) = endpoint_control_rx.recv() => {
221 self.handle_endpoint_control(command).await;
222 }
223 packet = packet_rx.recv() => {
224 match packet {
225 Some(p) => {
226 let latency_packet = p.is_transport_priority();
227 let mut firsts = crate::packet_mover2::PacketMover2LiveTurnFirsts {
228 raw_packet: Some(p),
229 ..Default::default()
230 };
231 if let Ok(packet) = tun_outbound_rx.try_recv() {
232 firsts.tun_packet = Some(packet);
233 }
234 let latency_work_ready = latency_packet
235 || packet_rx.priority_ready_packets() > 0;
236 if latency_work_ready {
237 let packet_budget = packet_drain_budget(true);
238 let endpoint_budget = endpoint_drain_budget(packet_budget);
239 let tun_budget = tun_drain_budget(packet_budget);
240 let crypto_budget = mixed_dataplane_crypto_budget(
241 packet_budget,
242 endpoint_budget,
243 tun_budget,
244 );
245 let mut turn = self.drain_packet_mover2_turn_with_firsts(
246 &mut packet_rx,
247 firsts,
248 packet_budget,
249 &mut endpoint_data_rx,
250 endpoint_budget,
251 &mut tun_outbound_rx,
252 tun_budget,
253 &packet_mover2_tun_tx,
254 &packet_mover2_endpoint_tx,
255 crypto_budget,
256 ).await;
257 self.finish_packet_mover2_turn(
258 &mut turn,
259 &mut maintenance_state,
260 &mut control_query_rx,
261 CONTROL_QUERY_INTERLEAVE_BUDGET,
262 ).await;
263 } else {
264 firsts.raw_ingress_prefetch = true;
265 self.service_packet_mover2_bulk_turns(
266 &mut packet_rx,
267 &mut packet_mover2_fast_ingress_rx,
268 firsts,
269 &mut endpoint_data_rx,
270 &mut tun_outbound_rx,
271 &packet_mover2_tun_tx,
272 &packet_mover2_endpoint_tx,
273 &mut maintenance_state,
274 &mut control_query_rx,
275 ).await;
276 }
277 }
278 None => break, }
280 }
281 Some(fast_ingress) = packet_mover2_fast_ingress_rx.recv() => {
282 self.service_packet_mover2_bulk_turns(
283 &mut packet_rx,
284 &mut packet_mover2_fast_ingress_rx,
285 crate::packet_mover2::PacketMover2LiveTurnFirsts {
286 fast_ingress: Some(fast_ingress),
287 ..Default::default()
288 },
289 &mut endpoint_data_rx,
290 &mut tun_outbound_rx,
291 &packet_mover2_tun_tx,
292 &packet_mover2_endpoint_tx,
293 &mut maintenance_state,
294 &mut control_query_rx,
295 ).await;
296 }
297 _ = packet_mover2_completion_notify.notified() => {
298 let mut turn = self.drain_packet_mover2_completion_turn(
299 &mut endpoint_data_rx,
300 &mut tun_outbound_rx,
301 &packet_mover2_tun_tx,
302 &packet_mover2_endpoint_tx,
303 LATENCY_PACKET_DRAIN_BUDGET,
304 ).await;
305 self.finish_packet_mover2_turn(
306 &mut turn,
307 &mut maintenance_state,
308 &mut control_query_rx,
309 0,
310 ).await;
311 }
312 Some(ipv6_packet) = tun_outbound_rx.recv() => {
313 let tun_budget = tun_drain_budget(LATENCY_PACKET_DRAIN_BUDGET);
314 let mut turn = self.drain_packet_mover2_turn_with_firsts(
315 &mut packet_rx,
316 crate::packet_mover2::PacketMover2LiveTurnFirsts {
317 tun_packet: Some(ipv6_packet),
318 ..Default::default()
319 },
320 0,
321 &mut endpoint_data_rx,
322 0,
323 &mut tun_outbound_rx,
324 tun_budget,
325 &packet_mover2_tun_tx,
326 &packet_mover2_endpoint_tx,
327 tun_budget,
328 ).await;
329 self.finish_packet_mover2_turn(
330 &mut turn,
331 &mut maintenance_state,
332 &mut control_query_rx,
333 0,
334 ).await;
335 }
336 Some(identity) = dns_identity_rx.recv() => {
337 debug!(
338 node_addr = %identity.node_addr,
339 "Registering identity from DNS resolution"
340 );
341 self.register_identity(identity.node_addr, identity.pubkey);
342 }
343 Some(batch) = endpoint_data_rx.recv() => {
344 let mut turn = self.drain_packet_mover2_turn_with_firsts(
345 &mut packet_rx,
346 crate::packet_mover2::PacketMover2LiveTurnFirsts {
347 endpoint_data_batch: Some(batch),
348 ..Default::default()
349 },
350 0,
351 &mut endpoint_data_rx,
352 ENDPOINT_DRAIN_BUDGET,
353 &mut tun_outbound_rx,
354 0,
355 &packet_mover2_tun_tx,
356 &packet_mover2_endpoint_tx,
357 PACKET_DRAIN_BUDGET,
358 ).await;
359 self.finish_packet_mover2_turn(
360 &mut turn,
361 &mut maintenance_state,
362 &mut control_query_rx,
363 0,
364 ).await;
365 }
366 Some((request, response_tx)) = control_command_rx.recv() => {
367 let response = commands::dispatch(
368 self,
369 &request.command,
370 request.params.as_ref(),
371 ).await;
372 let _ = response_tx.send(response);
373 }
374 }
375 }
376
377 info!("RX event loop stopped (channel closed)");
378 Ok(())
379 }
380
381 async fn drain_rx_loop_data_queues(
382 &mut self,
383 packet_rx: &mut PacketRx,
384 packet_mover2_fast_ingress_rx: &mut crate::packet_mover2::PacketMover2FastIngressRx,
385 tun_outbound_rx: &mut TunOutboundRx,
386 endpoint_data_rx: &mut EndpointDataBatchRx,
387 tun_tx: &crate::upper::tun::TunTx,
388 endpoint_tx: &EndpointEventSender,
389 budget: usize,
390 ) -> RxLoopDataDrainStats {
391 let fast_ingress =
392 Self::take_packet_mover2_fast_ingress_batch(packet_mover2_fast_ingress_rx, budget);
393 let packet_budget = budget.max(
394 fast_ingress
395 .as_ref()
396 .map_or(0, |fast_ingress| fast_ingress.len()),
397 );
398 let endpoint_budget = endpoint_drain_budget(packet_budget);
399 let tun_budget = tun_drain_budget(packet_budget);
400 let crypto_budget =
401 mixed_dataplane_crypto_budget(packet_budget, endpoint_budget, tun_budget);
402 let mut turn = self
403 .drain_packet_mover2_turn_with_firsts(
404 packet_rx,
405 crate::packet_mover2::PacketMover2LiveTurnFirsts {
406 fast_ingress,
407 ..Default::default()
408 },
409 packet_budget,
410 endpoint_data_rx,
411 endpoint_budget,
412 tun_outbound_rx,
413 tun_budget,
414 tun_tx,
415 endpoint_tx,
416 crypto_budget,
417 )
418 .await;
419 let drained_packets = Self::packet_mover2_packet_activity(&turn);
420 let control_drained = self.process_packet_mover2_control_ingress(&mut turn).await;
421 RxLoopDataDrainStats::new(
422 drained_packets,
423 turn.tun_source_drained(),
424 turn.endpoint_source_drained(),
425 control_drained,
426 )
427 }
428
429 fn take_packet_mover2_fast_ingress_batch(
430 packet_mover2_fast_ingress_rx: &mut crate::packet_mover2::PacketMover2FastIngressRx,
431 limit: usize,
432 ) -> Option<crate::packet_mover2::PacketMover2FastIngressBatch> {
433 let fast_ingress = packet_mover2_fast_ingress_rx.try_recv().ok()?;
434 Some(Self::coalesce_packet_mover2_fast_ingress(
435 fast_ingress,
436 packet_mover2_fast_ingress_rx,
437 limit,
438 ))
439 }
440
441 fn coalesce_packet_mover2_fast_ingress(
442 mut fast_ingress: crate::packet_mover2::PacketMover2FastIngressBatch,
443 packet_mover2_fast_ingress_rx: &mut crate::packet_mover2::PacketMover2FastIngressRx,
444 limit: usize,
445 ) -> crate::packet_mover2::PacketMover2FastIngressBatch {
446 while fast_ingress.len() < limit {
447 let Ok(next) = packet_mover2_fast_ingress_rx.try_recv() else {
448 break;
449 };
450 fast_ingress.absorb(next);
451 }
452 fast_ingress
453 }
454
455 #[allow(clippy::too_many_arguments)]
456 async fn service_packet_mover2_bulk_turns(
457 &mut self,
458 packet_rx: &mut PacketRx,
459 packet_mover2_fast_ingress_rx: &mut crate::packet_mover2::PacketMover2FastIngressRx,
460 firsts: crate::packet_mover2::PacketMover2LiveTurnFirsts,
461 endpoint_data_rx: &mut EndpointDataBatchRx,
462 tun_outbound_rx: &mut TunOutboundRx,
463 tun_tx: &crate::upper::tun::TunTx,
464 endpoint_tx: &EndpointEventSender,
465 maintenance_state: &mut RxLoopMaintenanceState,
466 control_query_rx: &mut Receiver<ControlMessage>,
467 ) {
468 let started = Instant::now();
469 let mut firsts = Some(firsts);
470 let mut turns = 0usize;
471
472 loop {
473 if turns > 0 {
474 if turns >= RX_LOOP_BULK_SERVICE_MAX_TURNS
475 || started.elapsed() >= RX_LOOP_BULK_SERVICE_MAX_ELAPSED
476 || packet_rx.priority_ready_packets() > 0
477 {
478 break;
479 }
480 }
481
482 let packet_budget = PACKET_DRAIN_BUDGET;
483 let mut turn_firsts = firsts.take().unwrap_or_default();
484 turn_firsts.raw_ingress_prefetch = true;
485 turn_firsts.fast_ingress = match turn_firsts.fast_ingress.take() {
486 Some(fast_ingress) => Some(Self::coalesce_packet_mover2_fast_ingress(
487 fast_ingress,
488 packet_mover2_fast_ingress_rx,
489 packet_budget,
490 )),
491 None => Self::take_packet_mover2_fast_ingress_batch(
492 packet_mover2_fast_ingress_rx,
493 packet_budget,
494 ),
495 };
496 let packet_budget = packet_budget.max(
497 turn_firsts
498 .fast_ingress
499 .as_ref()
500 .map_or(0, |fast_ingress| fast_ingress.len()),
501 );
502 let endpoint_budget = endpoint_drain_budget(packet_budget);
503 let tun_budget = tun_drain_budget(packet_budget);
504 let crypto_budget =
505 mixed_dataplane_crypto_budget(packet_budget, endpoint_budget, tun_budget);
506
507 let mut turn = self
508 .drain_packet_mover2_turn_with_firsts(
509 packet_rx,
510 turn_firsts,
511 packet_budget,
512 endpoint_data_rx,
513 endpoint_budget,
514 tun_outbound_rx,
515 tun_budget,
516 tun_tx,
517 endpoint_tx,
518 crypto_budget,
519 )
520 .await;
521 let raw_drained = Self::packet_mover2_raw_ingress_activity(&turn);
522 let control_activity = Self::packet_mover2_control_activity(&turn);
523 let completions_drained = turn.summary().completions();
524 let keep_servicing = raw_drained >= packet_budget
525 || completions_drained >= crypto_budget
526 || turn.tun_source_drained() >= tun_budget
527 || turn.endpoint_source_drained() >= endpoint_budget;
528 let control_drained = self
529 .finish_packet_mover2_turn(
530 &mut turn,
531 maintenance_state,
532 control_query_rx,
533 CONTROL_QUERY_INTERLEAVE_BUDGET,
534 )
535 .await;
536 turns += 1;
537
538 if !keep_servicing || control_activity > 0 || control_drained > 0 {
539 break;
540 }
541 }
542 }
543
544 async fn finish_packet_mover2_turn(
545 &mut self,
546 turn: &mut crate::packet_mover2::PacketMover2LiveNodeTurn,
547 maintenance_state: &mut RxLoopMaintenanceState,
548 control_query_rx: &mut Receiver<ControlMessage>,
549 control_query_budget: usize,
550 ) -> usize {
551 let had_activity = turn.has_activity();
552 let control_drained = self.process_packet_mover2_control_ingress(turn).await;
553 if control_drained > 0 && self.packet_mover2.has_deferred_raw_ingress() {
554 self.packet_mover2.completion_notify().notify_one();
555 }
556 let query_drained = if control_query_budget > 0 {
557 self.drain_control_queries(control_query_rx, None, control_query_budget)
558 .await
559 } else {
560 0
561 };
562 if had_activity || control_drained > 0 {
563 maintenance_state.record_data_activity(Instant::now());
564 }
565 control_drained.saturating_add(query_drained)
566 }
567
568 async fn drain_control_queries(
569 &mut self,
570 control_query_rx: &mut Receiver<ControlMessage>,
571 first_message: Option<ControlMessage>,
572 budget: usize,
573 ) -> usize {
574 let mut drain = SingleLaneDrainCursor::new(first_message, budget);
575 while let Some((request, response_tx)) = drain.next(control_query_rx) {
576 let response = queries::dispatch(self, &request.command, request.params.as_ref());
577 let _ = response_tx.send(response);
578 }
579
580 drain.drained()
581 }
582
583 async fn run_rx_loop_maintenance_tick(&mut self, plan: RxLoopMaintenancePlan) -> bool {
584 self.check_timeouts();
585 let now_ms = Self::now_ms();
586 self.check_link_heartbeats().await;
590 self.reload_peer_acl();
591 self.resend_pending_handshakes(now_ms).await;
592 self.resend_pending_rekeys(now_ms).await;
593 self.resend_pending_session_handshakes(now_ms).await;
594 self.resend_pending_session_msg3(now_ms).await;
595 self.purge_idle_sessions(now_ms);
596 self.purge_learned_routes(now_ms);
597 self.check_mmp_reports().await;
598 self.check_session_mmp_reports().await;
599 self.check_rekey().await;
600 self.check_session_rekey().await;
601 self.check_pending_lookups(now_ms).await;
602 self.poll_pending_connects().await;
603 self.process_pending_retries(now_ms).await;
604 self.poll_transport_discovery().await;
605 self.sample_transport_congestion();
606
607 let Some(slow_timeout) = plan.slow_timeout() else {
608 crate::perf_profile::record_event(
609 crate::perf_profile::Event::RxLoopSlowMaintenanceSkipped,
610 );
611 return false;
612 };
613
614 if tokio::time::timeout(slow_timeout, self.run_rx_loop_slow_maintenance_tick())
615 .await
616 .is_err()
617 {
618 crate::perf_profile::record_event(
619 crate::perf_profile::Event::RxLoopSlowMaintenanceTimeout,
620 );
621 self.mark_rx_loop_maintenance_timeout();
622 warn!(
623 timeout_ms = slow_timeout.as_millis() as u64,
624 data_pressure = plan.data_pressure(),
625 "RX loop slow maintenance timed out; continuing packet processing"
626 );
627 return true;
628 }
629 false
630 }
631
632 async fn run_rx_loop_slow_maintenance_tick(&mut self) {
633 if let Some(delay) = rx_loop_slow_maintenance_fault_delay() {
634 tokio::time::sleep(delay).await;
635 }
636
637 self.poll_nostr_discovery().await;
642 self.poll_lan_discovery().await;
643 self.poll_local_instance_discovery().await;
644 self.check_tree_state().await;
645 self.check_bloom_state().await;
646 self.compute_mesh_size();
647 self.record_stats_history();
648 }
649}