Impl own pin project

DataDog · Apr 13, 2022 · 623ed01 · 623ed01
1 parent 6924d1c
commit 623ed01
Showing 1 changed file with 45 additions and 29 deletions.
diff --git a/ddprof-exporter/src/connector/conn_stream.rs b/ddprof-exporter/src/connector/conn_stream.rs
@@ -8,37 +8,53 @@ use std::{
 
 use futures::{future, Future, FutureExt, TryFutureExt};
 use hyper_rustls::HttpsConnector;
-use pin_project_lite::pin_project;
 
-#[cfg(unix)]
-pin_project! {
-    #[derive(Debug)]
-    #[project = ConnStreamProj]
-    pub enum ConnStream {
-        Tcp {
-            #[pin] transport: tokio::net::TcpStream,
-        },
-        Tls {
-            #[pin] transport: tokio_rustls::client::TlsStream<tokio::net::TcpStream>,
-        },
-        // Tokio doesn't handle unix sockets on windows
-        Udp {
-            #[pin] transport: tokio::net::UnixStream,
-        },
-    }
+#[derive(Debug)]
+pub enum ConnStream {
+    Tcp {
+        transport: tokio::net::TcpStream,
+    },
+    Tls {
+        transport: Box<tokio_rustls::client::TlsStream<tokio::net::TcpStream>>,
+    },
+    #[cfg(unix)]
+    Udp {
+        transport: tokio::net::UnixStream,
+    },
+}
+
+pub enum ConnStreamProj<'pin>
+where
+    ConnStream: 'pin,
+{
+    Tcp {
+        transport: Pin<&'pin mut tokio::net::TcpStream>,
+    },
+    Tls {
+        transport: Pin<&'pin mut tokio_rustls::client::TlsStream<tokio::net::TcpStream>>,
+    },
+    #[cfg(unix)]
+    Udp {
+        transport: Pin<&'pin mut tokio::net::UnixStream>,
+    },
 }
 
-#[cfg(not(unix))]
-pin_project! {
-    #[derive(Debug)]
-    #[project = ConnStreamProj]
-    pub enum ConnStream {
-        Tcp {
-            #[pin] transport: tokio::net::TcpStream,
-        },
-        Tls {
-            #[pin] transport: tokio_rustls::client::TlsStream<tokio::net::TcpStream>,
-        },
+impl ConnStream {
+    pub(crate) fn project<'__pin>(self: Pin<&'__pin mut Self>) -> ConnStreamProj<'__pin> {
+        unsafe {
+            match self.get_unchecked_mut() {
+                Self::Tcp { transport } => ConnStreamProj::Tcp {
+                    transport: Pin::new_unchecked(transport),
+                },
+                Self::Tls { transport } => ConnStreamProj::Tls {
+                    transport: Pin::new_unchecked(transport),
+                },
+                #[cfg(unix)]
+                Self::Udp { transport } => ConnStreamProj::Udp {
+                    transport: Pin::new_unchecked(transport),
+                },
+            }
+        }
     }
 }
 
@@ -87,7 +103,7 @@ impl ConnStream {
                 }
             }
             hyper_rustls::MaybeHttpsStream::Https(t) => {
-                future::ready(Ok(ConnStream::Tls { transport: t }))
+                future::ready(Ok(ConnStream::Tls { transport: Box::from(t) }))
             }
         })
     }