stream_partition/

lib.rs

//! Stream partitioning utilities for splitting a single stream into multiple streams based on keys.
//!
//! This module provides functionality to partition a stream into multiple sub-streams, where each
//! sub-stream contains only items that match a specific key determined by an async function.
//!
//! # Example
//!
//! ```rust
//! use futures::{stream, StreamExt};
//! use futures::future::ready;
//! use stream_partition::StreamPartitionExt;
//!
//! # #[tokio::main]
//! # async fn main() {
//! let stream = stream::iter(vec![1, 2, 3, 4, 5, 6]);
//! let mut partitioner = stream.partition_by(|x| ready(x % 2));
//!
//! // Get odd numbers
//! let mut odd_stream = partitioner.lock().unwrap().get_partition(1);
//! let first_odd = odd_stream.next().await.unwrap();
//! assert_eq!(first_odd, 1);
//! # }
//! ```

use std::{
    collections::{HashMap, VecDeque},
    hash::Hash,
    ops::DerefMut,
    pin::Pin,
    sync::{self, Arc, Mutex},
    task::{Context, Poll, Waker},
};

use futures::{
    Stream,
    future::{self, Future},
};
use pin_project_lite::pin_project;

pin_project! {
    /// A stream that partitions items from an underlying stream into multiple sub-streams
    /// based on keys generated by an async function.
    ///
    /// This struct implements `Stream` and yields `(K, Partitioned<St, K>)` tuples, where
    /// each tuple represents a new partition with its associated key and the sub-stream
    /// for that partition.
    ///
    /// Items from the underlying stream are processed through the key function `f`, and
    /// items with the same key are grouped together into the same partition stream.
    pub struct PartitionBy<St, Fut, F, K>
    where
        St: Stream,
        K: Clone,
    {
        me: sync::Weak<Mutex<PartitionBy<St, Fut, F, K>>>,
        #[pin]
        stream: St,
        f: F,
        #[pin]
        pending_fut: Option<Fut>,
        // Item being processed by the partitioning function
        pending_item: Option<St::Item>,
        // Queues of items pending for each partition key
        pending_items: HashMap<K, VecDeque<St::Item>>,
        // Whether the underlying stream has finished

        #[pin]
        stream_finished: bool,

        // Partitions waiting for items
        partition_wakers: HashMap<K, Vec<Waker>>,

        // Configuration: whether to create new queues for unknown keys or drop items
        allow_new_queues: bool,
    }
}

/// A sub-stream that yields only items matching a specific key from the partitioned stream.
///
/// This stream is created by `PartitionBy` and contains only items from the original stream
/// that produced the associated key when passed through the partitioning function.
///
/// Multiple `Partitioned` streams can exist simultaneously, each filtering for different keys.
/// The streams coordinate through shared state to ensure each item goes to the correct partition.
pub struct Partitioned<St, Fut, F, K>
where
    St: Stream,
    K: Clone,
{
    key: K, // The key this partition represents
    shared_state: Arc<Mutex<PartitionBy<St, Fut, F, K>>>,
}

impl<St, Fut, F, K> Clone for Partitioned<St, Fut, F, K>
where
    St: Stream,
    K: Clone,
{
    fn clone(&self) -> Self {
        Self {
            key: self.key.clone(),
            shared_state: Arc::clone(&self.shared_state),
        }
    }
}

