Skip to main content

ruda_runtime/runtime/stream/
base.rs

1use alloc::vec::Vec;
2use ruda_core::stream_id::StreamId;
3
4/// Trait for creating streams, used by the stream pool to generate streams as needed.
5pub trait StreamFactory {
6    /// The type of stream produced by this factory.
7    type Stream;
8    /// Creates a new stream instance.
9    fn create(&mut self) -> Self::Stream;
10}
11
12/// Represents a pool of streams, managing a collection of streams created by a factory.
13#[derive(Debug)]
14pub struct StreamPool<F: StreamFactory> {
15    /// Vector storing optional streams, where None indicates an uninitialized stream.
16    streams: Vec<Option<F::Stream>>,
17    /// The factory used to create new streams when needed.
18    factory: F,
19    /// Maximum number of regular streams (excludes special streams).
20    max_streams: usize,
21}
22
23impl<F: StreamFactory> StreamPool<F> {
24    /// Creates a new stream pool with the given backend factory and capacity constraints.
25    pub fn new(backend: F, max_streams: u8, num_special: u8) -> Self {
26        // Initialize a vector with capacity for regular and special streams.
27        let mut streams = Vec::with_capacity(max_streams as usize);
28        // Pre-populate the vector with None to reserve space for all streams.
29        for _ in 0..(max_streams.saturating_add(num_special)) {
30            streams.push(None);
31        }
32
33        Self {
34            streams,
35            factory: backend,
36            max_streams: max_streams as usize,
37        }
38    }
39
40    /// Retrieves a mutable reference to a stream for a given stream ID.
41    pub fn get_mut(&mut self, stream_id: &StreamId) -> &mut F::Stream {
42        // Calculate the index for the stream ID.
43        let index = self.stream_index(stream_id);
44
45        self.get_mut_index(index)
46    }
47
48    /// Initialize an unused regular slot with a backend-specific stream. This
49    /// does not modulo-map IDs or replace a stream that may have pending work.
50    /// Returns false for an occupied slot or an ID outside the regular pool.
51    pub fn try_initialize_with<E>(
52        &mut self, stream_id: &StreamId,
53        initialize: impl FnOnce(&mut F) -> Result<F::Stream, E>,
54    ) -> Result<bool, E> {
55        let Ok(index) = usize::try_from(stream_id.value) else { return Ok(false); };
56        if index >= self.max_streams || self.streams[index].is_some() { return Ok(false); }
57        let stream = initialize(&mut self.factory)?;
58        self.streams[index] = Some(stream);
59        Ok(true)
60    }
61
62    /// Retrieves a stream by index, initializing it if needed.
63    pub fn get_mut_index(&mut self, index: usize) -> &mut F::Stream {
64        let entry = &mut self.streams[index];
65        entry.get_or_insert_with(|| self.factory.create())
66    }
67
68    /// Retrieves a special stream by index, initializing it if needed.
69    pub fn get_special(&mut self, index: u8) -> &mut F::Stream {
70        self.get_mut_index(self.max_streams + index as usize)
71    }
72
73    /// Calculates the index for a given stream ID, mapping it to the pool's capacity.
74    pub fn stream_index(&mut self, id: &StreamId) -> usize {
75        stream_index(id, self.max_streams)
76    }
77}
78
79/// Maps a stream ID to an index within the pool's capacity using modulo arithmetic.
80pub fn stream_index(stream_id: &StreamId, max_streams: usize) -> usize {
81    stream_id.value as usize % max_streams
82}