Refactor runtime decisions into a functional core and fix wake/timeout edge cases

src/async_event.rs

@@ -30,18 +30,103 @@ impl AsyncEvent {
 
     #[cfg_attr(feature = "profile", hotpath::measure(impl_type = "AsyncEvent"))]
     pub fn notify_all(&self) {
-        let mut waiters = self.waiters.lock().expect("poisoned async event waiters");
-        let drained = waiters.drain(..);
-        for waiter in drained {
+        let mut pending = {
+            let mut waiters = self.waiters.lock().expect("poisoned async event waiters");
+            if waiters.is_empty() {
+                return;
+            }
+            std::mem::take(&mut *waiters)
+        };
+        // Sending can invoke an arbitrary waker synchronously. Notify this
+        // snapshot outside the lock; reentrant listeners belong to the next
+        // one.
+        for waiter in pending.drain(..) {
             let _ = waiter.send(());
         }
+        // Reuse the allocation when callbacks/concurrent threads have not
+        // registered more waiters. Never overwrite their live registrations.
+        let mut waiters = self.waiters.lock().expect("poisoned async event waiters");
+        if waiters.is_empty() && pending.capacity() > waiters.capacity() {
+            *waiters = pending;
+        }
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use std::{
+        future::Future,
+        pin::Pin,
+        sync::{
+            Arc, Mutex,
+            atomic::{AtomicBool, Ordering},
+        },
+        task::{Context, Wake, Waker},
+    };
+
     use super::AsyncEvent;
 
+    #[test]
+    fn notification_callbacks_can_register_the_next_listener() {
+        struct RegisterListener {
+            event: Arc<AsyncEvent>,
+            unlocked: AtomicBool,
+            next: Mutex<Option<futures::channel::oneshot::Receiver<()>>>,
+        }
+        impl Wake for RegisterListener {
+            fn wake(self: Arc<Self>) {
+                self.wake_by_ref();
+            }
+
+            fn wake_by_ref(self: &Arc<Self>) {
+                // Detect the old deadlock without making the test hang.
+                let unlocked = self.event.waiters.try_lock().is_ok();
+                self.unlocked.store(unlocked, Ordering::Relaxed);
+                if unlocked {
+                    *self.next.lock().unwrap() = Some(self.event.listen());
+                }
+            }
+        }
+
+        let event = Arc::new(AsyncEvent::new());
+        let callback = Arc::new(RegisterListener {
+            event: Arc::clone(&event),
+            unlocked: AtomicBool::new(false),
+            next: Mutex::new(None),
+        });
+        let waker = Waker::from(Arc::clone(&callback));
+        let mut first = event.listen();
+        assert!(
+            Pin::new(&mut first)
+                .poll(&mut Context::from_waker(&waker))
+                .is_pending()
+        );
+
+        event.notify_all();
+        assert!(callback.unlocked.load(Ordering::Relaxed));
+        assert_eq!(first.try_recv().unwrap(), Some(()));
+        let mut next = callback.next.lock().unwrap().take().unwrap();
+        assert_eq!(next.try_recv().unwrap(), None);
+        event.notify_all();
+        assert_eq!(next.try_recv().unwrap(), Some(()));
+    }
+
+    #[test]
+    fn notification_reuses_waiter_capacity_between_bursts() {
+        let event = AsyncEvent::new();
+        let mut listeners: Vec<_> = (0..16).map(|_| event.listen()).collect();
+        let capacity = event.waiters.lock().unwrap().capacity();
+        event.notify_all();
+        assert!(
+            listeners
+                .iter_mut()
+                .all(|listener| listener.try_recv().unwrap() == Some(()))
+        );
+        assert_eq!(event.waiters.lock().unwrap().capacity(), capacity);
+        event.notify_all();
+        assert_eq!(event.waiters.lock().unwrap().capacity(), capacity);
+    }
+
     #[test]
     fn notification_wakes_all_current_listeners() {
         let event = AsyncEvent::new();

src/transport/stream/tls_session.rs

@@ -178,7 +178,14 @@ pub(super) fn complete_tls_handshake(
     timeout: Duration,
     server: Option<&Weak<ServerCore>>,
 ) -> io::Result<()> {
-    let deadline = std::time::Instant::now() + timeout;
+    let deadline = std::time::Instant::now()
+        .checked_add(timeout)
+        .ok_or_else(|| {
+            io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "TLS handshake timeout is too large",
+            )
+        })?;
     loop {
         if tls_server_closed(server) {
             return Err(io::Error::new(
@@ -457,6 +464,10 @@ mod tests {
             .expect_err("zero timeout should fail before handshake I/O");
         assert_eq!(err.kind(), io::ErrorKind::TimedOut);
 
+        let err = complete_tls_handshake(&tls_state, Duration::MAX, None)
+            .expect_err("unrepresentable deadline should fail before handshake I/O");
+        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
+
         crate::initialize_python_for_tests();
         Python::attach(|py| {
             let (server, loop_core) = build_test_server(py);

src/transport/tls/mod.rs

@@ -137,26 +137,44 @@ pub fn server_tls_settings(
     })
 }
 
+#[derive(Debug, PartialEq, Eq)]
+enum TimeoutValueError {
+    NotPositiveFinite,
+    TooLarge,
+}
+
+/// Convert scalar input without Python interaction or panicking on overflow.
 #[cfg_attr(feature = "profile", hotpath::measure)]
-fn handshake_timeout(value: Option<f64>) -> PyResult<Duration> {
-    let secs = value.unwrap_or(DEFAULT_HANDSHAKE_TIMEOUT_SECS);
-    if !secs.is_finite() || secs <= 0.0 {
-        return Err(PyValueError::new_err(
-            "ssl_handshake_timeout must be a positive finite number",
-        ));
+fn timeout_duration(value: f64) -> Result<Duration, TimeoutValueError> {
+    if !value.is_finite() || value <= 0.0 {
+        return Err(TimeoutValueError::NotPositiveFinite);
     }
-    Ok(Duration::from_secs_f64(secs))
+    Duration::try_from_secs_f64(value).map_err(|_| TimeoutValueError::TooLarge)
+}
+
+#[cfg_attr(feature = "profile", hotpath::measure)]
+fn python_timeout(value: Option<f64>, default: f64, parameter: &str) -> PyResult<Duration> {
+    timeout_duration(value.unwrap_or(default)).map_err(|error| {
+        let detail = match error {
+            TimeoutValueError::NotPositiveFinite => "must be a positive finite number",
+            TimeoutValueError::TooLarge => "is too large",
+        };
+        PyValueError::new_err(format!("{parameter} {detail}"))
+    })
+}
+
+#[cfg_attr(feature = "profile", hotpath::measure)]
+fn handshake_timeout(value: Option<f64>) -> PyResult<Duration> {
+    python_timeout(
+        value,
+        DEFAULT_HANDSHAKE_TIMEOUT_SECS,
+        "ssl_handshake_timeout",
+    )
 }
 
 #[cfg_attr(feature = "profile", hotpath::measure)]
 fn shutdown_timeout(value: Option<f64>) -> PyResult<Duration> {
-    let secs = value.unwrap_or(DEFAULT_SHUTDOWN_TIMEOUT_SECS);
-    if !secs.is_finite() || secs <= 0.0 {
-        return Err(PyValueError::new_err(
-            "ssl_shutdown_timeout must be a positive finite number",
-        ));
-    }
-    Ok(Duration::from_secs_f64(secs))
+    python_timeout(value, DEFAULT_SHUTDOWN_TIMEOUT_SECS, "ssl_shutdown_timeout")
 }
 
 #[cfg_attr(feature = "profile", hotpath::measure)]
@@ -372,3 +390,60 @@ fn build_server_config(py: Python<'_>, ssl_context: &Py<PyAny>) -> PyResult<Serv
 fn default_protocol_versions() -> &'static [&'static SupportedProtocolVersion] {
     rustls::DEFAULT_VERSIONS
 }
+
+#[cfg(test)]
+mod timeout_tests {
+    use super::*;
+
+    #[test]
+    fn timeout_values_preserve_defaults_rounding_and_invalid_classification() {
+        assert_eq!(handshake_timeout(None).unwrap(), Duration::from_secs(60));
+        assert_eq!(shutdown_timeout(None).unwrap(), Duration::from_secs(30));
+        for value in [f64::MIN_POSITIVE, 0.000_000_000_1, 0.25, 1.5, 1e10] {
+            assert_eq!(
+                timeout_duration(value).unwrap(),
+                Duration::from_secs_f64(value)
+            );
+        }
+        for value in [0.0, -0.0, -1.0, f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
+            assert_eq!(
+                timeout_duration(value),
+                Err(TimeoutValueError::NotPositiveFinite)
+            );
+        }
+        assert_eq!(timeout_duration(f64::MAX), Err(TimeoutValueError::TooLarge));
+    }
+
+    #[test]
+    fn invalid_timeout_messages_keep_the_parameter_name() {
+        crate::initialize_python_for_tests();
+        Python::attach(|py| {
+            type ParseTimeout = fn(Option<f64>) -> PyResult<Duration>;
+            for (parse, name) in [
+                (handshake_timeout as ParseTimeout, "ssl_handshake_timeout"),
+                (shutdown_timeout, "ssl_shutdown_timeout"),
+            ] {
+                let error = parse(Some(-1.0)).unwrap_err();
+                assert!(error.is_instance_of::<PyValueError>(py));
+                assert_eq!(
+                    error.value(py).to_string(),
+                    format!("{name} must be a positive finite number")
+                );
+            }
+        });
+    }
+
+    #[test]
+    fn oversized_timeouts_report_python_errors_instead_of_panicking() {
+        crate::initialize_python_for_tests();
+        Python::attach(|py| {
+            for parse in [handshake_timeout, shutdown_timeout] {
+                let outcome = std::panic::catch_unwind(|| parse(Some(f64::MAX)));
+                assert!(outcome.is_ok(), "timeout conversion must not panic");
+                let error = outcome.unwrap().unwrap_err();
+                assert!(error.is_instance_of::<PyValueError>(py));
+                assert!(error.to_string().contains("too large"));
+            }
+        });
+    }
+}