/// Shared state between the main partitioning stream and all partition sub-streams.
///
/// This structure coordinates the distribution of items from the source stream to the
/// appropriate partition streams. It maintains queues of pending items for each key
/// and tracks the overall state of the partitioning operation.
impl<St, Fut, F, K> PartitionBy<St, Fut, F, K>
where
    St: Stream,
    St::Item: Clone,
    F: Fn(&St::Item) -> Fut,
    Fut: Future<Output = K>,
    K: Hash + Eq + Clone,
{
    /// Creates a new `PartitionBy` stream that partitions items using the provided function.
    ///
    /// # Arguments
    ///
    /// * `stream` - The source stream to partition
    /// * `f` - An async function that takes stream items and returns a key for partitioning
    ///
    /// # Type Parameters
    ///
    /// * `St` - The source stream type
    /// * `Fut` - The future type returned by the partitioning function
    /// * `F` - The partitioning function type
    /// * `K` - The key type used for partitioning (must be `Hash + Eq + Clone`)
    fn new(stream: St, f: F) -> Arc<Mutex<Self>> {
        Self::new_with_config(stream, f, true)
    }

    /// Creates a new `PartitionBy` stream with configuration options.
    ///
    /// # Arguments
    ///
    /// * `stream` - The source stream to partition
    /// * `f` - An async function that takes stream items and returns a key for partitioning
    /// * `allow_new_queues` - If true, creates new queues for unknown keys; if false, drops items for unknown keys
    fn new_with_config(stream: St, f: F, allow_new_queues: bool) -> Arc<Mutex<Self>> {
        Arc::new_cyclic(|me| {
            let me = me.clone();
            Mutex::new(Self {
                me,
                pending_items: HashMap::new(),
                stream_finished: false,
                stream,
                f,
                pending_fut: None,
                pending_item: None,
                partition_wakers: HashMap::new(),
                allow_new_queues,
            })
        })
    }

    /// Gets a partition stream for a specific key.
    ///
    /// **Must not be called with the same key twice**, otherwise items will be missed.
    ///
    /// # Arguments
    ///
    /// * `key` - The key for which to get a partition stream
    ///
    /// # Returns
    ///
    /// A `Partitioned` stream that yields only items matching the specified key.
    pub fn get_partition(&mut self, key: K) -> Partitioned<St, Fut, F, K> {
        // Ensures that the buffer for new items for this key exists
        self.pending_items.entry(key.clone()).or_default();

        Partitioned {
            key: key.clone(),
            shared_state: self.me.upgrade().unwrap(),
        }
    }

    /// Pre-creates queues for multiple keys.
    ///
    /// # Arguments
    ///
    /// * `keys` - An iterator of keys to pre-register
    pub fn register_keys<I>(&mut self, keys: I)
    where
        I: IntoIterator<Item = K>,
    {
        for key in keys {
            self.pending_items.entry(key).or_default();
        }
    }

    /// Returns whether new partitions are allowed to be created.
    pub fn allows_new_partitions(&self) -> bool {
        self.allow_new_queues
    }

    /// Sets whether new queues are allowed to be created.
    ///
    /// If set to false, only keys that already have queues (either from previous
    /// calls to `get_partition` or `register_keys`) will receive items. Items
    /// for unknown keys will be dropped.
    pub fn set_allow_new_queues(&mut self, allow: bool) {
        self.allow_new_queues = allow;
    }
}

impl<St, Fut, F, K> PartitionBy<St, Fut, F, K>
where
    St: Stream,
    St::Item: Clone,
    F: Fn(&St::Item) -> Fut,
    Fut: Future<Output = K>,
    K: Hash + Eq + Clone,
{
    fn poll_item(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let mut this = self.project();

        loop {
            // If we have a pending future, poll it first
            // We don't get a new item until the current item is processed, as we can only hold one item at a time for the filter
            if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() {
                match fut.poll(cx) {
                    Poll::Ready(key) => {
                        this.pending_fut.set(None);

                        // Store the pending item with its key
                        {
                            if let Some(item) = this.pending_item.take() {
                                // Only store the item if we allow new partitions or the partition already exists
                                if *this.allow_new_queues || this.pending_items.contains_key(&key) {
                                    this.pending_items
                                        .entry(key.clone())
                                        .or_default()
                                        .push_back(item);
                                }
                            }
                        }

                        if let Some(wakers) = this.partition_wakers.remove(&key) {
                            for waker in wakers {
                                waker.wake_by_ref();
                            }
                        }

                        return Poll::Ready(());
                    }
                    Poll::Pending => return Poll::Pending,
                }
            }

            // Poll the underlying stream for the next item
            match this.stream.as_mut().poll_next(cx) {
                Poll::Ready(Some(item)) => {
                    // Store the item and create a future to determine its key
                    let fut = {
                        *this.pending_item = Some(item);
                        (this.f)(this.pending_item.as_ref().unwrap())
                    };
                    this.pending_fut.set(Some(fut));
                }
                Poll::Ready(None) => {
                    // Stream is finished
                    {
                        this.stream_finished.set(true);

                        // Wake up all waiting partitions
                        for wakers in this.partition_wakers.values() {
                            for waker in wakers {
                                waker.wake_by_ref();
                            }
                        }
                        this.partition_wakers.clear();
                    }
                    return Poll::Ready(());
                }
                Poll::Pending => {
                    // If partitions are waiting but main stream has no more items right now,
                    // we need to return Pending so this task can be woken when more items arrive
                    return Poll::Pending;
                }
            }
        }
    }
}

