feat: cancel native requests on disconnect

transport/native/server.go

@@ -334,7 +334,7 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 		state.active.Store(true)
 		writer := s.writerPool.Get().(*responseWriter)
 		writer.reset(request.Method)
-		writer.configure(request, state.conn, closeAfter, s.WriteTimeout)
+		writer.configure(request, state.conn, closeAfter, s.WriteTimeout, cancel)
 		requestSlot := false
 		if slots := s.requestSlots; slots != nil {
 			select {
@@ -363,7 +363,9 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
 			}
 		}
 		handler := s.handler
+		disconnectWatch := newClientDisconnectWatch(reader, state.conn, cancel)
 		handler.ServeHTTP(writer, request)
+		disconnectWatch.stop()
 		cancel()
 
 		closing := s.isClosing() || state.closeWhenIdle.Load()
@@ -711,6 +713,48 @@ func isTimeout(err error) bool {
 	return errors.As(err, &netErr) && netErr.Timeout()
 }
 
+// Most handlers complete before this delay. A delayed, blocking read avoids a
+// watcher goroutine for those requests and does not poll idle client sockets.
+const clientDisconnectWatchDelay = 5 * time.Millisecond
+
+type clientDisconnectWatch struct {
+	conn    net.Conn
+	done    chan struct{}
+	stopped atomic.Bool
+	timer   *time.Timer
+}
+
+func newClientDisconnectWatch(reader *bufio.Reader, conn net.Conn, cancel context.CancelFunc) *clientDisconnectWatch {
+	watch := &clientDisconnectWatch{conn: conn, done: make(chan struct{})}
+	watch.timer = time.AfterFunc(clientDisconnectWatchDelay, func() {
+		defer close(watch.done)
+		if watch.stopped.Load() {
+			return
+		}
+		// Only this goroutine may read until stop joins it. Peek preserves
+		// pipelined bytes for the next request and bounds read-ahead to the
+		// existing reader buffer. Once data is available, stop watching:
+		// draining it to detect EOF would require an unbounded pipeline queue.
+		_, err := reader.Peek(1)
+		if err != nil && !(watch.stopped.Load() && isTimeout(err)) {
+			cancel()
+		}
+	})
+	return watch
+}
+
+func (w *clientDisconnectWatch) stop() {
+	if w.timer.Stop() {
+		return
+	}
+	w.stopped.Store(true)
+	// The watcher never changes deadlines: this interrupt cannot be
+	// overwritten by a callback starting concurrently with stop.
+	_ = w.conn.SetReadDeadline(time.Now())
+	<-w.done
+	_ = w.conn.SetReadDeadline(time.Time{})
+}
+
 func writeSimple(conn net.Conn, status int, body string, closeAfter bool, timeout time.Duration) error {
 	writer := newResponseWriter("")
 	writer.status = status
@@ -735,6 +779,7 @@ type responseWriter struct {
 	chunked        bool
 	streamingClose bool
 	writeErr       error
+	cancel         context.CancelFunc
 }
 
 func newResponseWriter(method string) *responseWriter {
@@ -762,9 +807,11 @@ func (w *responseWriter) reset(method string) {
 	w.chunked = false
 	w.streamingClose = false
 	w.writeErr = nil
+	w.cancel = nil
 }
 
-func (w *responseWriter) configure(request *http.Request, conn net.Conn, closeAfter bool, timeout time.Duration) {
+func (w *responseWriter) configure(request *http.Request, conn net.Conn, closeAfter bool, timeout time.Duration, cancel context.CancelFunc) {
+	w.cancel = cancel
 	w.protoMajor = request.ProtoMajor
 	w.protoMinor = request.ProtoMinor
 	w.conn = conn
@@ -807,9 +854,7 @@ func (w *responseWriter) Flush() {
 		w.status == http.StatusResetContent || w.status == http.StatusNotModified || (w.status >= 100 && w.status < 200) {
 		return
 	}
-	if err := w.beginStream(); err != nil {
-		w.writeErr = err
-	}
+	_ = w.recordWriteError(w.beginStream())
 }
 
 func (w *responseWriter) beginStream() error {
@@ -881,15 +926,25 @@ func (w *responseWriter) writeBytes(data []byte) error {
 	if w.writeTimeout > 0 {
 		_ = w.conn.SetWriteDeadline(time.Now().Add(w.writeTimeout))
 	}
-	return writeAll(w.conn, data)
+	return w.recordWriteError(writeAll(w.conn, data))
+}
+
+func (w *responseWriter) recordWriteError(err error) error {
+	if err != nil {
+		w.writeErr = err
+		if w.cancel != nil {
+			w.cancel()
+		}
+	}
+	return err
 }
 
 func (w *responseWriter) writeTo(conn net.Conn, closeAfter bool, timeout time.Duration) error {
+	if w.writeErr != nil {
+		_ = conn.SetWriteDeadline(time.Time{})
+		return w.writeErr
+	}
 	if w.streaming {
-		if w.writeErr != nil {
-			_ = conn.SetWriteDeadline(time.Time{})
-			return w.writeErr
-		}
 		if w.chunked {
 			if err := w.writeBytes([]byte("0\r\n\r\n")); err != nil {
 				w.writeErr = err
@@ -979,7 +1034,7 @@ func (w *responseWriter) writeHeaders(conn net.Conn) error {
 		}
 	}
 	output.WriteString("\r\n")
-	return writeAll(conn, output.Bytes())
+	return w.recordWriteError(writeAll(conn, output.Bytes()))
 }
 
 func writeAll(conn net.Conn, data []byte) error {