Enable 0rtt (#25055)

* Enable 0rtt connections

client/src/connection_cache.rs

@@ -63,6 +63,18 @@ impl ConnectionCacheStats {
             client_stats.connection_reuse.load(Ordering::Relaxed),
             Ordering::Relaxed,
         );
+        self.total_client_stats.connection_errors.fetch_add(
+            client_stats.connection_errors.load(Ordering::Relaxed),
+            Ordering::Relaxed,
+        );
+        self.total_client_stats.zero_rtt_accepts.fetch_add(
+            client_stats.zero_rtt_accepts.load(Ordering::Relaxed),
+            Ordering::Relaxed,
+        );
+        self.total_client_stats.zero_rtt_rejects.fetch_add(
+            client_stats.zero_rtt_rejects.load(Ordering::Relaxed),
+            Ordering::Relaxed,
+        );
         self.sent_packets
             .fetch_add(num_packets as u64, Ordering::Relaxed);
         self.total_batches.fetch_add(1, Ordering::Relaxed);
@@ -130,6 +142,27 @@ impl ConnectionCacheStats {
                     .swap(0, Ordering::Relaxed),
                 i64
             ),
+            (
+                "connection_errors",
+                self.total_client_stats
+                    .connection_errors
+                    .swap(0, Ordering::Relaxed),
+                i64
+            ),
+            (
+                "zero_rtt_accepts",
+                self.total_client_stats
+                    .zero_rtt_accepts
+                    .swap(0, Ordering::Relaxed),
+                i64
+            ),
+            (
+                "zero_rtt_rejects",
+                self.total_client_stats
+                    .zero_rtt_rejects
+                    .swap(0, Ordering::Relaxed),
+                i64
+            ),
             (
                 "congestion_events",
                 self.total_client_stats.congestion_events.load_and_reset(),

client/src/quic_client.rs

@@ -159,10 +159,11 @@ impl QuicClient {
     pub fn new(client_socket: UdpSocket, addr: SocketAddr) -> Self {
         let _guard = RUNTIME.enter();
 
-        let crypto = rustls::ClientConfig::builder()
+        let mut crypto = rustls::ClientConfig::builder()
             .with_safe_defaults()
             .with_custom_certificate_verifier(SkipServerVerification::new())
             .with_no_client_auth();
+        crypto.enable_early_data = true;
 
         let create_endpoint = QuicClient::create_endpoint(EndpointConfig::default(), client_socket);
 
@@ -217,6 +218,33 @@ impl QuicClient {
         Ok(Arc::new(connection))
     }
 
+    // Attempts to make a faster connection by taking advantage of pre-existing key material.
+    // Only works if connection to this endpoint was previously established.
+    async fn make_connection_0rtt(
+        &self,
+        stats: &ClientStats,
+    ) -> Result<Arc<NewConnection>, WriteError> {
+        let connecting = self.endpoint.connect(self.addr, "connect").unwrap();
+        stats.total_connections.fetch_add(1, Ordering::Relaxed);
+        let connection = match connecting.into_0rtt() {
+            Ok((connection, zero_rtt)) => {
+                if zero_rtt.await {
+                    stats.zero_rtt_accepts.fetch_add(1, Ordering::Relaxed);
+                } else {
+                    stats.zero_rtt_rejects.fetch_add(1, Ordering::Relaxed);
+                }
+                connection
+            }
+            Err(connecting) => {
+                stats.connection_errors.fetch_add(1, Ordering::Relaxed);
+                let connecting = connecting.await;
+                connecting?
+            }
+        };
+
+        Ok(Arc::new(connection))
+    }
+
     // Attempts to send data, connecting/reconnecting as necessary
     // On success, returns the connection used to successfully send the data
     async fn _send_buffer(
@@ -244,7 +272,7 @@ impl QuicClient {
             Ok(()) => Ok(connection),
             _ => {
                 let connection = {
-                    let connection = self.make_connection(stats).await?;
+                    let connection = self.make_connection_0rtt(stats).await?;
                     let mut conn_guard = self.connection.lock().await;
                     *conn_guard = Some(connection.clone());
                     connection

client/src/tpu_connection.rs

@@ -13,6 +13,8 @@ pub struct ClientStats {
     pub total_connections: AtomicU64,
     pub connection_reuse: AtomicU64,
     pub connection_errors: AtomicU64,
+    pub zero_rtt_accepts: AtomicU64,
+    pub zero_rtt_rejects: AtomicU64,
 
     // these will be the last values of these stats
     pub congestion_events: MovingStat,