Lingering close for HTTP and WebSocket servers; half-close API (#21)

crates/turnloop-io/README.md

@@ -25,6 +25,23 @@ one absolute backend-clock deadline for a whole exchange.
 5. Allocate retained buffers at connection creation/warm-up, never per I/O call.
    Application-owned headers/messages and crypto state have separate accounting.
 
+## Half-close and lingering close
+
+`close` releases the whole transport. `shutdown` (the `HalfClose` trait, implemented
+by `AsyncIo`, `TlsStream` and `Transport`) flushes accepted writes and ends only the
+write direction: TLS sends close_notify, then TCP `shutdown(SHUT_WR)` (IOCP
+`SD_SEND`, WASI `shutdown(send)`). Reads continue until the peer's EOF and the handle
+stays open until `close` or drop. UDP, Windows named pipes and non-socket Unix
+descriptors return `Unsupported`.
+
+A server should not close a socket while peer bytes are unread: that sends RST, and
+on macOS and Windows a received RST discards data the peer has not read yet (the
+tail of the final response). `linger_close(stream, scratch, deadline)` flushes and
+half-closes, reads and discards into caller-retained scratch until the peer's EOF
+or the deadline, then closes. It waits only on one read and one executor timer (no
+spin) and allocates nothing per read; `Lingered` reports how it ended and how many
+bytes were discarded.
+
 Native TCP/pipes and WASI sockets use the identical generic stream code. Browser
 raw sockets/listeners return Unsupported; browser HTTP uses host fetch. Windows
 uses the production Backend implementation when integrated.

crates/turnloop-io/src/lib.rs

@@ -6,6 +6,7 @@ use std::{
     io,
     net::SocketAddr,
     pin::Pin,
+    task::{Context, Poll},
 };
 pub use turnloop::{self, AsyncIo, ExecutorHandle, Instant, backend::Backend};
 pub mod dns;
@@ -53,6 +54,233 @@ pub async fn flush<S: AsyncWrite + Unpin>(stream: &mut S) -> io::Result<()> {
 pub async fn close<S: AsyncWrite + Unpin>(stream: &mut S) -> io::Result<()> {
     poll_fn(|cx| Pin::new(&mut *stream).poll_close(cx)).await
 }
+/// A stream whose write direction can end while its read direction stays open.
+///
+/// futures-io's `poll_close` releases the whole transport. A half-close instead
+/// sends the peer EOF and keeps reading, which a server needs for a lingering
+/// close: closing a socket with unread peer bytes sends RST, and a received RST
+/// discards data the peer has not read yet on macOS and Windows. `AsyncIo`,
+/// `turnloop_tls::asynchronous::TlsStream` and `Transport` implement it.
+pub trait HalfClose: Stream {
+    /// Backend of the executor that owns the native handle.
+    type Backend: Backend;
+    /// The executor owning the native handle, whose timers bound a lingering close.
+    /// `None` once the transport has been released.
+    fn executor(&self) -> Option<ExecutorHandle<Self::Backend>>;
+    /// Flush accepted writes, then end the write direction without closing the
+    /// handle (TLS sends close_notify first). Reads continue until the peer's EOF;
+    /// later writes fail. Completed shutdowns stay complete on repeated polls.
+    /// Transports without a separate write direction return
+    /// `io::ErrorKind::Unsupported` (see `AsyncIo::poll_shutdown`).
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>>;
+}
+impl<B: Backend> HalfClose for AsyncIo<B> {
+    type Backend = B;
+    fn executor(&self) -> Option<ExecutorHandle<B>> {
+        Some(AsyncIo::executor(self))
+    }
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.get_mut().poll_shutdown(cx)
+    }
+}
+/// Half-close: flush accepted writes and end the write direction, keeping the
+/// transport open for reads. [`close`] keeps its full-close semantics.
+pub async fn shutdown<S: HalfClose>(stream: &mut S) -> io::Result<()> {
+    poll_fn(|cx| Pin::new(&mut *stream).poll_shutdown(cx)).await
+}
+/// How the discard phase of a lingering close ended.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum LingerEnd {
+    /// The peer ended its stream; all its remaining input was read and discarded.
+    Eof,
+    /// The linger deadline passed before the peer's EOF.
+    Deadline,
+    /// The transport cannot half-close, so it closed without lingering.
+    Unsupported,
+    /// The half-close, a discard read or the deadline timer failed (for example
+    /// the peer reset the connection); lingering stopped and the stream closed.
+    Failed(io::ErrorKind),
+}
+/// Outcome of a lingering close, including proof of the discarded input.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct Lingered {
+    /// Why discarding stopped.
+    pub end: LingerEnd,
+    /// Non-empty reads performed after the half-close.
+    pub reads: u64,
+    /// Peer bytes read and discarded after the half-close.
+    pub discarded: u64,
+}
+enum LingerPhase {
+    Shutdown,
+    Discard,
+    Close,
+    Done,
+}
+/// Future for [`linger_close`]. It borrows the stream and caller-retained scratch,
+/// allocates nothing, and waits only on the stream's own read and one executor
+/// timer, so an idle peer costs no turns before the deadline.
+pub struct LingerClose<'a, S: HalfClose> {
+    stream: &'a mut S,
+    scratch: &'a mut [u8],
+    executor: Option<ExecutorHandle<S::Backend>>,
+    deadline: Option<Instant>,
+    sleep: Option<(Instant, turnloop::Sleep<S::Backend>)>,
+    phase: LingerPhase,
+    lingered: Lingered,
+}
+/// Lingering close, as nginx `lingering_close`: flush and half-close, read and
+/// discard peer input until its EOF or `deadline`, then close.
+///
+/// Flush the final response before calling this. Closing with unread peer bytes
+/// (a pipelined request, HTTP/2 WINDOW_UPDATE or SETTINGS acknowledgements) sends
+/// RST, which can discard the peer's unread copy of that response. `scratch` must
+/// be non-empty; reuse a connection's retained input buffer. `None` lingers until
+/// the peer's EOF. The deadline bounds the half-close too. Transports that cannot
+/// half-close close at once ([`LingerEnd::Unsupported`]); half-close or read
+/// failures stop lingering and close ([`LingerEnd::Failed`]).
+///
+/// The result is `Err` when scratch is empty, when the half-close has not finished
+/// by the deadline (`TimedOut`: the peer stopped reading, so pending output cannot
+/// flush) or when the final `poll_close` fails. The stream is then not closed: drop
+/// it, which releases the handle. Dropping the future mid-way keeps in-flight
+/// operations owned by the stream, so a new lingering close resumes them.
+pub fn linger_close<'a, S: HalfClose>(
+    stream: &'a mut S,
+    scratch: &'a mut [u8],
+    deadline: Option<Instant>,
+) -> LingerClose<'a, S> {
+    LingerClose {
+        executor: stream.executor(),
+        stream,
+        scratch,
+        deadline,
+        sleep: None,
+        phase: LingerPhase::Shutdown,
+        lingered: Lingered {
+            end: LingerEnd::Deadline,
+            reads: 0,
+            discarded: 0,
+        },
+    }
+}
+impl<S: HalfClose> LingerClose<'_, S> {
+    /// Move the deadline, for example to honour an earlier server shutdown
+    /// deadline. The timer is re-armed on the next poll only if it changed.
+    pub fn set_deadline(&mut self, deadline: Option<Instant>) {
+        self.deadline = deadline;
+    }
+    /// Discard progress so far; `end` is meaningful once the future completes.
+    pub fn lingered(&self) -> Lingered {
+        self.lingered
+    }
+    /// Whether the deadline has passed. Otherwise its timer is armed (re-armed
+    /// only when the deadline moved) and registered with `cx`.
+    fn poll_expired(&mut self, cx: &mut Context<'_>) -> Result<bool, io::ErrorKind> {
+        let Some(deadline) = self.deadline else {
+            self.sleep = None;
+            return Ok(false);
+        };
+        if self.sleep.as_ref().is_none_or(|(at, _)| *at != deadline) {
+            let Some(executor) = &self.executor else {
+                return Ok(true);
+            };
+            if executor.now() >= deadline {
+                self.sleep = None;
+                return Ok(true);
+            }
+            // Replacing the timer drops (and closes) the previous one.
+            self.sleep = Some((deadline, executor.sleep_until(deadline)));
+        }
+        let Some((_, sleep)) = self.sleep.as_mut() else {
+            return Ok(false);
+        };
+        match Pin::new(sleep).poll(cx) {
+            Poll::Ready(Ok(())) => Ok(true),
+            Poll::Ready(Err(e)) => Err(error(e).kind()),
+            Poll::Pending => Ok(false),
+        }
+    }
+    fn poll_discard(&mut self, cx: &mut Context<'_>) -> Option<LingerEnd> {
+        loop {
+            match Pin::new(&mut *self.stream).poll_read(cx, self.scratch) {
+                Poll::Ready(Ok(0)) => return Some(LingerEnd::Eof),
+                Poll::Ready(Ok(n)) => {
+                    self.lingered.reads += 1;
+                    self.lingered.discarded += n as u64;
+                }
+                Poll::Ready(Err(e)) => return Some(LingerEnd::Failed(e.kind())),
+                Poll::Pending => break,
+            }
+        }
+        match self.poll_expired(cx) {
+            Ok(false) => None,
+            Ok(true) => Some(LingerEnd::Deadline),
+            Err(kind) => Some(LingerEnd::Failed(kind)),
+        }
+    }
+}
+impl<S: HalfClose> Future for LingerClose<'_, S> {
+    type Output = io::Result<Lingered>;
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+        loop {
+            match this.phase {
+                LingerPhase::Shutdown => {
+                    if this.scratch.is_empty() {
+                        return Poll::Ready(Err(io::Error::new(
+                            io::ErrorKind::InvalidInput,
+                            "lingering close needs non-empty scratch storage",
+                        )));
+                    }
+                    match Pin::new(&mut *this.stream).poll_shutdown(cx) {
+                        Poll::Pending => {
+                            let expired = match this.poll_expired(cx) {
+                                Ok(false) => return Poll::Pending,
+                                Ok(true) => io::ErrorKind::TimedOut,
+                                Err(kind) => kind,
+                            };
+                            this.lingered.end = if expired == io::ErrorKind::TimedOut {
+                                LingerEnd::Deadline
+                            } else {
+                                LingerEnd::Failed(expired)
+                            };
+                            this.sleep = None;
+                            this.phase = LingerPhase::Done;
+                            return Poll::Ready(Err(io::Error::new(
+                                expired,
+                                "half-close did not complete before the linger deadline",
+                            )));
+                        }
+                        Poll::Ready(Ok(())) => this.phase = LingerPhase::Discard,
+                        Poll::Ready(Err(e)) => {
+                            this.lingered.end = if e.kind() == io::ErrorKind::Unsupported {
+                                LingerEnd::Unsupported
+                            } else {
+                                LingerEnd::Failed(e.kind())
+                            };
+                            this.phase = LingerPhase::Close;
+                        }
+                    }
+                }
+                LingerPhase::Discard => {
+                    let Some(end) = this.poll_discard(cx) else {
+                        return Poll::Pending;
+                    };
+                    this.lingered.end = end;
+                    this.sleep = None;
+                    this.phase = LingerPhase::Close;
+                }
+                LingerPhase::Close => {
+                    let result = std::task::ready!(Pin::new(&mut *this.stream).poll_close(cx));
+                    this.phase = LingerPhase::Done;
+                    return Poll::Ready(result.map(|()| this.lingered));
+                }
+                LingerPhase::Done => return Poll::Ready(Ok(this.lingered)),
+            }
+        }
+    }
+}
 /// One deadline shared across every phase of an exchange; uses the executor timer.
 pub async fn deadline<B: Backend, F: Future<Output = io::Result<T>>, T>(
     executor: &ExecutorHandle<B>,