Add max_retransmissions flag on each message (#5642)

This commit is contained in:
Jon Häggblad
2025-03-20 19:54:07 +01:00
committed by GitHub
parent 0abc07c96f
commit 1a37e60483
11 changed files with 218 additions and 74 deletions
+2
View File
@@ -11,6 +11,8 @@ on:
jobs:
build:
runs-on: arc-ubuntu-20.04
env:
RUSTUP_PERMIT_COPY_RENAME: 1
defaults:
run:
working-directory: documentation/docs
@@ -28,6 +28,7 @@ pub enum InputMessage {
recipient: Recipient,
data: Vec<u8>,
lane: TransmissionLane,
max_retransmissions: Option<u32>,
},
/// Creates a message used for a duplex anonymous communication where the recipient
@@ -43,6 +44,7 @@ pub enum InputMessage {
data: Vec<u8>,
reply_surbs: u32,
lane: TransmissionLane,
max_retransmissions: Option<u32>,
},
/// Attempt to use our internally received and stored `ReplySurb` to send the message back
@@ -53,6 +55,7 @@ pub enum InputMessage {
recipient_tag: AnonymousSenderTag,
data: Vec<u8>,
lane: TransmissionLane,
max_retransmissions: Option<u32>,
},
MessageWrapper {
@@ -92,6 +95,7 @@ impl InputMessage {
recipient,
data,
lane,
max_retransmissions: None,
};
if let Some(packet_type) = packet_type {
InputMessage::new_wrapper(message, packet_type)
@@ -112,28 +116,7 @@ impl InputMessage {
data,
reply_surbs,
lane,
};
if let Some(packet_type) = packet_type {
InputMessage::new_wrapper(message, packet_type)
} else {
message
}
}
// IMHO `new_anonymous` should take `mix_hops: Option<u8>` as an argument instead of creating
// this function, but that would potentially break backwards compatibility with the current API
pub fn new_anonymous_with_custom_hops(
recipient: Recipient,
data: Vec<u8>,
reply_surbs: u32,
lane: TransmissionLane,
packet_type: Option<PacketType>,
) -> Self {
let message = InputMessage::Anonymous {
recipient,
data,
reply_surbs,
lane,
max_retransmissions: None,
};
if let Some(packet_type) = packet_type {
InputMessage::new_wrapper(message, packet_type)
@@ -152,6 +135,7 @@ impl InputMessage {
recipient_tag,
data,
lane,
max_retransmissions: None,
};
if let Some(packet_type) = packet_type {
InputMessage::new_wrapper(message, packet_type)
@@ -169,4 +153,34 @@ impl InputMessage {
InputMessage::MessageWrapper { message, .. } => message.lane(),
}
}
pub fn set_max_retransmissions(&mut self, max_retransmissions: u32) -> &mut Self {
match self {
InputMessage::Regular {
max_retransmissions: m,
..
}
| InputMessage::Anonymous {
max_retransmissions: m,
..
}
| InputMessage::Reply {
max_retransmissions: m,
..
} => {
*m = Some(max_retransmissions);
}
InputMessage::Premade { .. } => {}
InputMessage::MessageWrapper { message, .. } => {
message.set_max_retransmissions(max_retransmissions);
}
}
self
}
pub fn with_max_retransmissions(mut self, max_retransmissions: u32) -> Self {
self.set_max_retransmissions(max_retransmissions);
self
}
}
@@ -222,9 +222,6 @@ impl ActionController {
// note: when the entry expires it's automatically removed from pending_acks_timers
fn handle_expired_ack_timer(&mut self, expired_ack: Expired<FragmentIdentifier>) {
// I'm honestly not sure how to handle it, because getting it means other things in our
// system are already misbehaving. If we ever see this panic, then I guess we should worry
// about it. Perhaps just reschedule it at later point?
let frag_id = expired_ack.into_inner();
trace!("{frag_id} has expired");
@@ -65,11 +65,12 @@ where
recipient_tag: AnonymousSenderTag,
data: Vec<u8>,
lane: TransmissionLane,
max_retransmissions: Option<u32>,
) {
// offload reply handling to the dedicated task
if let Err(err) = self
.reply_controller_sender
.send_reply(recipient_tag, data, lane)
if let Err(err) =
self.reply_controller_sender
.send_reply(recipient_tag, data, lane, max_retransmissions)
{
if !self.task_client.is_shutdown_poll() {
error!("failed to send a reply - {err}");
@@ -83,10 +84,11 @@ where
content: Vec<u8>,
lane: TransmissionLane,
packet_type: PacketType,
max_retransmissions: Option<u32>,
) {
if let Err(err) = self
.message_handler
.try_send_plain_message(recipient, content, lane, packet_type)
.try_send_plain_message(recipient, content, lane, packet_type, max_retransmissions)
.await
{
warn!("failed to send a plain message - {err}")
@@ -100,10 +102,18 @@ where
reply_surbs: u32,
lane: TransmissionLane,
packet_type: PacketType,
max_retransmissions: Option<u32>,
) {
if let Err(err) = self
.message_handler
.try_send_message_with_reply_surbs(recipient, content, reply_surbs, lane, packet_type)
.try_send_message_with_reply_surbs(
recipient,
content,
reply_surbs,
lane,
packet_type,
max_retransmissions,
)
.await
{
warn!("failed to send a repliable message - {err}")
@@ -116,25 +126,42 @@ where
recipient,
data,
lane,
max_retransmissions,
} => {
self.handle_plain_message(recipient, data, lane, PacketType::Mix)
.await
self.handle_plain_message(
recipient,
data,
lane,
PacketType::Mix,
max_retransmissions,
)
.await
}
InputMessage::Anonymous {
recipient,
data,
reply_surbs,
lane,
max_retransmissions,
} => {
self.handle_repliable_message(recipient, data, reply_surbs, lane, PacketType::Mix)
.await
self.handle_repliable_message(
recipient,
data,
reply_surbs,
lane,
PacketType::Mix,
max_retransmissions,
)
.await
}
InputMessage::Reply {
recipient_tag,
data,
lane,
max_retransmissions,
} => {
self.handle_reply(recipient_tag, data, lane).await;
self.handle_reply(recipient_tag, data, lane, max_retransmissions)
.await;
}
InputMessage::Premade { msgs, lane } => self.handle_premade_packets(msgs, lane).await,
InputMessage::MessageWrapper {
@@ -145,25 +172,42 @@ where
recipient,
data,
lane,
max_retransmissions,
} => {
self.handle_plain_message(recipient, data, lane, packet_type)
.await
self.handle_plain_message(
recipient,
data,
lane,
packet_type,
max_retransmissions,
)
.await
}
InputMessage::Anonymous {
recipient,
data,
reply_surbs,
lane,
max_retransmissions,
} => {
self.handle_repliable_message(recipient, data, reply_surbs, lane, packet_type)
.await
self.handle_repliable_message(
recipient,
data,
reply_surbs,
lane,
packet_type,
max_retransmissions,
)
.await
}
InputMessage::Reply {
recipient_tag,
data,
lane,
max_retransmissions,
} => {
self.handle_reply(recipient_tag, data, lane).await;
self.handle_reply(recipient_tag, data, lane, max_retransmissions)
.await;
}
InputMessage::Premade { msgs, lane } => {
self.handle_premade_packets(msgs, lane).await
@@ -72,6 +72,7 @@ pub struct PendingAcknowledgement {
delay: SphinxDelay,
destination: PacketDestination,
retransmissions: u32,
max_retransmissions: Option<u32>,
}
impl PendingAcknowledgement {
@@ -80,12 +81,14 @@ impl PendingAcknowledgement {
message_chunk: Fragment,
delay: SphinxDelay,
recipient: Recipient,
max_retransmissions: Option<u32>,
) -> Self {
PendingAcknowledgement {
message_chunk,
delay,
destination: PacketDestination::KnownRecipient(recipient.into()),
retransmissions: 0,
max_retransmissions,
}
}
@@ -94,6 +97,7 @@ impl PendingAcknowledgement {
delay: SphinxDelay,
recipient_tag: AnonymousSenderTag,
extra_surb_request: bool,
max_retransmissions: Option<u32>,
) -> Self {
PendingAcknowledgement {
message_chunk,
@@ -103,6 +107,7 @@ impl PendingAcknowledgement {
extra_surb_request,
},
retransmissions: 0,
max_retransmissions,
}
}
@@ -118,6 +123,18 @@ impl PendingAcknowledgement {
self.delay = new_delay;
self.retransmissions += 1;
}
pub(crate) fn reached_max_retransmissions(
&self,
global_max_retransmissions: Option<u32>,
) -> bool {
let reached_local_max = self
.max_retransmissions
.is_some_and(|limit| self.retransmissions >= limit);
let reached_global_max =
global_max_retransmissions.is_some_and(|limit| self.retransmissions >= limit);
reached_local_max || reached_global_max
}
}
/// AcknowledgementControllerConnectors represents set of channels for communication with
@@ -79,17 +79,15 @@ where
let frag_id = timed_out_ack.message_chunk.fragment_identifier();
if let Some(limit) = self.maximum_retransmissions {
if timed_out_ack.retransmissions >= limit {
warn!("reached maximum number of allowed retransmissions for the packet");
if let Err(err) = self
.action_sender
.unbounded_send(Action::new_remove(frag_id))
{
error!("Failed to send remove action to the controller: {err}");
}
return;
if timed_out_ack.reached_max_retransmissions(self.maximum_retransmissions) {
debug!("reached maximum number of allowed retransmissions for the packet");
if let Err(err) = self
.action_sender
.unbounded_send(Action::new_remove(frag_id))
{
error!("Failed to send remove action to the controller: {err}");
}
return;
}
let maybe_prepared_fragment = match &timed_out_ack.destination {
@@ -6,6 +6,7 @@ use crate::client::real_messages_control::real_traffic_stream::{
BatchRealMessageSender, RealMessage,
};
use crate::client::real_messages_control::{AckActionSender, Action};
use crate::client::replies::reply_controller::MaxRetransmissions;
use crate::client::replies::reply_storage::{ReceivedReplySurbsMap, SentReplyKeys, UsedSenderTags};
use crate::client::topology_control::{TopologyAccessor, TopologyReadPermit};
use log::{debug, error, info, trace, warn};
@@ -142,6 +143,12 @@ impl Config {
}
}
#[derive(Clone)]
pub(crate) struct FragmentWithMaxRetransmissions {
pub(crate) fragment: Fragment,
pub(crate) max_retransmissions: MaxRetransmissions,
}
#[derive(Clone)]
pub(crate) struct MessageHandler<R> {
config: Config,
@@ -294,8 +301,14 @@ where
Some(chunk.fragment_identifier()),
);
let delay = prepared_fragment.total_delay;
let pending_ack =
PendingAcknowledgement::new_anonymous(chunk, delay, target, is_extra_surb_request);
let max_retransmissions = None;
let pending_ack = PendingAcknowledgement::new_anonymous(
chunk,
delay,
target,
is_extra_surb_request,
max_retransmissions,
);
let lane = if is_extra_surb_request {
TransmissionLane::ReplySurbRequest
@@ -350,7 +363,7 @@ where
pub(crate) async fn try_send_reply_chunks_on_lane(
&mut self,
target: AnonymousSenderTag,
fragments: Vec<Fragment>,
fragments: Vec<FragmentWithMaxRetransmissions>,
reply_surbs: Vec<ReplySurb>,
lane: TransmissionLane,
) -> Result<(), SurbWrappedPreparationError> {
@@ -367,12 +380,12 @@ where
pub(crate) async fn try_send_reply_chunks(
&mut self,
target: AnonymousSenderTag,
fragments: Vec<(TransmissionLane, Fragment)>,
fragments: Vec<(TransmissionLane, FragmentWithMaxRetransmissions)>,
reply_surbs: Vec<ReplySurb>,
) -> Result<(), SurbWrappedPreparationError> {
let prepared_fragments = self
.prepare_reply_chunks_for_sending(
fragments.iter().map(|(_, f)| f.clone()).collect(),
fragments.iter().map(|(_, f)| f.fragment.clone()).collect(),
reply_surbs,
)
.await?;
@@ -382,12 +395,21 @@ where
for (raw, prepared) in fragments.into_iter().zip(prepared_fragments.into_iter()) {
let lane = raw.0;
let fragment = raw.1;
let FragmentWithMaxRetransmissions {
fragment,
max_retransmissions,
} = raw.1;
let real_message =
RealMessage::new(prepared.mix_packet, Some(prepared.fragment_identifier));
let delay = prepared.total_delay;
let pending_ack = PendingAcknowledgement::new_anonymous(fragment, delay, target, false);
let pending_ack = PendingAcknowledgement::new_anonymous(
fragment,
delay,
target,
false,
max_retransmissions,
);
let entry = to_forward.entry(lane).or_default();
entry.push(real_message);
@@ -416,10 +438,17 @@ where
message: Vec<u8>,
lane: TransmissionLane,
packet_type: PacketType,
max_retransmissions: Option<u32>,
) -> Result<(), PreparationError> {
let message = NymMessage::new_plain(message);
self.try_split_and_send_non_reply_message(message, recipient, lane, packet_type)
.await
self.try_split_and_send_non_reply_message(
message,
recipient,
lane,
packet_type,
max_retransmissions,
)
.await
}
pub(crate) async fn try_split_and_send_non_reply_message(
@@ -428,6 +457,7 @@ where
recipient: Recipient,
lane: TransmissionLane,
packet_type: PacketType,
max_retransmissions: Option<u32>,
) -> Result<(), PreparationError> {
debug!("Sending non-reply message with packet type {packet_type}");
// TODO: I really dislike existence of this assertion, it implies code has to be re-organised
@@ -467,7 +497,8 @@ where
Some(fragment.fragment_identifier()),
);
let delay = prepared_fragment.total_delay;
let pending_ack = PendingAcknowledgement::new_known(fragment, delay, recipient);
let pending_ack =
PendingAcknowledgement::new_known(fragment, delay, recipient, max_retransmissions);
real_messages.push(real_message);
pending_acks.push(pending_ack);
@@ -495,11 +526,15 @@ where
reply_surbs,
));
// When sending SURBs we want to retransmit
let max_retransmissions = None;
self.try_split_and_send_non_reply_message(
message,
recipient,
TransmissionLane::AdditionalReplySurbs,
packet_type,
max_retransmissions,
)
.await?;
@@ -516,6 +551,7 @@ where
num_reply_surbs: u32,
lane: TransmissionLane,
packet_type: PacketType,
max_retransmissions: Option<u32>,
) -> Result<(), SurbWrappedPreparationError> {
debug!("Sending message with reply SURBs with packet type {packet_type}");
let sender_tag = self.get_or_create_sender_tag(&recipient);
@@ -526,8 +562,14 @@ where
let message =
NymMessage::new_repliable(RepliableMessage::new_data(message, sender_tag, reply_surbs));
self.try_split_and_send_non_reply_message(message, recipient, lane, packet_type)
.await?;
self.try_split_and_send_non_reply_message(
message,
recipient,
lane,
packet_type,
max_retransmissions,
)
.await?;
log::trace!("storing {} reply keys", reply_keys.len());
self.reply_key_storage.insert_multiple(reply_keys);
@@ -2,7 +2,9 @@
// SPDX-License-Identifier: Apache-2.0
use crate::client::real_messages_control::acknowledgement_control::PendingAcknowledgement;
use crate::client::real_messages_control::message_handler::{MessageHandler, PreparationError};
use crate::client::real_messages_control::message_handler::{
FragmentWithMaxRetransmissions, MessageHandler, PreparationError,
};
use crate::client::replies::reply_storage::CombinedReplyStorage;
use futures::channel::oneshot;
use futures::StreamExt;
@@ -10,7 +12,7 @@ use log::{debug, error, info, trace, warn};
use nym_sphinx::addressing::clients::Recipient;
use nym_sphinx::anonymous_replies::requests::AnonymousSenderTag;
use nym_sphinx::anonymous_replies::ReplySurb;
use nym_sphinx::chunking::fragment::{Fragment, FragmentIdentifier};
use nym_sphinx::chunking::fragment::FragmentIdentifier;
use nym_task::connections::{ConnectionId, TransmissionLane};
use nym_task::TaskClient;
use rand::{CryptoRng, Rng};
@@ -49,6 +51,8 @@ impl Config {
// - replies to "give additional surbs" requests
// - will reply to future heartbeats
pub type MaxRetransmissions = Option<u32>;
// TODO: this should be split into ingress and egress controllers
// because currently its trying to perform two distinct jobs
pub struct ReplyController<R> {
@@ -59,7 +63,8 @@ pub struct ReplyController<R> {
// of surbs required to send the message through
// expected_reliability: f32,
request_receiver: ReplyControllerReceiver,
pending_replies: HashMap<AnonymousSenderTag, TransmissionBuffer<Fragment>>,
pending_replies:
HashMap<AnonymousSenderTag, TransmissionBuffer<FragmentWithMaxRetransmissions>>,
/// Retransmission packets that have already timed out and are waiting for additional reply SURBs
/// so that they could be sent back to the network. Once we receive more SURBs, we should send them ASAP.
@@ -96,7 +101,7 @@ where
}
}
fn insert_pending_replies<I: IntoIterator<Item = Fragment>>(
fn insert_pending_replies<I: IntoIterator<Item = FragmentWithMaxRetransmissions>>(
&mut self,
recipient: &AnonymousSenderTag,
fragments: I,
@@ -112,7 +117,7 @@ where
fn re_insert_pending_replies(
&mut self,
recipient: &AnonymousSenderTag,
fragments: Vec<(TransmissionLane, Fragment)>,
fragments: Vec<(TransmissionLane, FragmentWithMaxRetransmissions)>,
) {
trace!("re-inserting pending replies for {recipient}");
// the buffer should ALWAYS exist at this point, if it doesn't, it's a bug...
@@ -205,6 +210,7 @@ where
recipient_tag: AnonymousSenderTag,
data: Vec<u8>,
lane: TransmissionLane,
max_retransmissions: Option<u32>,
) {
if !self
.full_reply_storage
@@ -242,7 +248,14 @@ where
.get_reply_surbs(&recipient_tag, max_to_send);
if let Some(reply_surbs) = surbs {
let to_send = fragments.drain(..max_to_send).collect::<Vec<_>>();
let to_send = fragments
.drain(..max_to_send)
.map(|f| FragmentWithMaxRetransmissions {
fragment: f,
max_retransmissions,
})
.collect::<Vec<_>>();
if let Err(err) = self
.message_handler
.try_send_reply_chunks_on_lane(
@@ -276,6 +289,13 @@ where
"buffering {no_fragments} fragments for {recipient_tag}",
no_fragments = fragments.len()
);
let fragments: Vec<_> = fragments
.into_iter()
.map(|fragment| FragmentWithMaxRetransmissions {
fragment,
max_retransmissions,
})
.collect();
self.insert_pending_replies(&recipient_tag, fragments, lane);
}
@@ -409,7 +429,7 @@ where
&mut self,
from: &AnonymousSenderTag,
amount: usize,
) -> Option<Vec<(TransmissionLane, Fragment)>> {
) -> Option<Vec<(TransmissionLane, FragmentWithMaxRetransmissions)>> {
// if possible, pop all pending replies, if not, pop only entries for which we'd have a reply surb
let total = self.pending_replies.get(from)?.total_size();
trace!("pending queue has {total} elements");
@@ -689,7 +709,11 @@ where
recipient,
message,
lane,
} => self.handle_send_reply(recipient, message, lane).await,
max_retransmissions,
} => {
self.handle_send_reply(recipient, message, lane, max_retransmissions)
.await
}
ReplyControllerMessage::AdditionalSurbs {
sender_tag,
reply_surbs,
@@ -66,12 +66,14 @@ impl ReplyControllerSender {
recipient: AnonymousSenderTag,
message: Vec<u8>,
lane: TransmissionLane,
max_retransmissions: Option<u32>,
) -> Result<(), ReplyControllerSenderError> {
self.0
.unbounded_send(ReplyControllerMessage::SendReply {
recipient,
message,
lane,
max_retransmissions,
})
.map_err(ReplyControllerSenderError::SendReply)
}
@@ -160,6 +162,7 @@ pub enum ReplyControllerMessage {
recipient: AnonymousSenderTag,
message: Vec<u8>,
lane: TransmissionLane,
max_retransmissions: Option<u32>,
},
AdditionalSurbs {
@@ -182,10 +182,11 @@ impl MixnetMessageSinkTranslator for ToIprDataResponse {
let response_packet = create_ip_packet_response(bundled_ip_packets, self.client_version)?;
// Wrap the response packet in a mixnet input message
Ok(crate::util::create_message::create_input_message(
&self.send_to,
response_packet,
))
let input_message =
crate::util::create_message::create_input_message(&self.send_to, response_packet)
.with_max_retransmissions(0);
Ok(input_message)
}
}
@@ -191,6 +191,7 @@ impl MixnetAddress {
recipient: *recipient,
data: message,
lane: TransmissionLane::ConnectionId(connection_id),
max_retransmissions: None,
}),
packet_type,
},
@@ -199,6 +200,7 @@ impl MixnetAddress {
recipient_tag: sender_tag,
data: message,
lane: TransmissionLane::ConnectionId(connection_id),
max_retransmissions: None,
}),
packet_type,
},