use self.drop.take()

Author: ssrlive

src/lib.rs

@@ -164,7 +164,7 @@ impl Drop for UdpStream {
             handler.abort()
         }
 
-        if let Some(drop) = &self.drop {
+        if let Some(drop) = self.drop.take() {
             let _ = drop.try_send(self.peer_addr);
         };
     }
@@ -267,11 +267,11 @@ impl AsyncRead for UdpStream {
 }
 
 impl AsyncWrite for UdpStream {
-    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<std::io::Result<usize>> {
+    fn poll_write(mut self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<std::io::Result<usize>> {
         match self.socket.poll_send_to(cx, buf, self.peer_addr) {
             Poll::Ready(Ok(r)) => Poll::Ready(Ok(r)),
             Poll::Ready(Err(e)) => {
-                if let Some(drop) = &self.drop {
+                if let Some(drop) = self.drop.take() {
                     let _ = drop.try_send(self.peer_addr);
                 };
                 Poll::Ready(Err(e))
@@ -282,14 +282,15 @@ impl AsyncWrite for UdpStream {
     fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<std::io::Result<()>> {
         Poll::Ready(Ok(()))
     }
-    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<std::io::Result<()>> {
-        let Some(drop) = &self.drop else {
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<std::io::Result<()>> {
+        let Some(drop) = self.drop.take() else {
             return Poll::Ready(Ok(()));
         };
-        match Box::pin(drop.send(self.peer_addr)).as_mut().poll(cx) {
+        let val = match Box::pin(drop.send(self.peer_addr)).as_mut().poll(cx) {
             Poll::Ready(Ok(_)) => Poll::Ready(Ok(())),
             Poll::Ready(Err(e)) => Poll::Ready(Err(std::io::Error::other(e))),
             Poll::Pending => Poll::Pending,
-        }
+        };
+        val
     }
 }