diff --git a/quic-client/src/nonblocking/quic_client.rs b/quic-client/src/nonblocking/quic_client.rs index 4cd6635563..dd5b858865 100644 --- a/quic-client/src/nonblocking/quic_client.rs +++ b/quic-client/src/nonblocking/quic_client.rs @@ -35,7 +35,7 @@ use { thread, }, thiserror::Error, - tokio::{sync::RwLock, time::timeout}, + tokio::{sync::OnceCell, time::timeout}, }; struct SkipServerVerification; @@ -67,7 +67,7 @@ pub struct QuicClientCertificate { /// A lazy-initialized Quic Endpoint pub struct QuicLazyInitializedEndpoint { - endpoint: RwLock>>, + endpoint: OnceCell>, client_certificate: Arc, client_endpoint: Option, } @@ -94,7 +94,7 @@ impl QuicLazyInitializedEndpoint { client_endpoint: Option, ) -> Self { Self { - endpoint: RwLock::new(None), + endpoint: OnceCell::>::new(), client_certificate, client_endpoint, } @@ -139,26 +139,10 @@ impl QuicLazyInitializedEndpoint { } async fn get_endpoint(&self) -> Arc { - let lock = self.endpoint.read().await; - let endpoint = lock.as_ref(); - - match endpoint { - Some(endpoint) => endpoint.clone(), - None => { - drop(lock); - let mut lock = self.endpoint.write().await; - let endpoint = lock.as_ref(); - - match endpoint { - Some(endpoint) => endpoint.clone(), - None => { - let connection = Arc::new(self.create_endpoint()); - *lock = Some(connection.clone()); - connection - } - } - } - } + self.endpoint + .get_or_init(|| async { Arc::new(self.create_endpoint()) }) + .await + .clone() } }