Remove custom WebSocketStream::close() implementation

Instead simply send an owned Closed message. This simplifies the code
and among other things also handles errors like WouldBlock correctly
instead of handling them like a real error.
pull/3/head
Sebastian Dröge 5 years ago
parent e86a2f4ae7
commit ae0a324116
  1. 34
      src/lib.rs

@ -29,10 +29,8 @@ use std::io::{Read, Write};
use compat::{cvt, AllowStd, ContextWaker}; use compat::{cvt, AllowStd, ContextWaker};
use futures::io::{AsyncRead, AsyncWrite}; use futures::io::{AsyncRead, AsyncWrite};
use futures::{Sink, Stream}; use futures::{Sink, SinkExt, Stream};
use log::*; use log::*;
use pin_project::pin_project;
use std::future::Future;
use std::pin::Pin; use std::pin::Pin;
use std::task::{Context, Poll}; use std::task::{Context, Poll};
@ -248,11 +246,8 @@ impl<S> WebSocketStream<S> {
where where
S: AsyncRead + AsyncWrite + Unpin, S: AsyncRead + AsyncWrite + Unpin,
{ {
let f = CloseFuture { let msg = msg.map(|msg| msg.into_owned());
stream: self, self.send(Message::Close(msg)).await
message: Some(msg),
};
f.await
} }
} }
@ -324,29 +319,6 @@ where
} }
} }
#[pin_project]
struct CloseFuture<'a, T> {
stream: &'a mut WebSocketStream<T>,
message: Option<Option<CloseFrame<'a>>>,
}
impl<'a, T> Future for CloseFuture<'a, T>
where
T: AsyncRead + AsyncWrite + Unpin,
AllowStd<T>: Read + Write,
{
type Output = Result<(), WsError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
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)] #[cfg(test)]
mod tests { mod tests {
use crate::compat::AllowStd; use crate::compat::AllowStd;

Loading…
Cancel
Save