Feature/ack messing (#313)

* Added Display for FragmentIdentifier * AckDelayQueue as wrapper for DelayQueue * First pass on ack controller * Controller very initially incorporated into the client * Cleanup * Using socks5 build.rs * Extra (temporary) log statements * Decreased logging levels
2020-08-25 10:45:31 +01:00
parent 7b2036bd33
commit a0484ce082
13 changed files with 769 additions and 388 deletions
@@ -0,0 +1,68 @@
+// Copyright 2020 Nym Technologies SA
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use futures::Stream;
+use std::pin::Pin;
+use std::task::{Context, Poll, Waker};
+use std::time::Duration;
+use tokio::time::{
+    delay_queue::{self, Expired},
+    DelayQueue,
+};
+
+// works under assumption that it will be used inside a loop, where we never want a `None`
+// TODO: perhaps this should/could be renamed and moved to common/utils (and expose all inner methods?)
+pub struct AckDelayQueue<T> {
+    inner: DelayQueue<T>,
+    waker: Option<Waker>,
+}
+
+// more methods of underlying DelayQueue will get exposed as we need them
+impl<T> AckDelayQueue<T> {
+    pub fn new() -> Self {
+        AckDelayQueue {
+            inner: DelayQueue::new(),
+            waker: None,
+        }
+    }
+
+    pub fn insert(&mut self, value: T, timeout: Duration) -> delay_queue::Key {
+        let key = self.inner.insert(value, timeout);
+        if let Some(waker) = self.waker.take() {
+            // we were waiting for an item - wake the executor!
+            waker.wake()
+        }
+        key
+    }
+
+    pub fn remove(&mut self, key: &delay_queue::Key) -> Expired<T> {
+        self.inner.remove(key)
+    }
+}
+
+impl<T> Stream for AckDelayQueue<T> {
+    type Item = <DelayQueue<T> as Stream>::Item;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match Pin::new(&mut self.inner).poll_next(cx) {
+            Poll::Pending => Poll::Pending,
+            Poll::Ready(Some(item)) => Poll::Ready(Some(item)),
+            Poll::Ready(None) => {
+                // we'll need to keep the waker to notify the executor once we get new item
+                self.waker = Some(cx.waker().clone());
+                Poll::Pending
+            }
+        }
+    }
+}
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

-use super::PendingAcksMap;
+use super::action_controller::{Action, ActionSender};
 use futures::StreamExt;
 use gateway_client::AcknowledgementReceiver;
 use log::*;
@@ -22,67 +22,62 @@ use nymsphinx::{
 };
 use std::sync::Arc;

