cs_utils/utils/futures/

generic_lines_codec.rs

use serde::{Serialize, de::DeserializeOwned};
use tokio_util::codec::LinesCodec;
use std::{marker::PhantomData, fmt::Debug};

mod encoder;
mod decoder;

#[derive(Debug, Clone)]
pub struct GenericLinesCodec<T: Serialize + DeserializeOwned> {
    length_delimited_codec: LinesCodec,
    _phantom: PhantomData<T>,
}

impl<T: Serialize + DeserializeOwned> GenericLinesCodec<T> {
    pub fn new() -> Self {
        return GenericLinesCodec {
            length_delimited_codec: LinesCodec::new(),
            _phantom: PhantomData,
        };
    }
}

impl<T: Serialize + DeserializeOwned> Default for GenericLinesCodec<T> {
    fn default() -> Self {
        return GenericLinesCodec::new();
    }
}

#[cfg(test)]
mod tests {
    use super::GenericLinesCodec;

    mod object {
        use futures::{StreamExt, SinkExt};
        use rstest::rstest;
        use tokio::io::duplex;
        use tokio_util::codec::Framed;
        use tokio::try_join;
        use serde::{Serialize, Deserialize};

        use crate::{swap, random_number, random_str, random_bool};
        use crate::futures::wait_random;

        use super::GenericLinesCodec;

        #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
        pub struct TestStruct {
            pub id: u64,
            pub name: String,
            pub is_flag: bool,
        }

        fn random_test_struct(data_len: usize) -> TestStruct {
            return TestStruct {
                id: random_number(0..=u64::MAX),
                name: random_str(data_len),
                is_flag: random_bool(),
            };
        }

        fn random_test_structs(count: usize, data_len: usize) -> Vec<TestStruct> {
            return (0..count).map(|_| random_test_struct(data_len)).collect();
        }

        #[rstest]
        #[case(1, true)]
        #[case(2, true)]
        #[case(3, true)]
        #[case(4, true)]
        #[case(5, true)]
        #[case(6, true)]
        #[case(7, true)]
        #[case(8, true)]
        #[case(1, false)]
        #[case(2, false)]
        #[case(3, false)]
        #[case(4, false)]
        #[case(5, false)]
        #[case(6, false)]
        #[case(7, false)]
        #[case(8, false)]
        #[tokio::test]
        async fn can_send_receive_messages_bidirectional(
            #[case] data_len: usize,
            #[case] is_forward_direction: bool,
        ) {
            let data_len = data_len * data_len * data_len * 1024;
            let (stream1, stream2) = duplex(random_number(256..=4096));
            
            let stream1_messages = random_test_structs(random_number(4..8), data_len);
            let stream1_messages1 = stream1_messages.clone();
            let stream1_messages2 = stream1_messages.clone();

            let stream2_messages = random_test_structs(random_number(4..8), data_len);
            let stream2_messages1 = stream2_messages.clone();
            let stream2_messages2 = stream2_messages.clone();

            // maybe swap the direction
            let (stream1, stream2) = if !is_forward_direction {
                swap(stream1, stream2)
            } else {
                (stream1, stream2)
            };

            try_join!(
                tokio::spawn(async move {
                    let framed_stream = Framed::new(
                        stream1,
                        GenericLinesCodec::<TestStruct>::new(),
                    );
                    let (
                        mut sink,
                        mut source,
                    ) = framed_stream.split();

                    try_join!(
                        tokio::spawn(async move {
                            for message in stream1_messages1 {
                                sink
                                    .send(message).await
                                    .expect("Cannot send message.");
            
                                wait_random(5..25).await;
                            }
                        }),
                        tokio::spawn(async move {
                            let mut received_messages = vec![];
            
                            while let Some(maybe_message) = source.next().await {
                                let message = maybe_message
                                    .expect("Failed to unwrap the message.");
            
                                received_messages.push(message);
                                
                                wait_random(5..25).await;

                                // otherwise the test never completes
                                if received_messages.len() == stream2_messages2.len() {
                                    break;
                                }
                            }
            
                            assert_eq!(
                                received_messages,
                                stream2_messages2,
                                "Sent and received messages must match.",
                            );
                        }),
                    ).unwrap();
                }),
                tokio::spawn(async move {
                    let framed_stream = Framed::new(
                        stream2,
                        GenericLinesCodec::<TestStruct>::new(),
                    );
                    let (
                        mut sink,
                        mut source,
                    ) = framed_stream.split();

                    try_join!(
                        tokio::spawn(async move {
                            for message in stream2_messages1 {
                                sink
                                    .send(message).await
                                    .expect("Cannot send message.");
            
                                wait_random(5..25).await;
                            }
                        }),
                        tokio::spawn(async move {
                            let mut received_messages = vec![];
            
                            while let Some(maybe_message) = source.next().await {
                                let message = maybe_message
                                    .expect("Failed to unwrap the message.");
            
                                received_messages.push(message);
                                
                                wait_random(5..25).await;

                                // otherwise the test never completes
                                if received_messages.len() == stream1_messages2.len() {
                                    break;
                                }
                            }
            
                            assert_eq!(
                                received_messages,
                                stream1_messages2,
                                "Sent and received messages must match.",
                            );
                        }),
                    ).unwrap();
                }),
            ).unwrap();
        }
    }

