From 91bc3a49bea5d582a2e6cc31ed22f092da1b95aa Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Sun, 17 Dec 2023 13:50:27 -0800
Subject: [PATCH] Switch to AtomicWaker to wake the endpoint driver

Groundwork for direct connection -> driver wake-ups, and yields one
less piece of state inside a monolithic mutex.
---
 quinn/Cargo.toml      |  1 +
 quinn/src/endpoint.rs | 14 +++++---------
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/quinn/Cargo.toml b/quinn/Cargo.toml
index b8bfb2073..b171f1266 100644
--- a/quinn/Cargo.toml
+++ b/quinn/Cargo.toml
@@ -35,6 +35,7 @@ maintenance = { status = "experimental" }
 [dependencies]
 async-io = { version = "2.0", optional = true }
 async-std = { version = "1.11", optional = true }
+atomic-waker = "1.1.2"
 bytes = "1"
 # Enables futures::io::{AsyncRead, AsyncWrite} support for streams
 futures-io = { version = "0.3.19", optional = true }
diff --git a/quinn/src/endpoint.rs b/quinn/src/endpoint.rs
index 17abdc7f2..f4ffca545 100644
--- a/quinn/src/endpoint.rs
+++ b/quinn/src/endpoint.rs
@@ -8,11 +8,12 @@ use std::{
     pin::Pin,
     str,
     sync::{Arc, Mutex},
-    task::{Context, Poll, Waker},
+    task::{Context, Poll},
     time::Instant,
 };
 
 use crate::runtime::{default_runtime, AsyncUdpSocket, Runtime};
+use atomic_waker::AtomicWaker;
 use bytes::{Bytes, BytesMut};
 use pin_project_lite::pin_project;
 use proto::{
@@ -319,10 +320,8 @@ impl Future for EndpointDriver {
 
     #[allow(unused_mut)] // MSRV
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
+        self.0.shared.driver.register(cx.waker());
         let mut endpoint = self.0.state.lock().unwrap();
-        if endpoint.driver.is_none() {
-            endpoint.driver = Some(cx.waker().clone());
-        }
 
         let now = Instant::now();
         let mut keep_going = false;
@@ -372,7 +371,6 @@ pub(crate) struct State {
     inner: proto::Endpoint,
     outgoing: VecDeque<udp::Transmit>,
     incoming: VecDeque<Connecting>,
-    driver: Option<Waker>,
     ipv6: bool,
     connections: ConnectionSet,
     events: mpsc::UnboundedReceiver<(ConnectionHandle, EndpointEvent)>,
@@ -391,6 +389,7 @@ pub(crate) struct State {
 pub(crate) struct Shared {
     incoming: Notify,
     idle: Notify,
+    driver: Arc<AtomicWaker>,
 }
 
 impl State {
@@ -705,7 +704,6 @@ impl EndpointRef {
                 events,
                 outgoing: VecDeque::new(),
                 incoming: VecDeque::new(),
-                driver: None,
                 connections: ConnectionSet {
                     senders: FxHashMap::default(),
                     sender,
@@ -738,9 +736,7 @@ impl Drop for EndpointRef {
             if x == 0 {
                 // If the driver is about to be on its own, ensure it can shut down if the last
                 // connection is gone.
-                if let Some(task) = endpoint.driver.take() {
-                    task.wake();
-                }
+                self.0.shared.driver.wake();
             }
         }
     }