1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
use std::{collections::HashMap, hash::Hash};

use crate::{
    messages::message_manager::MessageManager, types::PacketIndex,
    wrapping_number::sequence_greater_than, HostWorldManager, LocalWorldManager,
};

use super::{
    packet_notifiable::PacketNotifiable, packet_type::PacketType, sequence_buffer::SequenceBuffer,
    standard_header::StandardHeader,
};

pub const REDUNDANT_PACKET_ACKS_SIZE: u16 = 32;
const DEFAULT_SEND_PACKETS_SIZE: usize = 256;

/// Keeps track of sent & received packets, and contains ack information that is
/// copied into the standard header on each outgoing packet
pub struct AckManager {
    // Local packet index which we'll bump each time we send a new packet over the network.
    next_packet_index: PacketIndex,
    // The last acked packet index of the packets we've sent to the remote host.
    last_recv_packet_index: PacketIndex,
    // Using a `Hashmap` to track every packet we send out so we can ensure that we can resend when
    // dropped.
    sent_packets: HashMap<PacketIndex, SentPacket>,
    // However, we can only reasonably ack up to `REDUNDANT_PACKET_ACKS_SIZE + 1` packets on each
    // message we send so this should be that large.
    received_packets: SequenceBuffer<ReceivedPacket>,
}

impl AckManager {
    pub fn new() -> Self {
        Self {
            next_packet_index: 0,
            last_recv_packet_index: u16::MAX,
            sent_packets: HashMap::with_capacity(DEFAULT_SEND_PACKETS_SIZE),
            received_packets: SequenceBuffer::with_capacity(REDUNDANT_PACKET_ACKS_SIZE + 1),
        }
    }

    /// Get the index of the next outgoing packet
    pub fn next_sender_packet_index(&self) -> PacketIndex {
        self.next_packet_index
    }

    /// Process an incoming packet, handle notifications of delivered / dropped
    /// packets
    pub fn process_incoming_header<E: Copy + Eq + Hash + Send + Sync>(
        &mut self,
        header: &StandardHeader,
        message_manager: &mut MessageManager,
        host_world_manager: &mut HostWorldManager<E>,
        local_world_manager: &mut LocalWorldManager<E>,
        packet_notifiables: &mut [&mut dyn PacketNotifiable],
    ) {
        let sender_packet_index = header.sender_packet_index;
        let sender_ack_index = header.sender_ack_index;
        let mut sender_ack_bitfield = header.sender_ack_bitfield;

        self.received_packets
            .insert(sender_packet_index, ReceivedPacket {});

        // ensure that `self.sender_ack_index` is always increasing (with
        // wrapping)
        if sequence_greater_than(sender_ack_index, self.last_recv_packet_index) {
            self.last_recv_packet_index = sender_ack_index;
        }

        // the current `sender_ack_index` was (clearly) received so we should remove it
        if let Some(sent_packet) = self.sent_packets.get(&sender_ack_index) {
            if sent_packet.packet_type == PacketType::Data {
                self.notify_packet_delivered(
                    sender_ack_index,
                    message_manager,
                    host_world_manager,
                    local_world_manager,
                    packet_notifiables,
                );
            }

            self.sent_packets.remove(&sender_ack_index);
        }

        // The `sender_ack_bitfield` is going to include whether or not the past 32
        // packets have been received successfully.
        // If so, we have no need to resend old packets.
        for i in 1..=REDUNDANT_PACKET_ACKS_SIZE {
            let sent_packet_index = sender_ack_index.wrapping_sub(i);
            if let Some(sent_packet) = self.sent_packets.get(&sent_packet_index) {
                if sender_ack_bitfield & 1 == 1 {
                    if sent_packet.packet_type == PacketType::Data {
                        self.notify_packet_delivered(
                            sent_packet_index,
                            message_manager,
                            host_world_manager,
                            local_world_manager,
                            packet_notifiables,
                        );
                    }

                    self.sent_packets.remove(&sent_packet_index);
                } else {
                    self.sent_packets.remove(&sent_packet_index);
                }
            }

            sender_ack_bitfield >>= 1;
        }
    }

    /// Records the packet with the given packet index
    fn track_packet(&mut self, packet_type: PacketType, packet_index: PacketIndex) {
        self.sent_packets
            .insert(packet_index, SentPacket { packet_type });
    }

    /// Bumps the local packet index
    fn increment_local_packet_index(&mut self) {
        self.next_packet_index = self.next_packet_index.wrapping_add(1);
    }

    pub fn next_outgoing_packet_header(&mut self, packet_type: PacketType) -> StandardHeader {
        let next_packet_index = self.next_sender_packet_index();

        let outgoing = StandardHeader::new(
            packet_type,
            next_packet_index,
            self.last_received_packet_index(),
            self.ack_bitfield(),
        );

        self.track_packet(packet_type, next_packet_index);
        self.increment_local_packet_index();

        outgoing
    }

    fn notify_packet_delivered<E: Copy + Eq + Hash + Send + Sync>(
        &self,
        sent_packet_index: PacketIndex,
        message_manager: &mut MessageManager,
        host_world_manager: &mut HostWorldManager<E>,
        local_world_manager: &mut LocalWorldManager<E>,
        packet_notifiables: &mut [&mut dyn PacketNotifiable],
    ) {
        message_manager.notify_packet_delivered(sent_packet_index);
        host_world_manager.notify_packet_delivered(sent_packet_index, local_world_manager);
        for notifiable in packet_notifiables {
            notifiable.notify_packet_delivered(sent_packet_index);
        }
    }

    fn last_received_packet_index(&self) -> PacketIndex {
        self.received_packets.sequence_num().wrapping_sub(1)
    }

    fn ack_bitfield(&self) -> u32 {
        let last_received_remote_packet_index: PacketIndex = self.last_received_packet_index();
        let mut ack_bitfield: u32 = 0;
        let mut mask: u32 = 1;

        // iterate the past `REDUNDANT_PACKET_ACKS_SIZE` received packets and set the
        // corresponding bit for each packet which exists in the buffer.
        for i in 1..=REDUNDANT_PACKET_ACKS_SIZE {
            let received_packet_index = last_received_remote_packet_index.wrapping_sub(i);
            if self.received_packets.exists(received_packet_index) {
                ack_bitfield |= mask;
            }
            mask <<= 1;
        }

        ack_bitfield
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SentPacket {
    pub packet_type: PacketType,
}

#[derive(Clone, Debug, Default)]
pub struct ReceivedPacket;