runtime: add test for runtime drop issue

The test fails without the fix and passes with it.
It is run with rusty_fork in order to have a fresh process, so that we
can set a custom panic hook that aborts the process instead of
unwinding. This makes catching panics of tokio worker threads reliable.

Tokio worker threads panic if the runtime is dropped from an async
context, which is made more likely by the test's two properties:
1) all shared owners of the `Runtime` that are held by the main thread
   are dropped immediately, which makes more likely that the last owner
   is held by a tokio worker thread which executes the future.
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.

Author: wprzytula

scylla-rust-wrapper/src/runtime.rs

@@ -122,3 +122,105 @@ impl Runtimes {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use rusty_fork::rusty_fork_test;
+    use std::net::SocketAddr;
+    use tracing::instrument::WithSubscriber as _;
+
+    use scylla_proxy::RunningProxy;
+
+    use crate::{
+        argconv::str_to_c_str_n,
+        cass_error::CassError,
+        cluster::{cass_cluster_free, cass_cluster_new, cass_cluster_set_contact_points_n},
+        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,
+        },
+    };
+
+    rusty_fork_test! {
+        #![rusty_fork(timeout_ms = 20000)]
+        #[test]
+        /// This test aims to verify that the Tokio runtime used by the driver is never dropped
+        /// from the async context. Violation of this property results in a panic.
+        ///
+        /// The property could be violated for example if the last `Arc` reference to the runtime
+        /// was held by the tokio task that executes the future wrapped by `CassFuture`. Then,
+        /// dropping the future would drop the runtime, which is forbidden by Tokio as dropping
+        /// the runtime is a blocking operation.
+        ///
+        /// The test fails without the fix and passes with it.
+        /// It is run with rusty_fork in order to have a fresh process, so that we
+        /// can set a custom panic hook that aborts the process instead of
+        /// unwinding. This makes catching panics of tokio worker threads reliable.
+        ///
+        /// Tokio worker threads panic if the runtime is dropped from an async
+        /// context, which is made more likely by the test's two properties:
+        /// 1. all shared owners of the `Runtime` that are held by the main thread
+        ///    are dropped immediately, which makes more likely that the last owner
+        ///    is held by a tokio worker thread which executes the future.
+        /// 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() {
+            setup_tracing();
+
+            std::panic::set_hook(Box::new(|panic_info| {
+                let payload = panic_info.payload();
+                // Support both &str and String panic payloads.
+                let payload_str = payload
+                    .downcast_ref::<&str>().copied().or_else(|| payload.downcast_ref::<String>().map(String::as_str));
+                payload_str.inspect(|s| eprintln!("Panic occurred: {s}"));
+                std::process::abort();
+            }));
+
+            let runtime = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .unwrap();
+
+            runtime.block_on(test_with_one_proxy(
+                runtime_is_never_dropped_in_async_context_do,
+                mock_init_rules(),
+            )
+            .with_current_subscriber());
+        }
+    }
+
+    fn runtime_is_never_dropped_in_async_context_do(
+        node_addr: SocketAddr,
+        proxy: RunningProxy,
+    ) -> RunningProxy {
+        for _ in 0..100 {
+            unsafe {
+                let ip = node_addr.ip().to_string();
+                let mut cluster_raw = cass_cluster_new();
+                let (c_ip, c_ip_len) = str_to_c_str_n(ip.as_str());
+
+                assert_cass_error_eq!(
+                    cass_cluster_set_contact_points_n(cluster_raw.borrow_mut(), c_ip, c_ip_len),
+                    CassError::CASS_OK
+                );
+                let session_raw = cass_session_new();
+                cass_future_wait_check_and_free(cass_session_connect(
+                    session_raw.borrow(),
+                    cluster_raw.borrow().into_c_const(),
+                ));
+
+                let session_close_fut = cass_session_close(session_raw.borrow());
+                // Intentionally not waiting to increase likelihood of the tokio task dropping the runtime.
+                // cass_future_wait(session_close_fut.borrow());
+                cass_future_free(session_close_fut);
+
+                cass_cluster_free(cluster_raw);
+                cass_session_free(session_raw);
+            }
+        }
+        proxy
+    }
+}