fsm: streams

Author: nicoburniske

fsm/src/lib.rs

@@ -342,6 +342,15 @@ impl<M: StreamMeta> QlFsm<M> {
         fsm::stream(self, stream_id)
     }
 
+    /// returns the open streams
+    pub fn streams(&mut self) -> impl Iterator<Item = StreamOps<'_, M>> {
+        self.state
+            .link
+            .connected_mut()
+            .into_iter()
+            .flat_map(|connection| connection.session.streams())
+    }
+
     /// queues a ping on the active session
     pub fn queue_ping(&mut self) -> Result<(), NoSessionError> {
         fsm::queue_ping(self)

fsm/src/session/mod.rs

@@ -216,20 +216,31 @@ impl<M: StreamMeta> SessionFsm<M> {
             ),
         );
         let stream_index = self.state.streams.len() - 1;
-        Ok(StreamOps::new(self, stream_id, stream_index))
+        Ok(StreamOps::new(
+            stream_id,
+            self.state.streams.get_index_mut(stream_index).unwrap().1,
+        ))
     }
 
     pub fn stream(&mut self, stream_id: StreamId) -> Result<StreamOps<'_, M>, StreamError> {
-        let Some(stream_index) = (|| {
-            let index = self.state.streams.get_index_of(&stream_id)?;
-            // Event::Opened only fires after we receive the first frame of a stream
-            // prevent early access to streams
-            let _ = self.state.streams[index].io.header.as_ref()?;
-            Some(index)
-        })() else {
-            return Err(StreamError::MissingStream);
-        };
-        Ok(StreamOps::new(self, stream_id, stream_index))
+        // unopened remote streams are not available yet
+        let stream = self
+            .state
+            .streams
+            .get_mut(&stream_id)
+            .filter(|stream| stream.io.header.is_some())
+            .ok_or(StreamError::MissingStream)?;
+        Ok(StreamOps::new(stream_id, stream))
+    }
+
+    pub fn streams(&mut self) -> impl Iterator<Item = StreamOps<'_, M>> {
+        self.state
+            .streams
+            .iter_mut()
+            .filter_map(|(&stream_id, stream)| {
+                stream.io.header.as_ref()?;
+                Some(StreamOps::new(stream_id, stream))
+            })
     }
 
     pub fn queue_ping(&mut self) {

fsm/src/session/stream_ops.rs

@@ -1,51 +1,38 @@
 use ql_common::{ResetCode, StreamId};
 use ql_wire::StreamReset;
 
-use super::{
-    state::{InboundState, OutboundState, StreamIoState, StreamState},
-    SessionFsm,
-};
+use super::state::{InboundState, OutboundState, StreamIoState, StreamState};
 use crate::{CommitReadError, ResetOrigin, StreamMeta, StreamResetEvent, StreamResetTarget};
 
+/// mutable access to one stream
+/// terminal streams are reaped when the session is next polled
 pub struct StreamOps<'a, M: StreamMeta> {
-    session: &'a mut SessionFsm<M>,
     stream_id: StreamId,
-    stream_index: usize,
+    stream: &'a mut StreamState<M>,
 }
 
 impl<'a, M: StreamMeta> StreamOps<'a, M> {
-    pub(super) fn new(
-        session: &'a mut SessionFsm<M>,
-        stream_id: StreamId,
-        stream_index: usize,
-    ) -> Self {
-        Self {
-            session,
-            stream_id,
-            stream_index,
-        }
+    pub(super) fn new(stream_id: StreamId, stream: &'a mut StreamState<M>) -> Self {
+        Self { stream_id, stream }
     }
 
     pub fn metadata(&self) -> &M {
-        &self.stream().metadata
+        &self.stream.metadata
     }
 
     pub fn metadata_mut(&mut self) -> &mut M {
-        &mut self.stream_mut().metadata
+        &mut self.stream.metadata
     }
 
     pub fn io(&mut self) -> StreamIo<'_> {
-        let stream_id = self.stream_id;
-        StreamIo::new(stream_id, &mut self.stream_mut().io)
+        StreamIo::new(self.stream_id, &mut self.stream.io)
     }
 
     /// returns the metadata and stream I/O together
     pub fn split_mut(&mut self) -> (&mut M, StreamIo<'_>) {
-        let stream_id = self.stream_id;
-        let stream = self.stream_mut();
         (
-            &mut stream.metadata,
-            StreamIo::new(stream_id, &mut stream.io),
+            &mut self.stream.metadata,
+            StreamIo::new(self.stream_id, &mut self.stream.io),
         )
     }
 
@@ -64,7 +51,7 @@ impl<'a, M: StreamMeta> StreamOps<'a, M> {
     /// resets the stream and notifies its metadata
     pub fn reset(&mut self, target: StreamResetTarget, code: ResetCode) {
         let stream_id = self.stream_id;
-        let StreamState { metadata, io } = self.stream_mut();
+        let StreamState { metadata, io } = &mut *self.stream;
         let wire_target = match target {
             StreamResetTarget::Reader => io.role.inbound_target(),
             StreamResetTarget::Writer => io.role.outbound_target(),
@@ -91,22 +78,6 @@ impl<'a, M: StreamMeta> StreamOps<'a, M> {
             origin: ResetOrigin::Local,
         });
     }
-
-    #[inline]
-    fn stream(&self) -> &StreamState<M> {
-        &self.session.state.streams[self.stream_index]
-    }
-
-    #[inline]
-    fn stream_mut(&mut self) -> &mut StreamState<M> {
-        &mut self.session.state.streams[self.stream_index]
-    }
-}
-
-impl<M: StreamMeta> Drop for StreamOps<'_, M> {
-    fn drop(&mut self) {
-        self.session.try_reap_stream(self.stream_id);
-    }
 }
 
 /// I/O operations for one stream

fsm/src/session/tests.rs

@@ -976,7 +976,7 @@ fn one_sided_remote_reset_notifies_metadata_and_preserves_stream() {
 }
 
 #[test]
-fn draining_the_last_bytes_reaps_a_terminal_stream_on_drop() {
+fn draining_the_last_bytes_reaps_a_terminal_stream_on_next_poll() {
     let now = Instant::now();
     let mut fsm =
         super::SessionFsm::<TestMeta>::new(SessionConfig::default(), SessionParams::default(), now);
@@ -1010,6 +1010,8 @@ fn draining_the_last_bytes_reaps_a_terminal_stream_on_drop() {
     }
     drop(stream);
 
+    assert!(fsm.state.streams.contains_key(&stream_id));
+    fsm.take_next_write(now);
     assert!(fsm.stream(stream_id).is_err());
 }