Make PeerID generational

Side effects: The members of the PeerEvent now need to be private to
make sure the Peer is correctly cleaned up.

Author: LeonMatthes

examples/client.rs

@@ -34,12 +34,13 @@ fn main() -> anyhow::Result<()> {
 
         println!("[client] event: {:#?}", e);
 
-        match e.kind {
-            EventKind::Connect => break e.peer_id,
+        match e.kind() {
+            EventKind::Connect => break e.peer_id(),
             EventKind::Disconnect { data } => {
                 println!(
                     "connection NOT successful, peer: {:?}, reason: {}",
-                    e.peer_id, data
+                    e.peer_id(),
+                    data
                 );
                 std::process::exit(0);
             }

examples/server.rs

@@ -22,14 +22,17 @@ fn main() -> anyhow::Result<()> {
 
     loop {
         // Wait 500 ms for any events.
-        if let Some(Event { kind, .. }) = host
+        if let Some(event) = host
             .service(Duration::from_secs(1))
             .context("service failed")?
         {
-            match kind {
-                EventKind::Connect => println!("new connection!"),
-                EventKind::Disconnect { .. } => println!("disconnect!"),
-                EventKind::Receive { channel_id, packet } => println!(
+            match event.kind() {
+                &EventKind::Connect => println!("new connection!"),
+                &EventKind::Disconnect { .. } => println!("disconnect!"),
+                &EventKind::Receive {
+                    ref channel_id,
+                    ref packet,
+                } => println!(
                     "got packet on channel {}, content: '{}'",
                     channel_id,
                     std::str::from_utf8(packet.data()).unwrap()

src/event.rs

@@ -4,15 +4,14 @@ use enet_sys::{
     _ENetEventType_ENET_EVENT_TYPE_NONE, _ENetEventType_ENET_EVENT_TYPE_RECEIVE,
 };
 
-use crate::{Host, Packet, PeerID};
+use crate::{Host, Packet, Peer, PeerID};
 
 /// This struct represents an event that can occur when servicing an `Host`.
 #[derive(Debug)]
-pub struct Event {
-    /// The peer that this event happened on.
-    pub peer_id: PeerID,
-    /// The type of this event.
-    pub kind: EventKind,
+pub struct Event<'a, T> {
+    peer: &'a mut Peer<T>,
+    peer_id: PeerID,
+    kind: EventKind,
 }
 
 /// The type of an event.
@@ -22,7 +21,7 @@ pub enum EventKind {
     Connect,
     /// Peer has disconnected.
     //
-    /// The data of the peer will be dropped on the next call to Host::service or when the structure is dropped.
+    /// The data of the peer will be dropped when the received `Event` is dropped.
     Disconnect {
         /// The data associated with this event. Usually a reason for disconnection.
         data: u32,
@@ -36,12 +35,13 @@ pub enum EventKind {
     },
 }
 
-impl Event {
-    pub(crate) fn from_sys_event<T>(event_sys: ENetEvent, host: &Host<T>) -> Option<Event> {
+impl<'a, T> Event<'a, T> {
+    pub(crate) fn from_sys_event(event_sys: ENetEvent, host: &'a Host<T>) -> Option<Event<'a, T>> {
         if event_sys.type_ == _ENetEventType_ENET_EVENT_TYPE_NONE {
             return None;
         }
 
+        let peer = unsafe { Peer::new_mut(&mut *event_sys.peer) };
         let peer_id = unsafe { host.peer_id(event_sys.peer) };
         let kind = match event_sys.type_ {
             _ENetEventType_ENET_EVENT_TYPE_CONNECT => EventKind::Connect,
@@ -55,6 +55,61 @@ impl Event {
             _ => panic!("unrecognized event type: {}", event_sys.type_),
         };
 
-        Some(Event { peer_id, kind })
+        Some(Event {
+            peer,
+            peer_id,
+            kind,
+        })
+    }
+
+    /// The peer that this event happened on.
+    pub fn peer(&'_ self) -> &'_ Peer<T> {
+        &*self.peer
+    }
+
+    /// The peer that this event happened on.
+    pub fn peer_mut(&'_ mut self) -> &'_ mut Peer<T> {
+        self.peer
+    }
+
+    /// The `PeerID` of the peer that this event happened on.
+    pub fn peer_id(&self) -> PeerID {
+        self.peer_id
+    }
+
+    /// The type of this event.
+    pub fn kind(&self) -> &EventKind {
+        &self.kind
+    }
+
+    /// Take the EventKind out of this event.
+    /// If this peer is a Disconnect event, it will clean up the Peer.
+    /// See the `Drop` implementation
+    pub fn take_kind(mut self) -> EventKind {
+        // Unfortunately we can't simply take the `kind` out of the Event, as otherwise the `Drop`
+        // implementation would no longer work.
+        // We can however, swap the actual EventKind with an empty EventKind (in this case
+        // Connect).
+        // As the `Drop` implementation will then do nothing, we need to call cleanup_after_disconnect before we do the swap.
+        self.cleanup_after_disconnect();
+
+        let mut kind = EventKind::Connect;
+        std::mem::swap(&mut kind, &mut self.kind);
+        kind
+    }
+
+    fn cleanup_after_disconnect(&mut self) {
+        match self.kind {
+            EventKind::Disconnect { .. } => self.peer.cleanup_after_disconnect(),
+            EventKind::Connect | EventKind::Receive { .. } => {}
+        }
+    }
+}
+
+/// Dropping an `Event` with `EventKind::Disconnect` will clean up the Peer, by dropping
+/// the data associated with the `Peer`, as well as invalidating the `PeerID`.
+impl<'a, T> Drop for Event<'a, T> {
+    fn drop(&mut self) {
+        self.cleanup_after_disconnect();
     }
 }

src/host.rs

@@ -6,7 +6,7 @@ use enet_sys::{
     ENET_PROTOCOL_MAXIMUM_CHANNEL_COUNT,
 };
 
-use crate::{Address, EnetKeepAlive, Error, Event, EventKind, Peer, PeerID};
+use crate::{Address, EnetKeepAlive, Error, Event, Peer, PeerID};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 /// Represents a bandwidth limit or unlimited.
@@ -60,7 +60,6 @@ impl BandwidthLimit {
 /// transmission.
 pub struct Host<T> {
     inner: *mut ENetHost,
-    disconnect_drop: Option<PeerID>,
     _keep_alive: Arc<EnetKeepAlive>,
     _peer_data: PhantomData<*const T>,
 }
@@ -71,7 +70,6 @@ impl<T> Host<T> {
 
         Host {
             inner,
-            disconnect_drop: None,
             _keep_alive,
             _peer_data: PhantomData,
         }
@@ -134,33 +132,41 @@ impl<T> Host<T> {
         unsafe { (*self.inner).peerCount }
     }
 
-    /// Returns a mutable reference to a peer at the index, None if the index is invalid.
+    /// Returns a mutable reference to a peer at the given PeerID, None if the index is invalid.
     pub fn peer_mut(&mut self, idx: PeerID) -> Option<&mut Peer<T>> {
-        if idx.0 >= self.peer_count() {
+        if !(0..self.peer_count() as isize).contains(&idx.index) {
             return None;
         }
 
-        Some(Peer::new_mut(unsafe {
-            &mut *((*self.inner).peers.offset(idx.0 as isize))
-        }))
+        let peer = Peer::new_mut(unsafe { &mut *((*self.inner).peers.offset(idx.index)) });
+        if peer.generation() != idx.generation {
+            return None;
+        }
+
+        Some(peer)
     }
 
-    /// Returns a reference to a peer at the index, None if the index is invalid.
+    /// Returns a reference to a peer at the given PeerID, None if the index is invalid.
     pub fn peer(&self, idx: PeerID) -> Option<&Peer<T>> {
-        if idx.0 >= self.peer_count() {
+        if !(0..self.peer_count() as isize).contains(&idx.index) {
             return None;
         }
 
-        Some(Peer::new(unsafe {
-            &*((*self.inner).peers.offset(idx.0 as isize))
-        }))
+        let peer = Peer::new(unsafe { &*((*self.inner).peers.offset(idx.index)) });
+        if peer.generation() != idx.generation {
+            return None;
+        }
+
+        Some(peer)
     }
 
     pub(crate) unsafe fn peer_id(&self, peer: *mut ENetPeer) -> PeerID {
-        PeerID(
-            (peer as enet_sys::size_t - (*self.inner).peers as enet_sys::size_t)
-                / std::mem::size_of::<ENetPeer>() as enet_sys::size_t,
-        )
+        let index = peer.offset_from((*self.inner).peers);
+        let peer = Peer::<T>::new_mut(&mut *peer);
+        PeerID {
+            index,
+            generation: peer.generation(),
+        }
     }
 
     /// Returns an iterator over all peers connected to this `Host`.
@@ -187,39 +193,16 @@ impl<T> Host<T> {
         peers.iter().map(|peer| Peer::new(&*peer))
     }
 
-    fn drop_disconnected(&mut self) {
-        // Seemingly, the lifetime of an ENetPeer ends when the Disconnect event is received.
-        // However, this is *not really clear* in the ENet docs!
-        // It looks like the Peer *might* live longer, but not shorter, so it should be safe
-        // to destroy the associated data (if any) here.
-        if let Some(idx) = self.disconnect_drop.take() {
-            self.peer_mut(idx)
-                .expect("Invalid PeerID in disconnect_drop in enet::Host")
-                .set_data(None);
-        }
-    }
-
-    fn process_event(&mut self, sys_event: ENetEvent) -> Option<Event> {
-        self.drop_disconnected();
-
-        let event = Event::from_sys_event(sys_event, self);
-        if let Some(Event {
-            peer_id,
-            kind: EventKind::Disconnect { .. },
-        }) = event
-        {
-            self.disconnect_drop = Some(peer_id);
-        }
-
-        event
+    fn process_event(&'_ mut self, sys_event: ENetEvent) -> Option<Event<'_, T>> {
+        Event::from_sys_event(sys_event, self)
     }
 
     /// Maintains this host and delivers an event if available.
     ///
     /// This should be called regularly for ENet to work properly with good performance.
     ///
     /// The function won't block if `timeout` is less than 1ms.
-    pub fn service(&mut self, timeout: Duration) -> Result<Option<Event>, Error> {
+    pub fn service(&'_ mut self, timeout: Duration) -> Result<Option<Event<'_, T>>, Error> {
         // ENetEvent is Copy (aka has no Drop impl), so we don't have to make sure we `mem::forget` it later on
         let mut sys_event = MaybeUninit::uninit();
 
@@ -244,7 +227,7 @@ impl<T> Host<T> {
 
     /// Checks for any queued events on this `Host` and dispatches one if
     /// available
-    pub fn check_events(&mut self) -> Result<Option<Event>, Error> {
+    pub fn check_events(&'_ mut self) -> Result<Option<Event<'_, T>>, Error> {
         // ENetEvent is Copy (aka has no Drop impl), so we don't have to make sure we
         // `mem::forget` it later on
         let mut sys_event = MaybeUninit::uninit();
@@ -298,7 +281,7 @@ impl<T> Drop for Host<T> {
     /// Call the corresponding ENet cleanup-function(s).
     fn drop(&mut self) {
         for peer in self.peers_mut() {
-            peer.set_data(None);
+            peer.drop_raw_data();
         }
 
         unsafe {