Refactor proto::streams::store indices

- Removes incrementing counter, instead just using the StreamId as a
  slab ABA guard.
- Adjusts Ptr::deref to use Store::index, as before it was skipping to
  check the ABA guard.
- Rename fields and types to clarify it actually is an ABA guard.
- Improve panic message in case a dangling Ptr is accessed.

src/proto/streams/store.rs

@@ -11,9 +11,8 @@ use std::ops;
 /// Storage for streams
 #[derive(Debug)]
 pub(super) struct Store {
-    slab: slab::Slab<(StoreId, Stream)>,
+    slab: slab::Slab<Stream>,
-    ids: IndexMap<StreamId, (usize, StoreId)>,
+    ids: IndexMap<StreamId, SlabIndex>,
-    counter: StoreId,
 }
 
 /// "Pointer" to an entry in the store
@@ -25,11 +24,14 @@ pub(super) struct Ptr<'a> {
 /// References an entry in the store.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) struct Key {
-    index: usize,
+    index: SlabIndex,
-    store_id: StoreId,
+    /// Keep the stream ID in the key as an ABA guard, since slab indices
+    /// could be re-used with a new stream.
+    stream_id: StreamId,
 }
 
-type StoreId = usize;
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct SlabIndex(usize);
 
 #[derive(Debug)]
 pub(super) struct Queue<N> {
@@ -62,13 +64,12 @@ pub(super) enum Entry<'a> {
 }
 
 pub(super) struct OccupiedEntry<'a> {
-    ids: indexmap::map::OccupiedEntry<'a, StreamId, (usize, StoreId)>,
+    ids: indexmap::map::OccupiedEntry<'a, StreamId, SlabIndex>,
 }
 
 pub(super) struct VacantEntry<'a> {
-    ids: indexmap::map::VacantEntry<'a, StreamId, (usize, StoreId)>,
+    ids: indexmap::map::VacantEntry<'a, StreamId, SlabIndex>,
-    slab: &'a mut slab::Slab<(StoreId, Stream)>,
+    slab: &'a mut slab::Slab<Stream>,
-    counter: &'a mut usize,
 }
 
 pub(super) trait Resolve {
@@ -82,35 +83,32 @@ impl Store {
         Store {
             slab: slab::Slab::new(),
             ids: IndexMap::new(),
-            counter: 0,
         }
     }
 
     pub fn find_mut(&mut self, id: &StreamId) -> Option<Ptr> {
-        let key = match self.ids.get(id) {
+        let index = match self.ids.get(id) {
             Some(key) => *key,
             None => return None,
         };
 
         Some(Ptr {
             key: Key {
-                index: key.0,
+                index,
-                store_id: key.1,
+                stream_id: *id,
             },
             store: self,
         })
     }
 
     pub fn insert(&mut self, id: StreamId, val: Stream) -> Ptr {
-        let store_id = self.counter;
+        let index = SlabIndex(self.slab.insert(val));
-        self.counter = self.counter.wrapping_add(1);
+        assert!(self.ids.insert(id, index).is_none());
-        let key = self.slab.insert((store_id, val));
-        assert!(self.ids.insert(id, (key, store_id)).is_none());
 
         Ptr {
             key: Key {
-                index: key,
+                index,
-                store_id,
+                stream_id: id,
             },
             store: self,
         }
@@ -126,7 +124,6 @@ impl Store {
             Vacant(e) => Entry::Vacant(VacantEntry {
                 ids: e,
                 slab: &mut self.slab,
-                counter: &mut self.counter,
             }),
         }
     }
@@ -140,12 +137,15 @@ impl Store {
 
         while i < len {
             // Get the key by index, this makes the borrow checker happy
-            let key = *self.ids.get_index(i).unwrap().1;
+            let (stream_id, index) = {
+                let entry = self.ids.get_index(i).unwrap();
+                (*entry.0, *entry.1)
+            };
 
             f(Ptr {
                 key: Key {
-                    index: key.0,
+                    index,
-                    store_id: key.1,
+                    stream_id,
                 },
                 store: self,
             })?;
@@ -178,17 +178,23 @@ impl ops::Index<Key> for Store {
     type Output = Stream;
 
     fn index(&self, key: Key) -> &Self::Output {
-        let slot = self.slab.index(key.index);
+        self.slab
-        assert_eq!(slot.0, key.store_id);
+            .get(key.index.0)
-        &slot.1
+            .filter(|s| s.id == key.stream_id)
+            .unwrap_or_else(|| {
+                panic!("dangling store key for stream_id={:?}", key.stream_id);
+            })
     }
 }
 
 impl ops::IndexMut<Key> for Store {
     fn index_mut(&mut self, key: Key) -> &mut Self::Output {
-        let slot = self.slab.index_mut(key.index);
+        self.slab
-        assert_eq!(slot.0, key.store_id);
+            .get_mut(key.index.0)
-        &mut slot.1
+            .filter(|s| s.id == key.stream_id)
+            .unwrap_or_else(|| {
+                panic!("dangling store key for stream_id={:?}", key.stream_id);
+            })
     }
 }
 
@@ -329,10 +335,12 @@ impl<'a> Ptr<'a> {
     /// Remove the stream from the store
     pub fn remove(self) -> StreamId {
         // The stream must have been unlinked before this point
-        debug_assert!(!self.store.ids.contains_key(&self.id));
+        debug_assert!(!self.store.ids.contains_key(&self.key.stream_id));
 
         // Remove the stream state
-        self.store.slab.remove(self.key.index).1.id
+        let stream = self.store.slab.remove(self.key.index.0);
+        assert_eq!(stream.id, self.key.stream_id);
+        stream.id
     }
 
     /// Remove the StreamId -> stream state association.
@@ -340,7 +348,7 @@ impl<'a> Ptr<'a> {
     /// This will effectively remove the stream as far as the H2 protocol is
     /// concerned.
     pub fn unlink(&mut self) {
-        let id = self.id;
+        let id = self.key.stream_id;
         self.store.ids.remove(&id);
     }
 }
@@ -358,13 +366,13 @@ impl<'a> ops::Deref for Ptr<'a> {
     type Target = Stream;
 
     fn deref(&self) -> &Stream {
-        &self.store.slab[self.key.index].1
+        &self.store[self.key]
     }
 }
 
 impl<'a> ops::DerefMut for Ptr<'a> {
     fn deref_mut(&mut self) -> &mut Stream {
-        &mut self.store.slab[self.key.index].1
+        &mut self.store[self.key]
     }
 }
 
@@ -378,10 +386,11 @@ impl<'a> fmt::Debug for Ptr<'a> {
 
 impl<'a> OccupiedEntry<'a> {
     pub fn key(&self) -> Key {
-        let tup = self.ids.get();
+        let stream_id = *self.ids.key();
+        let index = *self.ids.get();
         Key {
-            index: tup.0,
+            index,
-            store_id: tup.1,
+            stream_id,
         }
     }
 }
@@ -391,16 +400,15 @@ impl<'a> OccupiedEntry<'a> {
 impl<'a> VacantEntry<'a> {
     pub fn insert(self, value: Stream) -> Key {
         // Insert the value in the slab
-        let store_id = *self.counter;
+        let stream_id = value.id;
-        *self.counter = store_id.wrapping_add(1);
+        let index = SlabIndex(self.slab.insert(value));
-        let index = self.slab.insert((store_id, value));
 
         // Insert the handle in the ID map
-        self.ids.insert((index, store_id));
+        self.ids.insert(index);
 
         Key {
             index,
-            store_id,
+            stream_id,
         }
     }
 }