// Do we need unpin here? see pinarcmutex crate

impl<St, Fut, F, K> Stream for Partitioned<St, Fut, F, K>
where
    St: Stream + std::marker::Unpin,
    St::Item: Clone,
    F: Fn(&St::Item) -> Fut,
    Fut: Future<Output = K> + std::marker::Unpin,
    K: Hash + Eq + Clone,
{
    type Item = St::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // Check shared state for items for this key
        loop {
            let mut state = self.shared_state.lock().unwrap();

            if let Some(queue) = state.pending_items.get_mut(&self.key) {
                if let Some(item) = queue.pop_front() {
                    return Poll::Ready(Some(item));
                }
            }

            // If stream is finished and no items, we're done
            if state.stream_finished {
                return Poll::Ready(None);
            }

            // Register our waker to be notified when items for our key are available
            state
                .partition_wakers
                .entry(self.key.clone())
                .or_default()
                .push(cx.waker().clone());
            let p = Pin::new(state.deref_mut());
            match p.poll_item(cx) {
                Poll::Ready(_) => continue,
                Poll::Pending => return Poll::Pending,
            }
        }
    }
}

/// Extension trait that adds partitioning functionality to any stream.
///
/// This trait provides the `partition_by` method that can be called on any stream
/// to create a partitioned stream that splits items based on keys generated by
/// an async function.
pub trait StreamPartitionExt: Stream {
    /// Partitions this stream into multiple sub-streams based on keys generated by an async function.
    ///
    /// Returns a stream of `(K, Partitioned<Self, K>)` tuples, where each tuple represents
    /// a new partition. The first element is the key, and the second is a stream that will
    /// yield only items from the original stream that produce that key.
    ///
    /// # Arguments
    ///
    /// * `f` - An async function that takes stream items and returns a partitioning key
    ///
    /// # Type Parameters
    ///
    /// * `F` - The partitioning function type
    /// * `Fut` - The future type returned by the partitioning function
    /// * `K` - The key type used for partitioning (must be `Hash + Eq + Clone`)
    ///
    /// # Example
    ///
    /// ```rust
    /// use futures::{stream, StreamExt};
    /// use futures::future::ready;
    /// use stream_partition::StreamPartitionExt;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let numbers = stream::iter(vec![1, 2, 3, 4, 5, 6]);
    /// let mut partitioner = numbers.partition_by(|x| ready(x % 2));
    /// # }
    ///
    /// ```
    fn partition_by<F, Fut, K>(self, f: F) -> Arc<Mutex<PartitionBy<Self, Fut, F, K>>>
    where
        Self: Sized,
        Self::Item: Clone,
        F: Fn(&Self::Item) -> Fut,
        Fut: future::Future<Output = K>,
        K: Hash + Eq + Clone,
    {
        PartitionBy::new(self, f)
    }

