testing: introduce rusty_fork_test_with_proxy!

The macro is similar to rusty_fork_test!, but it also passes a unique
IP address to the test function, which can be used to set up one proxy
node without risk of address collisions.

Without this macro, tests that use test_with_one_proxy() directly
and run in subprocesses would sometimes fail due to address collisions
with the main process or other subprocesses.

Author: wprzytula

scylla-rust-wrapper/src/runtime.rs

@@ -125,8 +125,7 @@ impl Runtimes {
 
 #[cfg(test)]
 mod tests {
-    use rusty_fork::rusty_fork_test;
-    use std::net::SocketAddr;
+    use std::{net::SocketAddr, process::Command};
     use tracing::instrument::WithSubscriber as _;
 
     use scylla_proxy::RunningProxy;
@@ -138,12 +137,12 @@ mod tests {
         future::cass_future_free,
         session::{cass_session_close, cass_session_connect, cass_session_free, cass_session_new},
         testing::{
-            assert_cass_error_eq, cass_future_wait_check_and_free, mock_init_rules, setup_tracing,
-            test_with_one_proxy,
+            assert_cass_error_eq, cass_future_wait_check_and_free, mock_init_rules,
+            rusty_fork_test_with_proxy, setup_tracing, test_with_one_proxy_at_ip,
         },
     };
 
-    rusty_fork_test! {
+    rusty_fork_test_with_proxy! {
         #![rusty_fork(timeout_ms = 20000)]
         #[test]
         /// This test aims to verify that the Tokio runtime used by the driver is never dropped
@@ -167,7 +166,7 @@ mod tests {
         /// 2. the case is repeated 100 times in a loop, which increases the
         ///    likelihood of the runtime being dropped from an async context.
         ///    Without the loop, the test could sometimes pass even without the fix.
-        fn runtime_is_never_dropped_in_async_context() {
+        fn runtime_is_never_dropped_in_async_context(ip: IpAddr) {
             setup_tracing();
 
             std::panic::set_hook(Box::new(|panic_info| {
@@ -184,9 +183,10 @@ mod tests {
                 .build()
                 .unwrap();
 
-            runtime.block_on(test_with_one_proxy(
+            runtime.block_on(test_with_one_proxy_at_ip(
                 runtime_is_never_dropped_in_async_context_do,
                 mock_init_rules(),
+                ip,
             )
             .with_current_subscriber());
         }

scylla-rust-wrapper/src/testing.rs

@@ -2,7 +2,12 @@ use scylla_proxy::{
     Condition, Node, Proxy, Reaction as _, RequestFrame, RequestOpcode, RequestReaction,
     RequestRule, ResponseFrame, RunningProxy,
 };
-use std::{collections::HashMap, ffi::c_char, net::SocketAddr, sync::Arc};
+use std::{
+    collections::HashMap,
+    ffi::c_char,
+    net::{IpAddr, SocketAddr},
+    sync::Arc,
+};
 
 use crate::{
     argconv::{CMut, CassOwnedSharedPtr, ptr_to_cstr_n},
@@ -111,11 +116,12 @@ pub(crate) fn mock_init_rules() -> impl IntoIterator<Item = RequestRule> {
         )))
 }
 
-pub(crate) async fn test_with_one_proxy(
+pub(crate) async fn test_with_one_proxy_at_ip(
     test: impl FnOnce(SocketAddr, RunningProxy) -> RunningProxy + Send + 'static,
     rules: impl IntoIterator<Item = RequestRule>,
+    ip: IpAddr,
 ) {
-    let proxy_addr = SocketAddr::new(scylla_proxy::get_exclusive_local_address(), 9042);
+    let proxy_addr = SocketAddr::new(ip, 9042);
 
     let proxy = Proxy::builder()
         .with_node(
@@ -137,3 +143,107 @@ pub(crate) async fn test_with_one_proxy(
 
     let _ = proxy.finish().await;
 }
+
+pub(crate) async fn test_with_one_proxy(
+    test: impl FnOnce(SocketAddr, RunningProxy) -> RunningProxy + Send + 'static,
+    rules: impl IntoIterator<Item = RequestRule>,
+) {
+    test_with_one_proxy_at_ip(test, rules, scylla_proxy::get_exclusive_local_address()).await
+}
+
+/// Run a Rust test in a subprocess, setting up the proxy
+/// with a correct unique address.
+///
+/// This is intended to coordinate in address allocation between
+/// the parent process and the subprocesses, so that no address
+/// collisions occur when tests are run in multiple processes.
+///
+/// The basic usage is to simply put this macro around your `#[test]`
+/// function and add one argument of type `IpAddr` to it, like so:
+///
+/// ```rust,nocompile
+/// use crate::testing::rusty_fork_test_with_proxy;
+///
+/// rusty_fork_test_with_proxy! {
+///     #[test]
+///     fn my_test(ip: IpAddr) {
+///         // Use the `ip` variable to set up the proxy.
+///         let fut = test_with_one_proxy_at_ip(
+///             test_body,
+///             proxy_rules,
+///             ip,
+///         );
+///
+///         // Run the future to completion.
+///     }
+/// }
+/// ```
+///
+/// The test will be run in its own process. If the subprocess exits
+/// unsuccessfully for any reason, including due to signals, the test fails.
+///
+/// It is also possible to specify a timeout which is applied to the test in
+/// the block, by adding the following line straight before the test:
+///
+/// ```rust,no_compile
+/// #![rusty_fork(timeout_ms = 1000)]
+/// ```
+///
+/// For more details, see the documentation of the `rusty_fork_test` macro
+/// in the [rusty_fork] crate.
+/// ```
+macro_rules! rusty_fork_test_with_proxy {
+    (#![rusty_fork(timeout_ms = $timeout:expr)]
+         $(#[$meta:meta])*
+         fn $test_name:ident($ip:ident: IpAddr) $body:block
+    ) => {
+        $(#[$meta])*
+        fn $test_name() {
+            // Eagerly convert everything to function pointers so that all
+            // tests use the same instantiation of `fork`.
+            let supervise_fn = |child: &mut rusty_fork::ChildWrapper, _file: &mut std::fs::File| {
+                rusty_fork::fork_test::supervise_child(child, $timeout);
+            };
+
+            const ENV_KEY: &str = "PROXY_RUSTY_FORK_IP";
+
+            let process_modifier = |child: &mut Command| {
+                child.env(
+                    ENV_KEY,
+                    scylla_proxy::get_exclusive_local_address().to_string(),
+                );
+            };
+
+            let test = || {
+                // Body expects the ip variable in scope.
+                let $ip: ::std::net::IpAddr = std::env::var(ENV_KEY)
+                    .expect("Parent should have set the proxy address as an env variable.")
+                    .parse()
+                    .expect("Error parsing the proxy address from env variable.");
+
+                $body
+            };
+
+            rusty_fork::fork(
+                rusty_fork::rusty_fork_test_name!($test_name),
+                rusty_fork::rusty_fork_id!(),
+                process_modifier,
+                supervise_fn,
+                test,
+            )
+            .expect("forking test failed")
+        }
+    };
+
+    (
+         $(#[$meta:meta])*
+         fn $test_name:ident() $body:block
+    ) => {
+        rusty_fork_test! {
+            #![rusty_fork(timeout_ms = 0)]
+
+            $(#[$meta])* fn $test_name() $body
+        }
+    };
+}
+pub(crate) use rusty_fork_test_with_proxy;