-// responsible for cancelling retransmission timers and removed entries from the map
+/// Module responsible for listening for any data resembling acknowledgements from the network
+/// and firing actions to remove them from the 'Pending' state.
 pub(super) struct AcknowledgementListener {
    ack_key: Arc<AckKey>,
    ack_receiver: AcknowledgementReceiver,
-    pending_acks: PendingAcksMap,
+    action_sender: ActionSender,
 }

 impl AcknowledgementListener {
    pub(super) fn new(
        ack_key: Arc<AckKey>,
        ack_receiver: AcknowledgementReceiver,
-        pending_acks: PendingAcksMap,
+        action_sender: ActionSender,
    ) -> Self {
        AcknowledgementListener {
            ack_key,
            ack_receiver,
-            pending_acks,
+            action_sender,
        }
    }

    async fn on_ack(&mut self, ack_content: Vec<u8>) {
        debug!("Received an ack");
-        let frag_id = match recover_identifier(&self.ack_key, &ack_content) {
-            None => {
+        let frag_id = match recover_identifier(&self.ack_key, &ack_content)
+            .map(FragmentIdentifier::try_from_bytes)
+        {
+            Some(Ok(frag_id)) => frag_id,
+            _ => {
                warn!("Received invalid ACK!"); // should we do anything else about that?
                return;
            }
-            Some(frag_id_bytes) => match FragmentIdentifier::try_from_bytes(frag_id_bytes) {
-                Ok(frag_id) => frag_id,
-                Err(err) => {
-                    warn!("Received invalid ACK! - {:?}", err); // should we do anything else about that?
-                    return;
-                }
-            },
        };

+        // if we received an ack for cover message or a reply there will be nothing to remove,
+        // because nothing was inserted in the first place
        if frag_id == COVER_FRAG_ID {
            trace!("Received an ack for a cover message - no need to do anything");
            return;
        } else if frag_id.is_reply() {
-            debug!("Received an ack for a reply message - no need to do anything!");
+            info!("Received an ack for a reply message - no need to do anything! (don't know what to do!)");
            // TODO: probably there will need to be some extra procedure here, something to notify
            // user that his reply reached the recipient (since we got an ack)
-            info!("We received an ack for one of the replies we sent!");
            return;
        }

-        if let Some(pending_ack) = self.pending_acks.write().await.remove(&frag_id) {
-            // cancel the retransmission future
-            pending_ack.retransmission_cancel.notify();
-        } else {
-            warn!("received ACK for packet we haven't stored! - {:?}", frag_id);
-        }
+        trace!("Received {} from the mix network", frag_id);
+
+        self.action_sender
+            .unbounded_send(Action::new_remove(frag_id))
+            .unwrap();
    }

    pub(super) async fn run(&mut self) {
        debug!("Started AcknowledgementListener");
        while let Some(acks) = self.ack_receiver.next().await {
-            // realistically we would only be getting one ack at the time, but if we managed to
-            // introduce batching in gateway client, this call should be improved to not re-acquire
-            // write permit on the map every loop iteration
+            // realistically we would only be getting one ack at the time
            for ack in acks {
                self.on_ack(ack).await;
            }
@@ -0,0 +1,277 @@
+// Copyright 2020 Nym Technologies SA
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::PendingAcknowledgement;
+use crate::client::real_messages_control::acknowledgement_control::ack_delay_queue::AckDelayQueue;
+use crate::client::real_messages_control::acknowledgement_control::RetransmissionRequestSender;
+use futures::channel::mpsc::{self, UnboundedReceiver, UnboundedSender};
+use log::*;
+use nymsphinx::chunking::fragment::FragmentIdentifier;
+use nymsphinx::Delay as SphinxDelay;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::stream::StreamExt;
+use tokio::time::delay_queue::{self, Expired};
+use tokio::time::Error as TimerError;
+
+pub(crate) type ActionSender = UnboundedSender<Action>;
+
+// The actual data being sent off as well as potential key to the delay queue
+type PendingAckEntry = (Arc<PendingAcknowledgement>, Option<delay_queue::Key>);
+
+// we can either:
+// - have a completely new set of packets we just sent and need to create entries for
+// - received an ack so we want to remove an entry
+// - start a retransmission timer for sending the packet into the network (on either first try or retransmission)
+// - update the internal sphinx delay of an expired packet
+pub(crate) enum Action {
+    /// Inserts new `PendingAcknowledgement`s into the 'shared' state.
+    /// Initiated by `InputMessageListener`
+    InsertPending(Vec<PendingAcknowledgement>),
+
+    /// Removes given `PendingAcknowledgement` from the 'shared' state. Also cancels the retransmission timer.
+    /// Initiated by `AcknowledgementListener`
+    RemovePending(FragmentIdentifier),
+
+    /// Starts the retransmission timer on given `PendingAcknowledgement` with the `Duration` based on
+    /// its internal data.
+    /// Initiated by `SentNotificationListener`
+    /// Can also be initiated by `RetransmissionRequestListener` in the rare cases of invalid Topology.
+    StartTimer(FragmentIdentifier),
+
+    /// Updates the expected delay of given `PendingAcknowledgement` with the new provided `SphinxDelay`.
+    /// Initiated by `RetransmissionRequestListener`
+    UpdateDelay(FragmentIdentifier, SphinxDelay),
+}
+
+impl Action {
+    pub(crate) fn new_insert(pending_acks: Vec<PendingAcknowledgement>) -> Self {
+        Action::InsertPending(pending_acks)
+    }
+
+    pub(crate) fn new_remove(frag_id: FragmentIdentifier) -> Self {
+        Action::RemovePending(frag_id)
+    }
+
+    pub(crate) fn new_start_timer(frag_id: FragmentIdentifier) -> Self {
+        Action::StartTimer(frag_id)
+    }
+
+    pub(crate) fn new_update_delay(frag_id: FragmentIdentifier, delay: SphinxDelay) -> Self {
+        Action::UpdateDelay(frag_id, delay)
+    }
+}
+
+/// Configurable parameters of the `ActionController`
+pub(super) struct Config {
+    /// Given ack timeout in the form a * BASE_DELAY + b, it specifies the additive part `b`
+    ack_wait_addition: Duration,
+
+    /// Given ack timeout in the form a * BASE_DELAY + b, it specifies the multiplier `a`
+    ack_wait_multiplier: f64,
+}
+
+impl Config {
+    pub(super) fn new(ack_wait_addition: Duration, ack_wait_multiplier: f64) -> Self {
+        Config {
+            ack_wait_addition,
+            ack_wait_multiplier,
+        }
+    }
+}
+
+pub(super) struct ActionController {
+    /// Configurable parameters of the `ActionController`
+    config: Config,
+
+    /// Contains a map between `FragmentIdentifier` and its full `PendingAcknowledgement` as well as
+    /// key to its `AckDelayQueue` entry if it was started.
+    pending_acks_data: HashMap<FragmentIdentifier, PendingAckEntry>,
+
+    // This structure ensures that we will EITHER handle expired timer or a received action and NEVER both
+    // at the same time hence getting rid of one possible race condition that we suffered from in the
+    // previous version.
+    /// DelayQueue with all `PendingAcknowledgement` that are waiting to be either received or
+    /// retransmitted if their timer fires up.
+    pending_acks_timers: AckDelayQueue<FragmentIdentifier>,
+
+    /// Channel for receiving `Action`s from other modules.
+    incoming_actions: UnboundedReceiver<Action>,
+
+    /// Channel for notifying `RetransmissionRequestListener` about expired acknowledgements.
+    retransmission_sender: RetransmissionRequestSender,
+}
+
+impl ActionController {
+    pub(super) fn new(
+        config: Config,
+        retransmission_sender: RetransmissionRequestSender,
+    ) -> (Self, ActionSender) {
+        let (sender, receiver) = mpsc::unbounded();
+        (
+            ActionController {
+                config,
+                pending_acks_data: HashMap::new(),
+                pending_acks_timers: AckDelayQueue::new(),
+                incoming_actions: receiver,
+                retransmission_sender,
+            },
+            sender,
+        )
+    }
+
+    fn handle_insert(&mut self, pending_acks: Vec<PendingAcknowledgement>) {
+        for pending_ack in pending_acks {
+            let frag_id = pending_ack.message_chunk.fragment_identifier();
+            trace!("{} is inserted", frag_id);
+
+            if self
+                .pending_acks_data
+                .insert(frag_id, (Arc::new(pending_ack), None))
+                .is_some()
+            {
+                panic!("Tried to insert duplicate pending ack")
+            }
+        }
+    }
+
+    fn handle_start_timer(&mut self, frag_id: FragmentIdentifier) {
+        trace!("{} is starting its timer", frag_id);
+
+        if let Some((pending_ack_data, queue_key)) = self.pending_acks_data.get_mut(&frag_id) {
+            if queue_key.is_some() {
+                // this branch should be IMPOSSIBLE under ANY condition. It would imply starting
+                // timer TWICE for the SAME PendingAcknowledgement
+                panic!("Tried to start an already started ack timer!")
+            }
+            let timeout = (pending_ack_data.delay.clone() * self.config.ack_wait_multiplier)
+                .to_duration()
+                + self.config.ack_wait_addition;
+
+            let new_queue_key = self.pending_acks_timers.insert(frag_id, timeout);
+            *queue_key = Some(new_queue_key)
+        } else {
+            debug!(
+                "Tried to START TIMER on pending ack that is already gone! - {}",
+                frag_id
+            );
+        }
+    }
+
+    fn handle_remove(&mut self, frag_id: FragmentIdentifier) {
+        trace!("{} is getting removed", frag_id);
+
+        match self.pending_acks_data.remove(&frag_id) {
+            None => {
+                debug!(
+                    "Tried to REMOVE pending ack that is already gone! - {}",
+                    frag_id
+                );
+            }
+            Some((_, queue_key)) => {
+                if let Some(queue_key) = queue_key {
+                    // there are no possible checks here, we must GUARANTEE that we NEVER try
+                    // to remove an entry that doesn't exist (and we MUST GUARANTEE that
+                    // we do not have a stale key)
+                    self.pending_acks_timers.remove(&queue_key);
+                // remove timer
+                } else {
+                    // I'm not 100% sure if having a `None` key is even possible here
+                    // (REMOVE would have to be called before START TIMER),
+                    debug!(
+                        "Tried to REMOVE pending ack without TIMER active - {}",
+                        frag_id
+                    );
+                }
+            }
+        }
+    }
+
+    // initiated basically as a first step of retransmission. At first data has its delay updated
+    // (as new sphinx packet was created with new expected delivery time)
+    fn handle_update_delay(&mut self, frag_id: FragmentIdentifier, delay: SphinxDelay) {
+        trace!("{} is updating its delay", frag_id);
+        // TODO: is it possible to solve this without either locking or temporarily removing the value?
+        if let Some((pending_ack_data, queue_key)) = self.pending_acks_data.remove(&frag_id) {
+            // this Action is triggered by `RetransmissionRequestListener` which held the other potential
+            // reference to this Arc. HOWEVER, before the Action was pushed onto the queue, the reference
+            // was dropped hence this unwrap is safe.
+            let mut inner_data = Arc::try_unwrap(pending_ack_data).unwrap();
+            inner_data.update_delay(delay);
+
+            self.pending_acks_data
+                .insert(frag_id, (Arc::new(inner_data), queue_key));
+        } else {
+            debug!(
+                "Tried to UPDATE TIMER on pending ack that is already gone! - {}",
+                frag_id
+            );
+        }
+    }
+
+    // note: when the entry expires it's automatically removed from pending_acks_timers
+    fn handle_expired_ack_timer(
+        &mut self,
+        expired_ack: Result<Expired<FragmentIdentifier>, TimerError>,
+    ) {
+        // 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
+            .expect("Tokio timer returned an error!")
+            .into_inner();
+
+        trace!("{} has expired", frag_id);
+
+        if let Some((pending_ack_data, queue_key)) = self.pending_acks_data.get_mut(&frag_id) {
+            if queue_key.is_none() {
+                // this branch should be IMPOSSIBLE under ANY condition. It would imply the timeout
+                // happened before it even started.
+                panic!("Ack expired before it was even scheduled!")
+            }
+            *queue_key = None;
+            // downgrading an arc and then upgrading vs cloning is difference of 30ns vs 15ns
+            // so it's literally a NO difference while it might prevent us from unnecessarily
+            // resending data (in maybe 1 in 1 million cases, but it's something)
+            self.retransmission_sender
+                .unbounded_send(Arc::downgrade(pending_ack_data))
+                .unwrap()
+        } else {
+            // this shouldn't cause any issues but shouldn't have happened to begin with!
+            error!("An already removed pending ack has expired")
+        }
+    }
+
+    fn process_action(&mut self, action: Action) {
+        match action {
+            Action::InsertPending(pending_acks) => self.handle_insert(pending_acks),
+            Action::RemovePending(frag_id) => self.handle_remove(frag_id),
+            Action::StartTimer(frag_id) => self.handle_start_timer(frag_id),
+            Action::UpdateDelay(frag_id, delay) => self.handle_update_delay(frag_id, delay),
+        }
+    }
+
+    pub(super) async fn run(&mut self) {
+        loop {
+            // at some point there will be a global shutdown signal here as the third option
+            tokio::select! {
+                // we NEVER expect for ANY sender to get dropped so unwrap here is fine
+                action = self.incoming_actions.next() => self.process_action(action.unwrap()),
+                // pending ack queue Stream CANNOT return a `None` so unwrap here is fine
+                expired_ack = self.pending_acks_timers.next() => self.handle_expired_ack_timer(expired_ack.unwrap())
+            }
+        }
+    }
+}
@@ -12,7 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

-use super::{PendingAcknowledgement, PendingAcksMap};
+use super::action_controller::{Action, ActionSender};
+use super::PendingAcknowledgement;
 use crate::client::reply_key_storage::ReplyKeyStorage;
 use crate::client::{
    inbound_messages::{InputMessage, InputMessageReceiver},
@@ -27,7 +28,9 @@ use nymsphinx::{acknowledgements::AckKey, addressing::clients::Recipient};
 use rand::{CryptoRng, Rng};
 use std::sync::Arc;

-// responsible for splitting received message and initial sending attempt
+/// Module responsible for dealing with the received messages: splitting them, creating acknowledgements,
+/// putting everything into sphinx packets, etc.
+/// It also makes an initial sending attempt for said messages.
 pub(super) struct InputMessageListener<R>
 where
    R: CryptoRng + Rng,
@@ -36,7 +39,7 @@ where
    ack_recipient: Recipient,
    input_receiver: InputMessageReceiver,
    message_preparer: MessagePreparer<R>,
-    pending_acks: PendingAcksMap,
+    action_sender: ActionSender,
    real_message_sender: RealMessageSender,
    topology_access: TopologyAccessor,
    reply_key_storage: ReplyKeyStorage,
@@ -51,7 +54,7 @@ where
        ack_recipient: Recipient,
        input_receiver: InputMessageReceiver,
        message_preparer: MessagePreparer<R>,
-        pending_acks: PendingAcksMap,
+        action_sender: ActionSender,
        real_message_sender: RealMessageSender,
        topology_access: TopologyAccessor,
        reply_key_storage: ReplyKeyStorage,
@@ -61,22 +64,23 @@ where
            ack_recipient,
            input_receiver,
            message_preparer,
-            pending_acks,
+            action_sender,
            real_message_sender,
            topology_access,
            reply_key_storage,
        }
    }

+    // we require topology for replies to generate surb_acks
    async fn handle_reply(&mut self, reply_surb: ReplySURB, data: Vec<u8>) -> Option<RealMessage> {
        let topology_permit = self.topology_access.get_read_permit().await;
-        let topology_ref_option =
-            topology_permit.try_get_valid_topology_ref(&self.ack_recipient, None);
-        if topology_ref_option.is_none() {
-            warn!("Could not process the message - the network topology is invalid");
-            return None;
-        }
-        let topology = topology_ref_option.unwrap();
+        let topology = match topology_permit.try_get_valid_topology_ref(&self.ack_recipient, None) {
+            Some(topology_ref) => topology_ref,
+            None => {
+                warn!("Could not process the message - the network topology is invalid");
+                return None;
+            }
+        };

        match self
            .message_preparer
@@ -86,7 +90,6 @@ where
                // TODO: later probably write pending ack here
                // and deal with them....
                // ... somehow
-
                Some(RealMessage::new(first_hop, sphinx_packet, reply_id))
            }
            Err(err) => {
@@ -105,13 +108,15 @@ where
        with_reply_surb: bool,
    ) -> Vec<RealMessage> {
        let topology_permit = self.topology_access.get_read_permit().await;
-        let topology_ref_option =
-            topology_permit.try_get_valid_topology_ref(&self.ack_recipient, Some(&recipient));
-        if topology_ref_option.is_none() {
-            warn!("Could not process the message - the network topology is invalid");
-            return Vec::new();
-        }
-        let topology = topology_ref_option.unwrap();
+        let topology = match topology_permit
+            .try_get_valid_topology_ref(&self.ack_recipient, Some(&recipient))
+        {
+            Some(topology_ref) => topology_ref,
+            None => {
+                warn!("Could not process the message - the network topology is invalid");
+                return Vec::new();
+            }
+        };

        // split the message, attach optional reply surb
        let (split_message, reply_key) = self
@@ -129,8 +134,6 @@ where
        let mut pending_acks = Vec::with_capacity(split_message.len());
        let mut real_messages = Vec::with_capacity(split_message.len());
        for message_chunk in split_message {
-            // since the paths can be constructed, this CAN'T fail, if it does, there's a bug somewhere
-            let frag_id = message_chunk.fragment_identifier();
            // we need to clone it because we need to keep it in memory in case we had to retransmit
            // it. And then we'd need to recreate entire ACK again.
            let chunk_clone = message_chunk.clone();
@@ -142,29 +145,20 @@ where
            real_messages.push(RealMessage::new(
                prepared_fragment.first_hop_address,
                prepared_fragment.sphinx_packet,
-                frag_id,
+                message_chunk.fragment_identifier(),
            ));

-            let pending_ack = PendingAcknowledgement::new(
+            pending_acks.push(PendingAcknowledgement::new(
                message_chunk,
                prepared_fragment.total_delay,
                recipient.clone(),
-            );
-
-            pending_acks.push((frag_id, pending_ack));
+            ));
        }

-        // first insert pending_acks only then request fragments to be sent, otherwise you might get
-        // some very nasty (and time-consuming to figure out...) race condition.
-        let mut pending_acks_map_write_guard = self.pending_acks.write().await;
-        for (frag_id, pending_ack) in pending_acks.into_iter() {
-            if pending_acks_map_write_guard
-                .insert(frag_id, pending_ack)
-                .is_some()
-            {
-                panic!("Tried to insert duplicate pending ack")
-            }
-        }
+        // tells the controller to put this into the hashmap
+        self.action_sender
+            .unbounded_send(Action::new_insert(pending_acks))
+            .unwrap();

        real_messages
    }
@@ -188,6 +182,7 @@ where
            }
        };

+        // tells real message sender (with the poisson timer) to send this to the mix network
        for real_message in real_messages {
            self.real_message_sender
                .unbounded_send(real_message)
@@ -13,7 +13,7 @@
 // limitations under the License.

 use self::{
-    acknowledgement_listener::AcknowledgementListener,
+    acknowledgement_listener::AcknowledgementListener, action_controller::ActionController,
    input_message_listener::InputMessageListener,
    retransmission_request_listener::RetransmissionRequestListener,
    sent_notification_listener::SentNotificationListener,
@@ -24,59 +24,81 @@ use crate::client::{inbound_messages::InputMessageReceiver, topology_control::To
 use futures::channel::mpsc;
 use gateway_client::AcknowledgementReceiver;
 use log::*;
-use nymsphinx::preparer::MessagePreparer;
 use nymsphinx::{
    acknowledgements::AckKey,
    addressing::clients::Recipient,
    chunking::fragment::{Fragment, FragmentIdentifier},
-    Delay,
+    preparer::MessagePreparer,
+    Delay as SphinxDelay,
 };
 use rand::{CryptoRng, Rng};
-use std::{collections::HashMap, sync::Arc, time::Duration};
-use tokio::{
-    sync::{Notify, RwLock},
-    task::JoinHandle,
+use std::{
+    sync::{Arc, Weak},
+    time::Duration,
 };
+use tokio::task::JoinHandle;

+mod ack_delay_queue;
 mod acknowledgement_listener;
+mod action_controller;
 mod input_message_listener;
 mod retransmission_request_listener;
 mod sent_notification_listener;

-type RetransmissionRequestSender = mpsc::UnboundedSender<FragmentIdentifier>;
-type RetransmissionRequestReceiver = mpsc::UnboundedReceiver<FragmentIdentifier>;
+/// Channel used for indicating that the particular `Fragment` should be retransmitted.
+type RetransmissionRequestSender = mpsc::UnboundedSender<Weak<PendingAcknowledgement>>;

+/// Channel used for receiving data about particular `Fragment` that should be retransmitted.
+type RetransmissionRequestReceiver = mpsc::UnboundedReceiver<Weak<PendingAcknowledgement>>;
+
+/// Channel used for signalling that the particular `Fragment` (associated with the `FragmentIdentifier`)
+/// is done being delayed and is about to be sent to the mix network.
 pub(super) type SentPacketNotificationSender = mpsc::UnboundedSender<FragmentIdentifier>;
+
+/// Channel used for receiving signals about the particular `Fragment` (associated with the `FragmentIdentifier`)
+/// that it is about to be sent to tbe mix network and its timeout timer should be started.
 type SentPacketNotificationReceiver = mpsc::UnboundedReceiver<FragmentIdentifier>;

-type PendingAcksMap = Arc<RwLock<HashMap<FragmentIdentifier, PendingAcknowledgement>>>;
-
-struct PendingAcknowledgement {
+/// Structure representing a data `Fragment` that is on-route to the specified `Recipient`
+#[derive(Debug)]
+pub(crate) struct PendingAcknowledgement {
    message_chunk: Fragment,
-    delay: Delay,
+    delay: SphinxDelay,
    recipient: Recipient,
-    retransmission_cancel: Arc<Notify>,
 }

 impl PendingAcknowledgement {
-    fn new(message_chunk: Fragment, delay: Delay, recipient: Recipient) -> Self {
+    /// Creates new instance of `PendingAcknowledgement` using the provided data.
+    fn new(message_chunk: Fragment, delay: SphinxDelay, recipient: Recipient) -> Self {
        PendingAcknowledgement {
            message_chunk,
            delay,
-            retransmission_cancel: Arc::new(Notify::new()),
            recipient,
        }
    }

-    fn update_delay(&mut self, new_delay: Delay) {
+    fn update_delay(&mut self, new_delay: SphinxDelay) {
        self.delay = new_delay;
    }
 }

+/// AcknowledgementControllerConnectors represents set of channels for communication with
+/// other parts of the system in order to support acknowledgements and retransmission.
 pub(super) struct AcknowledgementControllerConnectors {
+    /// Channel used for forwarding prepared sphinx messages into the poisson sender
+    /// to be sent to the mix network.
    real_message_sender: RealMessageSender,
+
+    /// Channel used for receiving raw messages from a client. The messages need to be put
+    /// into sphinx packets first.
    input_receiver: InputMessageReceiver,
+
+    /// Channel used for receiving notification about particular packet being sent off to the
+    /// mix network (i.e. it was done being delayed by whatever value was determined in the poisson
+    /// sender)
    sent_notifier: SentPacketNotificationReceiver,
+
+    /// Channel used for receiving acknowledgements from the mix network.
    ack_receiver: AcknowledgementReceiver,
 }

@@ -96,6 +118,37 @@ impl AcknowledgementControllerConnectors {
    }
 }

+/// Configurable parameters of the `AcknowledgementController`
+pub(super) struct Config {
+    /// Given ack timeout in the form a * BASE_DELAY + b, it specifies the additive part `b`
+    ack_wait_addition: Duration,
+
+    /// Given ack timeout in the form a * BASE_DELAY + b, it specifies the multiplier `a`
+    ack_wait_multiplier: f64,
+
+    /// Average delay an acknowledgement packet is going to get delayed at a single mixnode.
+    average_ack_delay: Duration,
+
+    /// Average delay a data packet is going to get delayed at a single mixnode.
+    average_packet_delay: Duration,
+}
+
+impl Config {
+    pub(super) fn new(
+        ack_wait_addition: Duration,
+        ack_wait_multiplier: f64,
+        average_ack_delay: Duration,
+        average_packet_delay: Duration,
+    ) -> Self {
+        Config {
+            ack_wait_addition,
+            ack_wait_multiplier,
+            average_ack_delay,
+            average_packet_delay,
+        }
+    }
+}
+
 pub(super) struct AcknowledgementController<R>
 where
    R: CryptoRng + Rng,
@@ -104,6 +157,7 @@ where
    input_message_listener: Option<InputMessageListener<R>>,
    retransmission_request_listener: Option<RetransmissionRequestListener<R>>,
    sent_notification_listener: Option<SentNotificationListener>,
+    action_controller: Option<ActionController>,
 }

 impl<R> AcknowledgementController<R>
@@ -111,67 +165,69 @@ where
    R: 'static + CryptoRng + Rng + Clone + Send,
 {
    pub(super) fn new(
+        config: Config,
        rng: R,
        topology_access: TopologyAccessor,
        ack_key: Arc<AckKey>,
        ack_recipient: Recipient,
        reply_key_storage: ReplyKeyStorage,
-        average_packet_delay: Duration,
-        average_ack_delay: Duration,
-        ack_wait_multiplier: f64,
-        ack_wait_addition: Duration,
        connectors: AcknowledgementControllerConnectors,
    ) -> Self {
-        let pending_acks = Arc::new(RwLock::new(HashMap::new()));
+        let (retransmission_tx, retransmission_rx) = mpsc::unbounded();
+
+        let action_config =
+            action_controller::Config::new(config.ack_wait_addition, config.ack_wait_multiplier);
+        let (action_controller, action_sender) =
+            ActionController::new(action_config, retransmission_tx);
+
        let message_preparer = MessagePreparer::new(
            rng,
            ack_recipient.clone(),
-            average_packet_delay,
-            average_ack_delay,
+            config.average_packet_delay,
+            config.average_ack_delay,
        );

+        // will listen for any acks coming from the network
        let acknowledgement_listener = AcknowledgementListener::new(
            Arc::clone(&ack_key),
            connectors.ack_receiver,
-            Arc::clone(&pending_acks),
+            action_sender.clone(),
        );

+        // will listen for any new messages from the client
        let input_message_listener = InputMessageListener::new(
            Arc::clone(&ack_key),
            ack_recipient.clone(),
            connectors.input_receiver,
            message_preparer.clone(),
-            Arc::clone(&pending_acks),
+            action_sender.clone(),
            connectors.real_message_sender.clone(),
            topology_access.clone(),
            reply_key_storage,
        );

-        let (retransmission_tx, retransmission_rx) = mpsc::unbounded();
-
+        // will listen for any ack timeouts and trigger retransmission
        let retransmission_request_listener = RetransmissionRequestListener::new(
            Arc::clone(&ack_key),
            ack_recipient,
            message_preparer,
-            Arc::clone(&pending_acks),
+            action_sender.clone(),
            connectors.real_message_sender,
            retransmission_rx,
            topology_access,
        );

-        let sent_notification_listener = SentNotificationListener::new(
-            ack_wait_multiplier,
-            ack_wait_addition,
-            connectors.sent_notifier,
-            pending_acks,
-            retransmission_tx,
-        );
+        // will listen for events indicating the packet was sent through the network so that
+        // the retransmission timer should be started.
+        let sent_notification_listener =
+            SentNotificationListener::new(connectors.sent_notifier, action_sender);

        AcknowledgementController {
            acknowledgement_listener: Some(acknowledgement_listener),
            input_message_listener: Some(input_message_listener),
            retransmission_request_listener: Some(retransmission_request_listener),
            sent_notification_listener: Some(sent_notification_listener),
+            action_controller: Some(action_controller),
        }
    }

@@ -181,11 +237,7 @@ where
        let mut retransmission_request_listener =
            self.retransmission_request_listener.take().unwrap();
        let mut sent_notification_listener = self.sent_notification_listener.take().unwrap();
-
-        // TODO: perhaps an extra 'DEBUG' task that would periodically check for stale entries in
-        // pending acks map?
-        // It would only be 'DEBUG' as I don't expect any stale entries to exist there to begin with,
-        // but when can bugs be expected to begin with?
+        let mut action_controller = self.action_controller.take().unwrap();

        // the below are log messages are errors as at the current stage we do not expect any of
        // the task to ever finish. This will of course change once we introduce
@@ -210,6 +262,11 @@ where
            error!("The sent notification listener has finished execution!");
            sent_notification_listener
        });
+        let action_controller_fut = tokio::spawn(async move {
+            action_controller.run().await;
+            error!("The controller has finished execution!");
+            action_controller
+        });

        // technically we don't have to bring `AcknowledgementController` back to a valid state
        // but we can do it, so why not? Perhaps it might be useful if we wanted to allow
@@ -218,6 +275,7 @@ where
        self.input_message_listener = Some(input_listener_fut.await.unwrap());
        self.retransmission_request_listener = Some(retransmission_req_fut.await.unwrap());
        self.sent_notification_listener = Some(sent_notification_fut.await.unwrap());
+        self.action_controller = Some(action_controller_fut.await.unwrap());
    }

    #[allow(dead_code)]
@@ -0,0 +1,141 @@
+// Copyright 2020 Nym Technologies SA
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::action_controller::{Action, ActionSender};
+use super::PendingAcknowledgement;
+use super::RetransmissionRequestReceiver;
+use crate::client::{
+    real_messages_control::real_traffic_stream::{RealMessage, RealMessageSender},
+    topology_control::TopologyAccessor,
+};
+use futures::StreamExt;
+use log::*;
+use nymsphinx::preparer::MessagePreparer;
+use nymsphinx::{acknowledgements::AckKey, addressing::clients::Recipient};
+use rand::{CryptoRng, Rng};
+use std::sync::{Arc, Weak};
+
+// responsible for packet retransmission upon fired timer
+pub(super) struct RetransmissionRequestListener<R>
+where
+    R: CryptoRng + Rng,
+{
+    ack_key: Arc<AckKey>,
+    ack_recipient: Recipient,
+    message_preparer: MessagePreparer<R>,
+    action_sender: ActionSender,
+    real_message_sender: RealMessageSender,
+    request_receiver: RetransmissionRequestReceiver,
+    topology_access: TopologyAccessor,
+}
+
+impl<R> RetransmissionRequestListener<R>
+where
+    R: CryptoRng + Rng,
+{
+    pub(super) fn new(
+        ack_key: Arc<AckKey>,
+        ack_recipient: Recipient,
+        message_preparer: MessagePreparer<R>,
+        action_sender: ActionSender,
+        real_message_sender: RealMessageSender,
+        request_receiver: RetransmissionRequestReceiver,
+        topology_access: TopologyAccessor,
+    ) -> Self {
+        RetransmissionRequestListener {
+            ack_key,
+            ack_recipient,
+            message_preparer,
+            action_sender,
+            real_message_sender,
+            request_receiver,
+            topology_access,
+        }
+    }
+
+    async fn on_retransmission_request(&mut self, timed_out_ack: Weak<PendingAcknowledgement>) {
+        let timed_out_ack = match timed_out_ack.upgrade() {
+            Some(timed_out_ack) => timed_out_ack,
+            None => {
+                debug!("We received an ack JUST as we were about to retransmit [1]");
+                return;
+            }
+        };
+        let packet_recipient = &timed_out_ack.recipient;
+        let chunk_clone = timed_out_ack.message_chunk.clone();
+        let frag_id = chunk_clone.fragment_identifier();
+
+        let topology_permit = self.topology_access.get_read_permit().await;
+        let topology_ref = match topology_permit
+            .try_get_valid_topology_ref(&self.ack_recipient, Some(packet_recipient))
+        {
+            Some(topology_ref) => topology_ref,
+            None => {
+                warn!("Could not retransmit the packet - the network topology is invalid");
+                // we NEED to start timer here otherwise we will have this guy permanently stuck in memory
+                self.action_sender
+                    .unbounded_send(Action::new_start_timer(frag_id))
+                    .unwrap();
+                return;
+            }
+        };
+
+        let prepared_fragment = self
+            .message_preparer
+            .prepare_chunk_for_sending(chunk_clone, topology_ref, &self.ack_key, packet_recipient)
+            .unwrap();
+
+        // if we have the ONLY strong reference to the ack data, it means it was removed from the
+        // pending acks
+        if Arc::strong_count(&timed_out_ack) == 1 {
+            // while we were messing with topology, wrapping data in sphinx, etc. we actually received
+            // this ack after all! no need to retransmit then
+            debug!("We received an ack JUST as we were about to retransmit [2]");
+            return;
+        }
+        // we no longer need the reference - let's drop it so that if somehow `UpdateTimer` action
+        // reached the controller before this function terminated, the controller would not panic.
+        drop(timed_out_ack);
+
+        let new_delay = prepared_fragment.total_delay;
+
+        // We know this update will be reflected by the `StartTimer` Action performed when this
+        // message is sent through the mix network.
+        // Reason being: UpdateTimer is now pushed onto the Action queue and `StartTimer` will
+        // only be pushed when the below `RealMessage` (which we are about to create)
+        // is sent to the `OutQueueControl` and has gone through its internal queue
+        // with the additional poisson delay.
+        // And since Actions are executed in order `UpdateTimer` will HAVE TO be executed before `StartTimer`
+        self.action_sender
+            .unbounded_send(Action::new_update_delay(frag_id, new_delay))
+            .unwrap();
+
+        // send to `OutQueueControl` to eventually send to the mix network
+        self.real_message_sender
+            .unbounded_send(RealMessage::new(
+                prepared_fragment.first_hop_address,
+                prepared_fragment.sphinx_packet,
+                frag_id,
+            ))
+            .unwrap();
+    }
+
+    pub(super) async fn run(&mut self) {
+        debug!("Started RetransmissionRequestListener");
+        while let Some(timed_out_ack) = self.request_receiver.next().await {
+            self.on_retransmission_request(timed_out_ack).await;
+        }
+        error!("TODO: error msg. Or maybe panic?")
+    }
+}
@@ -0,0 +1,63 @@
+// Copyright 2020 Nym Technologies SA
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::action_controller::{Action, ActionSender};
+use super::SentPacketNotificationReceiver;
+use futures::StreamExt;
+use log::*;
+use nymsphinx::chunking::fragment::{FragmentIdentifier, COVER_FRAG_ID};
+
+/// Module responsible for starting up retransmission timers.
+/// It is required because when we send our packet to the `real traffic stream` controlled
+/// by a poisson timer, there's no guarantee the message will be sent immediately, so we might
+/// accidentally fire retransmission way quicker than we should have.
+pub(super) struct SentNotificationListener {
+    sent_notifier: SentPacketNotificationReceiver,
+    action_sender: ActionSender,
+}
+
+impl SentNotificationListener {
+    pub(super) fn new(
+        sent_notifier: SentPacketNotificationReceiver,
+        action_sender: ActionSender,
+    ) -> Self {
+        SentNotificationListener {
+            sent_notifier,
+            action_sender,
+        }
+    }
+
+    async fn on_sent_message(&mut self, frag_id: FragmentIdentifier) {
+        if frag_id == COVER_FRAG_ID {
+            trace!("sent off a cover message - no need to start retransmission timer!");
+            return;
+        } else if frag_id.is_reply() {
+            debug!("sent off a reply message - no need to start retransmission timer!");
+            // TODO: probably there will need to be some extra procedure here, like it would
+            // be nice to know that our reply actually reached the recipient (i.e. we got the ack)
+            return;
+        }
+        self.action_sender
+            .unbounded_send(Action::new_start_timer(frag_id))
+            .unwrap();
+    }
+
+    pub(super) async fn run(&mut self) {
+        debug!("Started SentNotificationListener");
+        while let Some(frag_id) = self.sent_notifier.next().await {
+            self.on_sent_message(frag_id).await;
+        }
+        error!("TODO: error msg. Or maybe panic?")
+    }
+}
@@ -1,144 +0,0 @@
-// Copyright 2020 Nym Technologies SA
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use super::{PendingAcksMap, RetransmissionRequestReceiver};
-use crate::client::{
-    real_messages_control::real_traffic_stream::{RealMessage, RealMessageSender},
-    topology_control::TopologyAccessor,
-};
-use futures::StreamExt;
-use log::*;
-use nymsphinx::preparer::MessagePreparer;
-use nymsphinx::{
-    acknowledgements::AckKey, addressing::clients::Recipient,
-    chunking::fragment::FragmentIdentifier,
-};
-use rand::{CryptoRng, Rng};
-use std::sync::Arc;
-
-// responsible for packet retransmission upon fired timer
-pub(super) struct RetransmissionRequestListener<R>
-where
-    R: CryptoRng + Rng,
-{
-    ack_key: Arc<AckKey>,
-    ack_recipient: Recipient,
-    message_preparer: MessagePreparer<R>,
-    pending_acks: PendingAcksMap,
-    real_message_sender: RealMessageSender,
-    request_receiver: RetransmissionRequestReceiver,
-    topology_access: TopologyAccessor,
-}
-
-impl<R> RetransmissionRequestListener<R>
-where
-    R: CryptoRng + Rng,
-{
-    pub(super) fn new(
-        ack_key: Arc<AckKey>,
-        ack_recipient: Recipient,
-        message_preparer: MessagePreparer<R>,
-        pending_acks: PendingAcksMap,
-        real_message_sender: RealMessageSender,
-        request_receiver: RetransmissionRequestReceiver,
-        topology_access: TopologyAccessor,
-    ) -> Self {
-        RetransmissionRequestListener {
-            ack_key,
-            ack_recipient,
-            message_preparer,
-            pending_acks,
-            real_message_sender,
-            request_receiver,
-            topology_access,
-        }
-    }
-
-    async fn on_retransmission_request(&mut self, frag_id: FragmentIdentifier) {
-        let pending_acks_map_read_guard = self.pending_acks.read().await;
-
-        let unreceived_ack_fragment = match pending_acks_map_read_guard.get(&frag_id) {
-            Some(pending_ack) => pending_ack,
-            // this can actually happen when ack retransmission times out while `on_ack` is being processed
-            // 1. `retransmission_sender.unbounded_send(frag_id).unwrap()` happens thus triggering this function
-            // 2. at the same time ack is received and fully processed (which takes pending_acks *WRITE* lock!!) -> ack is removed from the map + `self.pending_acks.read()` blocks
-            // 3. `on_retransmission_request` manages to get read lock, but the entry was already removed
-            None => {
-                info!("wanted to retransmit ack'd fragment");
-                return;
-            }
-        };
-
-        let packet_recipient = unreceived_ack_fragment.recipient.clone();
-        let chunk_clone = unreceived_ack_fragment.message_chunk.clone();
-        let frag_id = unreceived_ack_fragment.message_chunk.fragment_identifier();
-
-        // TODO: we need some proper benchmarking here to determine whether it could
-        // be more efficient to just get write lock and keep it while doing sphinx computation,
-        // but my gut feeling tells me we should re-acquire it.
-        drop(pending_acks_map_read_guard);
-
-        let topology_permit = self.topology_access.get_read_permit().await;
-        let topology_ref_option = topology_permit
-            .try_get_valid_topology_ref(&self.ack_recipient, Some(&packet_recipient));
-        if topology_ref_option.is_none() {
-            warn!("Could not retransmit the packet - the network topology is invalid");
-            // TODO: perhaps put back into pending acks and reset the timer?
-            return;
-        }
-        let topology_ref = topology_ref_option.unwrap();
-
-        let prepared_fragment = self
-            .message_preparer
-            .prepare_chunk_for_sending(chunk_clone, topology_ref, &self.ack_key, &packet_recipient)
-            .unwrap();
-
-        // minor optimization to not hold the permit while we no longer need it and might have to block
-        // waiting for the write lock on `pending_acks`
-        drop(topology_permit);
-
-        // for this to actually return a None, the following sequence of events needs to happen:
-        // 0. recall that up until this point we're holding a READ lock, so nobody else can WRITE
-        // 1. `on_retransmission_request` is called - processing takes a while (we need to create SPHINX packet, etc.)
-        // 2. at the same time we receive DELAYED (i.e. post timeout) ack for the packet we are about to retransmit
-        // 3. the procedure to remove the pending ack waits for the WRITE lock and acquires it in the tiny window
-        // between when READ lock is dropped and WRITE lock is reacquired in this method
-        // 4. the pending ack is removed and when WRITE lock is acquired here, `None` is returned
-
-        // TODO: benchmark whether it wouldn't be potentially more efficient to acquire WRITE lock at the very beginning of the method
-        // one major drawback: nobody else could READ while we're preparing two sphinx packets, encrypting data, etc.
-        if let Some(pending_ack) = self.pending_acks.write().await.get_mut(&frag_id) {
-            pending_ack.update_delay(prepared_fragment.total_delay);
-
-            self.real_message_sender
-                .unbounded_send(RealMessage::new(
-                    prepared_fragment.first_hop_address,
-                    prepared_fragment.sphinx_packet,
-                    frag_id,
-                ))
-                .unwrap();
-        } else {
-            // later on we will want this to be decreased to 'debug' (or maybe not?)
-            info!("received an ack after timeout, but before retransmission went through")
-        }
-    }
-
-    pub(super) async fn run(&mut self) {
-        debug!("Started RetransmissionRequestListener");
-        while let Some(frag_id) = self.request_receiver.next().await {
-            self.on_retransmission_request(frag_id).await;
-        }
-        error!("TODO: error msg. Or maybe panic?")
-    }
-}
@@ -1,113 +0,0 @@
-// Copyright 2020 Nym Technologies SA
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use super::{PendingAcksMap, RetransmissionRequestSender, SentPacketNotificationReceiver};
-use futures::StreamExt;
-use log::*;
-use nymsphinx::chunking::fragment::{FragmentIdentifier, COVER_FRAG_ID};
-use std::sync::Arc;
-use std::time::Duration;
-
-// responsible for starting and controlling retransmission timers
-// it is required because when we send our packet to the `real traffic stream` controlled
-// with poisson timer, there's no guarantee the message will be sent immediately, so we might
-// accidentally fire retransmission way quicker than we would have wanted.
-pub(super) struct SentNotificationListener {
-    ack_wait_multiplier: f64,
-    ack_wait_addition: Duration,
-    sent_notifier: SentPacketNotificationReceiver,
-    pending_acks: PendingAcksMap,
-    retransmission_sender: RetransmissionRequestSender,
-}
-
-impl SentNotificationListener {
-    pub(super) fn new(
-        ack_wait_multiplier: f64,
-        ack_wait_addition: Duration,
-        sent_notifier: SentPacketNotificationReceiver,
-        pending_acks: PendingAcksMap,
-        retransmission_sender: RetransmissionRequestSender,
-    ) -> Self {
-        SentNotificationListener {
-            ack_wait_multiplier,
-            ack_wait_addition,
-            sent_notifier,
-            pending_acks,
-            retransmission_sender,
-        }
-    }
-
-    async fn on_sent_message(&mut self, frag_id: FragmentIdentifier) {
-        if frag_id == COVER_FRAG_ID {
-            trace!("sent off a cover message - no need to start retransmission timer!");
-            return;
-        } else if frag_id.is_reply() {
-            debug!("sent off a reply message - no need to start retransmission timer!");
-            // TODO: probably there will need to be some extra procedure here, like it would
-            // be nice to know that our reply actually reached the recipient (i.e. we got the ack)
-            return;
-        }
-
-        let pending_acks_map_read_guard = self.pending_acks.read().await;
-        // if the unwrap failed here, we have some weird bug somewhere
-        // although when I think about it, it *theoretically* could happen under extremely heavy client
-        // load that `on_sent_message()` is not called (and we do not receive the read permit)
-        // until we already received and processed an ack for the packet
-        // but this seems extremely unrealistic, but perhaps we should guard against that?
-        let pending_ack_data = match pending_acks_map_read_guard.get(&frag_id) {
-            Some(pending_ack) => pending_ack,
-            None => {
-                info!("on_sent_message: somehow we already received an ack for this packet?");
-                return;
-            }
-        };
-
-        // if this assertion ever fails, we have some bug due to some unintended leak.
-        // the only reason I see it could happen if the `tokio::select` in the spawned
-        // task below somehow did not drop it
-        debug_assert_eq!(
-            Arc::strong_count(&pending_ack_data.retransmission_cancel),
-            1
-        );
-
-        // TODO: read more about Arc::downgrade. it could be useful here
-        let retransmission_cancel = Arc::clone(&pending_ack_data.retransmission_cancel);
-
-        let retransmission_timeout = tokio::time::delay_for(
-            (pending_ack_data.delay.clone() * self.ack_wait_multiplier).to_duration()
-                + self.ack_wait_addition,
-        );
-
-        let retransmission_sender = self.retransmission_sender.clone();
-        tokio::spawn(async move {
-            tokio::select! {
-                _ = retransmission_cancel.notified() => {
-                    trace!("received ack for the fragment. Cancelling retransmission future");
-                }
-                _ = retransmission_timeout => {
-                    trace!("did not receive an ack - will retransmit the packet");
-                    retransmission_sender.unbounded_send(frag_id).unwrap();
-                }
-            }
-        });
-    }
-
-    pub(super) async fn run(&mut self) {
-        debug!("Started SentNotificationListener");
-        while let Some(frag_id) = self.sent_notifier.next().await {
-            self.on_sent_message(frag_id).await;
-        }
-        error!("TODO: error msg. Or maybe panic?")
-    }
-}
@@ -17,9 +17,9 @@
 // OUTPUT: MixMessage to mix traffic

 use self::{
-    acknowlegement_control::AcknowledgementController, real_traffic_stream::OutQueueControl,
+    acknowledgement_control::AcknowledgementController, real_traffic_stream::OutQueueControl,
 };
-use crate::client::real_messages_control::acknowlegement_control::AcknowledgementControllerConnectors;
+use crate::client::real_messages_control::acknowledgement_control::AcknowledgementControllerConnectors;
 use crate::client::reply_key_storage::ReplyKeyStorage;
 use crate::client::{
    inbound_messages::InputMessageReceiver, mix_traffic::MixMessageSender,
@@ -36,9 +36,10 @@ use std::time::Duration;
 use tokio::runtime::Handle;
 use tokio::task::JoinHandle;

-mod acknowlegement_control;
+mod acknowledgement_control;
 mod real_traffic_stream;

+// TODO: ack_key and self_recipient shouldn't really be part of this config
 pub struct Config {
    ack_key: Arc<AckKey>,
    ack_wait_multiplier: f64,
@@ -102,24 +103,32 @@ impl RealMessagesController<OsRng> {
            ack_receiver,
        );

+        let ack_control_config = acknowledgement_control::Config::new(
+            config.ack_wait_addition,
+            config.ack_wait_multiplier,
+            config.average_ack_delay_duration,
+            config.average_packet_delay_duration,
+        );
+
        let ack_control = AcknowledgementController::new(
+            ack_control_config,
            rng,
            topology_access.clone(),
            Arc::clone(&config.ack_key),
            config.self_recipient.clone(),
            reply_key_storage,
-            config.average_packet_delay_duration,
-            config.average_ack_delay_duration,
-            config.ack_wait_multiplier,
-            config.ack_wait_addition,
            ack_controller_connectors,
        );

-        let out_queue_control = OutQueueControl::new(
-            Arc::clone(&config.ack_key),
+        let out_queue_config = real_traffic_stream::Config::new(
            config.average_ack_delay_duration,
            config.average_packet_delay_duration,
            config.average_message_sending_delay,
+        );
+
+        let out_queue_control = OutQueueControl::new(
+            out_queue_config,
+            Arc::clone(&config.ack_key),
            sent_notifier_tx,
            mix_sender,
            real_message_receiver,
@@ -13,7 +13,7 @@
 // limitations under the License.

 use crate::client::mix_traffic::{MixMessage, MixMessageSender};
-use crate::client::real_messages_control::acknowlegement_control::SentPacketNotificationSender;
+use crate::client::real_messages_control::acknowledgement_control::SentPacketNotificationSender;
 use crate::client::topology_control::TopologyAccessor;
 use futures::channel::mpsc;
 use futures::task::{Context, Poll};
@@ -31,13 +31,8 @@ use std::sync::Arc;
 use std::time::Duration;
 use tokio::time;

-pub(crate) struct OutQueueControl<R>
-where
-    R: CryptoRng + Rng,
-{
-    /// Key used to encrypt and decrypt content of an ACK packet.
-    ack_key: Arc<AckKey>,
-
+/// Configurable parameters of the `OutQueueControl`
+pub(crate) struct Config {
    /// Average delay an acknowledgement packet is going to get delay at a single mixnode.
    average_ack_delay: Duration,

@@ -46,6 +41,31 @@ where

    /// Average delay between sending subsequent packets.
    average_message_sending_delay: Duration,
+}
+
+impl Config {
+    pub(crate) fn new(
+        average_ack_delay: Duration,
+        average_packet_delay: Duration,
+        average_message_sending_delay: Duration,
+    ) -> Self {
+        Config {
+            average_ack_delay,
+            average_packet_delay,
+            average_message_sending_delay,
+        }
+    }
+}
+
+pub(crate) struct OutQueueControl<R>
+where
+    R: CryptoRng + Rng,
+{
+    /// Configurable parameters of the `ActionController`
+    config: Config,
+
+    /// Key used to encrypt and decrypt content of an ACK packet.
+    ack_key: Arc<AckKey>,

    /// Channel used for notifying of a real packet being sent out. Used to start up retransmission timer.
    sent_notifier: SentPacketNotificationSender,
@@ -116,7 +136,7 @@ where

        // we know it's time to send a message, so let's prepare delay for the next one
        // Get the `now` by looking at the current `delay` deadline
-        let avg_delay = self.average_message_sending_delay;
+        let avg_delay = self.config.average_message_sending_delay;
        let now = self.next_delay.deadline();
        let next_poisson_delay = sample_poisson_duration(&mut self.rng, avg_delay);

@@ -145,10 +165,8 @@ where
    R: CryptoRng + Rng + Unpin,
 {
    pub(crate) fn new(
+        config: Config,
        ack_key: Arc<AckKey>,
-        average_ack_delay: Duration,
-        average_packet_delay: Duration,
-        average_message_sending_delay: Duration,
        sent_notifier: SentPacketNotificationSender,
        mix_tx: MixMessageSender,
        real_receiver: RealMessageReceiver,
@@ -157,10 +175,8 @@ where
        topology_access: TopologyAccessor,
    ) -> Self {
        OutQueueControl {
+            config,
            ack_key,
-            average_ack_delay,
-            average_packet_delay,
-            average_message_sending_delay,
            sent_notifier,
            next_delay: time::delay_for(Default::default()),
            mix_tx,
@@ -198,8 +214,8 @@ where
                    topology_ref,
                    &*self.ack_key,
                    &self.our_full_destination,
-                    self.average_ack_delay,
-                    self.average_packet_delay,
+                    self.config.average_ack_delay,
+                    self.config.average_packet_delay,
                )
                .expect("Somehow failed to generate a loop cover message with a valid topology");

@@ -209,6 +225,10 @@ where
                // well technically the message was not sent just yet, but now it's up to internal
                // queues and client load rather than the required delay. So realistically we can treat
                // whatever is about to happen as negligible additional delay.
+                trace!(
+                    "{} is about to get sent to the mixnet",
+                    real_message.fragment_id
+                );
                self.sent_notifier
                    .unbounded_send(real_message.fragment_id)
                    .unwrap();
@@ -234,7 +254,7 @@ where
        // we should set initial delay only when we actually start the stream
        self.next_delay = time::delay_for(sample_poisson_duration(
            &mut self.rng,
-            self.average_message_sending_delay,
+            self.config.average_message_sending_delay,
        ));

        info!("Starting out queue controller...");
@@ -1,4 +1,5 @@
 [package]
+build = "build.rs"
 name = "socks5"
 version = "0.8.0-dev"
 authors = ["Dave Hrycyszyn <futurechimp@users.noreply.github.com>"]
@@ -17,6 +17,7 @@ use crate::ChunkingError;
 use nymsphinx_params::{SerializedFragmentIdentifier, FRAG_ID_LEN};
 use rand::Rng;
 use std::convert::TryInto;
+use std::fmt::{self, Formatter};

 // Personal reflection: In hindsight I've spent too much time on relatively too little
 // gain here, as even though I might have saved couple of bytes per packet, the gain
@@ -76,6 +77,16 @@ pub struct FragmentIdentifier {
    fragment_position: u8,
 }

+impl fmt::Display for FragmentIdentifier {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Fragment Identifier: id: {} position: {}",
+            self.set_id, self.fragment_position
+        )
+    }
+}
+
 impl FragmentIdentifier {
    // I really dislike how 'hacky' this function seems
    // refer to: https://github.com/nymtech/nym/issues/294 for further discussion