    mod string {
        use futures::{StreamExt, SinkExt};
        use rstest::rstest;
        use tokio::io::duplex;
        use tokio_util::codec::Framed;
        use tokio::try_join;

        use crate::{swap, random_number, random_str_rg};
        use crate::futures::wait_random;

        use super::GenericLinesCodec;

        fn random_strings(count: usize, data_len: usize) -> Vec<String> {
            return (0..count)
                .map(|_| {
                    return format!(
                        "{}\n{}",
                        random_str_rg(1..=data_len/2),
                        random_str_rg(1..=data_len/2)
                    );
                })
                .collect();
        }

        #[rstest]
        #[case(1, true)]
        #[case(2, true)]
        #[case(3, true)]
        #[case(4, true)]
        #[case(5, true)]
        #[case(6, true)]
        #[case(7, true)]
        #[case(8, true)]
        #[case(1, false)]
        #[case(2, false)]
        #[case(3, false)]
        #[case(4, false)]
        #[case(5, false)]
        #[case(6, false)]
        #[case(7, false)]
        #[case(8, false)]
        #[tokio::test]
        async fn can_send_receive_messages_bidirectional(
            #[case] data_len: usize,
            #[case] is_forward_direction: bool,
        ) {
            let data_len = data_len * data_len * 1024;
            let (stream1, stream2) = duplex(random_number(256..=1024));
            
            let stream1_messages = random_strings(random_number(4..8), data_len);
            let stream1_messages1 = stream1_messages.clone();
            let stream1_messages2 = stream1_messages.clone();

            let stream2_messages = random_strings(random_number(4..8), data_len);

            let stream2_messages1 = stream2_messages.clone();
            let stream2_messages2 = stream2_messages.clone();

            // maybe swap the direction
            let (stream1, stream2) = if !is_forward_direction {
                swap(stream1, stream2)
            } else {
                (stream1, stream2)
            };

            try_join!(
                tokio::spawn(async move {
                    let framed_stream = Framed::new(
                        stream1,
                        GenericLinesCodec::<String>::new(),
                    );
                    let (
                        mut sink,
                        mut source,
                    ) = framed_stream.split();

                    try_join!(
                        tokio::spawn(async move {
                            for message in stream1_messages1 {
                                sink
                                    .send(message).await
                                    .expect("Cannot send message.");
            
                                wait_random(5..25).await;
                            }
                        }),
                        tokio::spawn(async move {
                            let mut received_messages = vec![];
            
                            while let Some(maybe_message) = source.next().await {
                                let message = maybe_message
                                    .expect("Failed to unwrap the message.");
            
                                received_messages.push(message);
                                
                                wait_random(5..25).await;

                                // otherwise the test never completes
                                if received_messages.len() == stream2_messages2.len() {
                                    break;
                                }
                            }
            
                            assert_eq!(
                                received_messages,
                                stream2_messages2,
                                "Sent and received messages must match.",
                            );
                        }),
                    ).unwrap();
                }),
                tokio::spawn(async move {
                    let framed_stream = Framed::new(
                        stream2,
                        GenericLinesCodec::<String>::new(),
                    );
                    let (
                        mut sink,
                        mut source,
                    ) = framed_stream.split();

                    try_join!(
                        tokio::spawn(async move {
                            for message in stream2_messages1 {
                                sink
                                    .send(message).await
                                    .expect("Cannot send message.");
            
                                wait_random(5..25).await;
                            }
                        }),
                        tokio::spawn(async move {
                            let mut received_messages = vec![];
            
                            while let Some(maybe_message) = source.next().await {
                                let message = maybe_message
                                    .expect("Failed to unwrap the message.");
            
                                received_messages.push(message);
                                
                                wait_random(5..25).await;

                                // otherwise the test never completes
                                if received_messages.len() == stream1_messages2.len() {
                                    break;
                                }
                            }
            
                            assert_eq!(
                                received_messages,
                                stream1_messages2,
                                "Sent and received messages must match.",
                            );
                        }),
                    ).unwrap();
                }),
            ).unwrap();
        }
    }
}