diff --git a/src/lib.rs b/src/lib.rs index 1ec8f09..874d9df 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -29,10 +29,8 @@ use std::io::{Read, Write}; use compat::{cvt, AllowStd, ContextWaker}; use futures::io::{AsyncRead, AsyncWrite}; -use futures::{Sink, Stream}; +use futures::{Sink, SinkExt, Stream}; use log::*; -use pin_project::pin_project; -use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; @@ -248,11 +246,8 @@ impl WebSocketStream { where S: AsyncRead + AsyncWrite + Unpin, { - let f = CloseFuture { - stream: self, - message: Some(msg), - }; - f.await + let msg = msg.map(|msg| msg.into_owned()); + self.send(Message::Close(msg)).await } } @@ -324,29 +319,6 @@ where } } -#[pin_project] -struct CloseFuture<'a, T> { - stream: &'a mut WebSocketStream, - message: Option>>, -} - -impl<'a, T> Future for CloseFuture<'a, T> -where - T: AsyncRead + AsyncWrite + Unpin, - AllowStd: Read + Write, -{ - type Output = Result<(), WsError>; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.project(); - let message = this.message.take().expect("Cannot poll twice"); - Poll::Ready( - this.stream - .with_context(Some((ContextWaker::Write, cx)), |s| s.close(message)), - ) - } -} - #[cfg(test)] mod tests { use crate::compat::AllowStd;