    /// Partitions this stream with configuration options.
    ///
    /// Like `partition_by`, but allows specifying whether new partitions should be
    /// created automatically or if items for unknown keys should be dropped.
    ///
    /// # Arguments
    ///
    /// * `f` - An async function that takes stream items and returns a partitioning key
    /// * `allow_new_queues` - If true, creates new queues for unknown keys; if false, drops items for unknown keys
    ///
    /// # Example
    ///
    /// ```rust
    /// use futures::{stream, StreamExt};
    /// use futures::future::ready;
    /// use stream_partition::StreamPartitionExt;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    /// let numbers = stream::iter(vec![1, 2, 3, 4, 5, 6]);
    /// // Only allow partitions for pre-registered keys
    /// let mut partitioner = numbers.partition_by_with_config(|x| ready(x % 2), false);
    ///
    /// // Pre-register the keys we want to allow
    /// partitioner.lock().unwrap().register_keys([0, 1]);
    /// # }
    /// ```
    fn partition_by_with_config<F, Fut, K>(
        self,
        f: F,
        allow_new_queues: bool,
    ) -> Arc<Mutex<PartitionBy<Self, Fut, F, K>>>
    where
        Self: Sized,
        Self::Item: Clone,
        F: Fn(&Self::Item) -> Fut,
        Fut: future::Future<Output = K>,
        K: Hash + Eq + Clone,
    {
        PartitionBy::new_with_config(self, f, allow_new_queues)
    }
}

