tuwunel_service/sending/sender/
netburst.rs

1use std::{collections::BTreeMap, time::Duration};
2
3use futures::StreamExt;
4use tuwunel_core::{
5	implement,
6	itertools::Itertools,
7	utils::{IterStream, ReadyExt, rand::secs as rand_secs, stream::TryIgnore},
8	warn,
9};
10
11use super::{
12	SendingFutures, TransactionStatus, TransactionStatuses, WakeQueue,
13	wake::{arm_wake, arm_wake_in, is_armed},
14};
15use crate::{
16	federation::ShouldAttempt,
17	sending::{Destination, Msg, SendingEvent, Service, data::QueueItem},
18};
19
20type Txns = BTreeMap<Destination, Vec<QueueItem>>;
21
22const BOOT_ARM_PACE: Duration = Duration::from_secs(2);
23
24#[implement(Service)]
25#[tracing::instrument(
26	name = "netburst",
27	level = "debug",
28	skip_all,
29	fields(
30		futures = %futures.len(),
31	),
32)]
33pub(super) async fn startup_netburst<'a>(
34	&'a self,
35	id: usize,
36	futures: &mut SendingFutures<'a>,
37	statuses: &mut TransactionStatuses,
38	wakes: &mut WakeQueue,
39) {
40	let netburst = self.server.config.startup_netburst;
41	let keep = usize::try_from(self.server.config.startup_netburst_keep).ok();
42	let txns = self
43		.db
44		.active_requests()
45		.ready_filter(|(_, _, dest)| self.shard_id(dest) == id)
46		.ready_fold(Txns::new(), |txns, (key, event, dest)| {
47			let len = txns.get(&dest).map_or(0, Vec::len);
48
49			match keep {
50				| Some(limit) if len >= limit => {
51					warn!(?dest, key = %String::from_utf8_lossy(&key), "Dropping unsent event");
52					self.db.delete_active_request(&key);
53					txns
54				},
55				| _ => batched(txns, dest, (key, event)),
56			}
57		})
58		.await;
59
60	txns.into_iter()
61		.filter(|(_, items)| !items.is_empty())
62		.for_each(|(dest, items)| {
63			let status = match netburst {
64				| true => TransactionStatus::Running { tries: 0 },
65				| false => TransactionStatus::Pending,
66			};
67
68			statuses.insert(dest.clone(), status);
69			if !netburst {
70				self.mark_pending(&dest);
71			}
72
73			if netburst {
74				futures.push(self.send_events(dest, items, None));
75			}
76		});
77
78	if !self.server.config.maintenance {
79		self.arm_inherited_destinations(id, statuses, wakes)
80			.await;
81	}
82
83	// Active transaction generations must own their queued successors before
84	// queued-only badge destinations are woken.
85	if !netburst || keep == Some(0) {
86		return;
87	}
88
89	let destinations: Vec<_> = self
90		.db
91		.queued_badge_refresh_destinations()
92		.ready_filter(|dest| self.shard_id(dest) == id)
93		.collect()
94		.await;
95
96	for dest in destinations.into_iter().sorted_unstable().dedup() {
97		let msg = Msg {
98			dest,
99			event: SendingEvent::BadgeRefresh,
100			queue_id: Vec::new(),
101		};
102
103		self.handle_request(msg, futures, statuses, wakes)
104			.await;
105	}
106}
107
108#[implement(Service)]
109fn mark_pending(&self, dest: &Destination) {
110	if let Destination::Federation(server) = dest {
111		self.stalled
112			.lock()
113			.expect("locked")
114			.insert(server.clone(), None);
115	}
116}
117
118#[implement(Service)]
119#[tracing::instrument(level = "trace", skip_all)]
120async fn arm_inherited_destinations(
121	&self,
122	id: usize,
123	statuses: &TransactionStatuses,
124	wakes: &mut WakeQueue,
125) {
126	let pending = statuses
127		.iter()
128		.filter(|(_, status)| matches!(status, TransactionStatus::Pending))
129		.filter(|(dest, _)| matches!(dest, Destination::Federation(_)))
130		.map(|(dest, _)| dest.clone())
131		.stream();
132
133	let queued = self
134		.db
135		.queued_federation_destinations(|server| self.federation_shard_id(server) == id)
136		.inspect(|result| {
137			if let Err(error) = result {
138				warn!(%error, "Queued federation discovery failed");
139			}
140		})
141		.ignore_err()
142		.ready_filter(|dest| !statuses.contains_key(dest));
143
144	pending
145		.chain(queued)
146		.fold((0_u64, wakes), async move |(index, wakes), dest| {
147			let index = self.arm_startup_wake(dest, index, wakes).await;
148
149			(index, wakes)
150		})
151		.await;
152}
153
154#[implement(Service)]
155#[tracing::instrument(level = "trace", skip_all)]
156pub(super) async fn arm_startup_wake(
157	&self,
158	dest: Destination,
159	index: u64,
160	wakes: &mut WakeQueue,
161) -> u64 {
162	if is_armed(wakes, &dest) {
163		return index;
164	}
165
166	let Destination::Federation(server) = &dest else {
167		return index;
168	};
169
170	match self
171		.services
172		.federation
173		.should_attempt(server)
174		.await
175	{
176		| ShouldAttempt::No { earliest_retry } => {
177			arm_wake(wakes, dest, earliest_retry);
178			index
179		},
180		| ShouldAttempt::Yes | ShouldAttempt::Deprioritize => {
181			let delay = boot_delay(self.server.config.sender_timeout, index);
182
183			arm_wake_in(wakes, dest, delay);
184			index.saturating_add(1)
185		},
186	}
187}
188
189fn batched(mut txns: Txns, dest: Destination, item: QueueItem) -> Txns {
190	txns.entry(dest).or_default().push(item);
191	txns
192}
193
194fn boot_delay(timeout: u64, index: u64) -> Duration {
195	let offset = match timeout {
196		| 0 => Duration::ZERO,
197		| _ => rand_secs(0..timeout),
198	};
199
200	let paced = Duration::from_secs(index.saturating_mul(BOOT_ARM_PACE.as_secs()));
201
202	offset.saturating_add(paced)
203}
tuwunel_service/sending/sender/netburst.rs

tuwunel_service/sending/sender/
netburst.rs
