From 020f25b2049f9f4bb2dbbddc9a0e876eadaf5389 Mon Sep 17 00:00:00 2001 From: Moshe Shababo <17073733+moshababo@users.noreply.github.com> Date: Sun, 21 Jan 2024 19:47:05 -0600 Subject: [PATCH] Upgrade test code --- node/Cargo.lock | 1 + node/libs/concurrency/Cargo.toml | 3 + .../src/sync/prunable_mpsc/tests.rs | 127 +++++++++++------- 3 files changed, 84 insertions(+), 47 deletions(-) diff --git a/node/Cargo.lock b/node/Cargo.lock index d632c84f..13b8f8cd 100644 --- a/node/Cargo.lock +++ b/node/Cargo.lock @@ -2585,6 +2585,7 @@ name = "zksync_concurrency" version = "0.1.0" dependencies = [ "anyhow", + "assert_matches", "once_cell", "pin-project", "rand 0.8.5", diff --git a/node/libs/concurrency/Cargo.toml b/node/libs/concurrency/Cargo.toml index ba8ed85d..c517e400 100644 --- a/node/libs/concurrency/Cargo.toml +++ b/node/libs/concurrency/Cargo.toml @@ -19,5 +19,8 @@ tracing.workspace = true tracing-subscriber.workspace = true vise.workspace = true +[dev-dependencies] +assert_matches.workspace = true + [lints] workspace = true \ No newline at end of file diff --git a/node/libs/concurrency/src/sync/prunable_mpsc/tests.rs b/node/libs/concurrency/src/sync/prunable_mpsc/tests.rs index e1a30748..29a81ce7 100644 --- a/node/libs/concurrency/src/sync/prunable_mpsc/tests.rs +++ b/node/libs/concurrency/src/sync/prunable_mpsc/tests.rs @@ -1,59 +1,92 @@ -use std::sync::Arc; +use tokio::time::{Duration, timeout}; use crate::ctx; +use assert_matches::assert_matches; // Test scenario: -// Send two sets of 0..1000 values, in conjunction, while pruning -// so that only one 0..1000 set is expected to remain in the buffer. -// Then, recv to assert the buffer's content. +// 1. Pre-send two sets of 1000 values, so that the first set is expected to be pruned. +// 2. Send a third set of 1000 values in parallel to receiving. #[tokio::test] async fn test_prunable_mpsc() { - use tokio::time::{timeout, Duration}; + let ctx = ctx::test_root(&ctx::RealClock); #[derive(Debug, Clone)] struct ValueType(usize, usize); - let ctx = ctx::test_root(&ctx::RealClock); - - let (sender, mut receiver) = super::channel(|a: &ValueType, b: &ValueType| a.0 == b.0); - - let sender1 = Arc::new(sender); - let sender2 = sender1.clone(); - - let handle1 = tokio::spawn(async move { - let set = 1; - let values = (0..1000).map(|i| ValueType(i, set)); - for value in values { - let _ = sender1.send(value).await; - tokio::task::yield_now().await; - } + let (send, mut recv): ( + super::Sender>, + super::Receiver>, + ) = super::channel(|a: &ValueType, b: &ValueType| { + // Prune values with the same i. + a.1 == b.1 }); - let handle2 = tokio::spawn(async move { - let set = 2; - let values = (0..1000).map(|i| ValueType(i, set)); - for value in values { - let _ = sender2.send(value).await; - tokio::task::yield_now().await; - } - }); - tokio::try_join!(handle1, handle2).unwrap(); - - tokio::spawn(async move { - let mut i = 0; - loop { - let (value, sender) = receiver.recv(&ctx).await.unwrap(); - assert_eq!(value.0, i); - let _ = sender.send(()); - i = i + 1; - if i == 1000 { - assert!( - timeout(Duration::from_secs(0), receiver.recv(&ctx)).await.is_err(), - "recv() is expected to hang since all values have been exhausted" - ); - break; - } + let res: Result<(), ctx::Canceled> = crate::scope::run!(&ctx, |ctx, s| async move { + // Pre-send sets 0 and 1, 1000 values each. + // Set 0 is expected to be pruned and dropped. + let values = (0..2000).map(|i| { + ValueType(i/1000, i%1000) + }); + for val in values { + let res_recv = send.send(val.clone()).await; + s.spawn(async move { + let res = res_recv.recv_or_disconnected(ctx).await; + match val.0 { + // set 0 values are expected to be pruned and dropped. + 0 => assert_matches!(res, Ok(Err(crate::sync::Disconnected))), + // set 1 values are expected to return `Ok(())`. + 1 => assert_matches!(res, Ok(Ok(Ok(())))), + _ => unreachable!() + } + Ok(()) + }); } - }) - .await - .unwrap(); -} + // Send set 2. + s.spawn(async move { + let values = (1000..2000).map(|i| ValueType(2, i)); + for val in values { + let res_recv = send.send(val.clone()).await; + s.spawn(async move { + let res = res_recv.recv_or_disconnected(ctx).await; + let i = val.1; + match val.0 { + // set 2 values are expected to return `Err(i)`. + 2 => assert_matches!(res, Ok(Ok(Err(err))) => { + assert_eq!(err, i); + }), + _ => unreachable!() + }; + Ok(()) + }); + } + Ok(()) + }); + // Receive. + s.spawn(async move { + let mut i = 0; + loop { + let (val, res_send) = recv.recv(ctx).await.unwrap(); + assert_eq!(val.1, i); + match val.0 { + // set 0 is expected to be pruned and dropped. + 0 => unreachable!(), + // Return `Ok(())` for set 1. + 1 => res_send.send(Ok(())).unwrap(), + // Return `Err(i)` for set 2. + 2 => res_send.send(Err(i)).unwrap(), + _ => unreachable!(), + }; + i = i + 1; + if i == 2000 { + assert!( + timeout(Duration::from_secs(0), recv.recv(ctx)).await.is_err(), + "recv() is expected to hang since all values have been exhausted" + ); + break; + } + } + Ok(()) + }); + Ok(()) + }).await; + assert_eq!(Ok(()), res); +} \ No newline at end of file