impl<St: Stream> StreamPartitionExt for St {}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use futures::{StreamExt, future::join, stream};
    use stream_throttle::{ThrottlePool, ThrottleRate, ThrottledStream};

    use super::*;

    #[tokio::test]
    async fn test_partition_single() {
        //! Test that demonstrates partitioning a stream of numbers into odd/even partitions.
        //!
        //! This test creates a stream of integers 1-6 and partitions them by their remainder
        //! when divided by 2 (i.e., odd vs even). It verifies that the first partition
        //! encountered (for odd numbers) correctly yields the first odd number (1).
        use futures::future::ready;

        let stream = stream::iter(vec![1, 2, 3, 4, 5, 6]);
        let partitioner = stream.partition_by(|x| ready(x % 2));
        println!("created PatritionBy");

        let mut partition_stream = partitioner.lock().unwrap().get_partition(1);
        println!("created Partition");

        let first_item = partition_stream.next().await.unwrap();
        assert_eq!(first_item, 1);
        println!("Got item");

        assert_eq!(partition_stream.next().await.unwrap(), 3);

        while let Some(v) = partition_stream.next().await {
            assert!(v % 2 == 1, "Expected odd number, got {}", v);
        }
    }

    #[tokio::test]
    async fn test_get_partition() {
        //! Test that demonstrates getting a specific partition by key.
        //!
        //! This test creates a stream of integers and uses get_partition to directly
        //! access the even numbers partition (key = 0) without waiting for it to
        //! appear naturally in the partitioner stream.
        use futures::future::ready;

        let rate = ThrottleRate::new(2, Duration::new(0, 10));
        let pool = ThrottlePool::new(rate);
        let stream = stream::iter(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).throttle(pool);
        let partitioner = stream.partition_by(|x| ready(x % 2));

        // Get the even numbers partition directly
        let mut even_partition = partitioner.lock().unwrap().get_partition(0);
        assert_eq!(even_partition.key, 0);
        let mut odd_partition = partitioner.lock().unwrap().get_partition(1);
        assert_eq!(odd_partition.key, 1);

        let a = tokio::spawn(async move {
            dbg!("a");
            while let Some(v) = even_partition.next().await {
                assert!(dbg!(v) % 2 == 0, "Expected even number, got {}", v);
            }
        });
        let b = tokio::spawn(async move {
            dbg!("b");
            while let Some(v) = odd_partition.next().await {
                assert!(dbg!(v) % 2 == 1, "Expected odd number, got {}", v);
            }
        });

        if tokio::time::timeout(Duration::from_millis(10), join(a, b))
            .await
            .is_err()
        {
            println!("did not complete within 10 ms");
        }
        dbg!("complete");
    }

    #[tokio::test]
    async fn test_get_partition_single_task() {
        use futures::future::ready;

        let rate = ThrottleRate::new(2, Duration::new(0, 10));
        let pool = ThrottlePool::new(rate);
        let stream = stream::iter(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).throttle(pool);
        let partitioner = stream.partition_by(|x| ready(x % 2));

        // Get the even numbers partition directly
        let mut even_partition = partitioner.lock().unwrap().get_partition(0);
        assert_eq!(even_partition.key, 0);
        let mut odd_partition = partitioner.lock().unwrap().get_partition(1);
        assert_eq!(odd_partition.key, 1);

        let a = async move {
            dbg!("a");
            while let Some(v) = even_partition.next().await {
                assert!(dbg!(v) % 2 == 0, "Expected even number, got {}", v);
            }
        };
        let b = async move {
            dbg!("b");
            while let Some(v) = odd_partition.next().await {
                assert!(dbg!(v) % 2 == 1, "Expected odd number, got {}", v);
            }
        };

        if tokio::time::timeout(Duration::from_millis(10), join(a, b))
            .await
            .is_err()
        {
            println!("did not complete within 10 ms");
        }
        dbg!("complete");
    }

    #[tokio::test]
    async fn test_drop_unknown_keys() {
        //! Test that items for unknown keys are dropped when allow_new_queues is false.
        use futures::future::ready;

        let stream = stream::iter(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
        let partitioner = stream.partition_by_with_config(|x| ready(x % 3), false);

        // Pre-register only keys 0 and 1, but not 2
        partitioner.lock().unwrap().register_keys([0, 1]);

        // Get partitions for registered keys
        let mut partition_0 = partitioner.lock().unwrap().get_partition(0);
        let mut partition_1 = partitioner.lock().unwrap().get_partition(1);

        // Collect all items from both partitions using separate tasks
        let task_0 = tokio::spawn(async move {
            let mut items = Vec::new();
            while let Some(item) = partition_0.next().await {
                items.push(item);
            }
            items
        });

        let task_1 = tokio::spawn(async move {
            let mut items = Vec::new();
            while let Some(item) = partition_1.next().await {
                items.push(item);
            }
            items
        });

        // Wait for both tasks to complete with a timeout
        let result = tokio::time::timeout(
            Duration::from_millis(10),
            futures::future::join(task_0, task_1),
        )
        .await;

        if result.is_err() {
            panic!("Collection timed out - streams did not complete");
        }

        let (items_0_result, items_1_result) = result.unwrap();
        let items_0 = items_0_result.unwrap();
        let items_1 = items_1_result.unwrap();

        // Verify we only got items for keys 0 and 1 (multiples of 3 and remainder 1)
        // Items with remainder 2 should have been dropped
        println!("Items for key 0: {:?}", items_0);
        println!("Items for key 1: {:?}", items_1);

        // Expected items: 3, 6, 9 for key 0; 1, 4, 7, 10 for key 1
        // Items 2, 5, 8 should be dropped
        assert!(!items_0.is_empty(), "Should have items for key 0");
        assert!(!items_1.is_empty(), "Should have items for key 1");

        for item in &items_0 {
            assert_eq!(
                item % 3,
                0,
                "All items in partition 0 should have remainder 0"
            );
        }

        for item in &items_1 {
            assert_eq!(
                item % 3,
                1,
                "All items in partition 1 should have remainder 1"
            );
        }

        let mut partition_2 = partitioner.lock().unwrap().get_partition(2);
        // This partition should yield no items since we didn't register key 2 before the stream was finished
        let mut items_2 = Vec::new();
        while let Some(item) = partition_2.next().await {
            items_2.push(item);
        }
        // Verify that no items were collected for key 2
        assert!(
            items_2.is_empty(),
            "Should not have collected any items for key 2"
        );
        println!("Items for key 2: {:?}", items_2);
    }
}

#[cfg(test)]
mod memory_tests {
    // #[global_allocator]
    // static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

    use super::*;
    use futures::{StreamExt, future::ready, stream};

    #[tokio::test]
    async fn test_memory_usage() {
        for _ in 0..10 {
            // Your partitioning code here
            let stream = stream::iter(0..10_000);
            let partitioner = stream.partition_by(|x| ready(x % 100));

            // Consume all partitions
            let mut handles = Vec::new();
            for key in 0..100 {
                let mut partition = partitioner.lock().unwrap().get_partition(key);
                handles.push(tokio::spawn(async move {
                    while (partition.next().await).is_some() {}
                }));
            }
            for handle in handles {
                handle.await.unwrap();
            }
        }
    }
}