net/http: reduce memory usage when hijacking

src/net/http/serve_test.go

@@ -6151,6 +6151,49 @@ func testServerHijackGetsBackgroundByte(t *testing.T, mode testMode) {
 	<-done
 }
 
+// Test that the bufio.Reader returned by Hijack yields the entire body.
+func TestServerHijackGetsFullBody(t *testing.T) {
+	run(t, testServerHijackGetsFullBody, []testMode{http1Mode})
+}
+func testServerHijackGetsFullBody(t *testing.T, mode testMode) {
+	if runtime.GOOS == "plan9" {
+		t.Skip("skipping test; see https://golang.org/issue/18657")
+	}
+	done := make(chan struct{})
+	needle := strings.Repeat("x", 4096)
+	ts := newClientServerTest(t, mode, HandlerFunc(func(w ResponseWriter, r *Request) {
+		defer close(done)
+
+		conn, buf, err := w.(Hijacker).Hijack()
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		defer conn.Close()
+
+		peek, err := buf.Reader.Peek(4096)
+		if string(peek) != needle || err != nil {
+			t.Errorf("Peek = %q, %v; want foo, nil", peek, err)
+		}
+	})).ts
+
+	cn, err := net.Dial("tcp", ts.Listener.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cn.Close()
+	buf := []byte("GET / HTTP/1.1\r\nHost: e.com\r\n\r\n")
+	buf = append(buf, []byte(needle)...)
+	if _, err := cn.Write(buf); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := cn.(*net.TCPConn).CloseWrite(); err != nil {
+		t.Fatal(err)
+	}
+	<-done
+}
+
 // Like TestServerHijackGetsBackgroundByte above but sending a
 // immediate 1MB of data to the server to fill up the server's 4KB
 // buffer.
@@ -6205,6 +6248,120 @@ func testServerHijackGetsBackgroundByte_big(t *testing.T, mode testMode) {
 	<-done
 }
 
+func BenchmarkServerHijackMemoryUsage(b *testing.B) {
+	echo := []byte("echo 01234")
+
+	// Client mode
+	if url := os.Getenv("TEST_BENCH_SERVER_URL"); url != "" {
+		n, err := strconv.Atoi(os.Getenv("TEST_BENCH_CLIENT_N"))
+		if err != nil {
+			panic(err)
+		}
+		clientBodies := make([]io.ReadWriteCloser, n)
+		for i := 0; i < n; i++ {
+			res, err := Get(url)
+			if err != nil {
+				log.Panicf("Get: %v", err)
+			}
+			clientBodies[i] = res.Body.(io.ReadWriteCloser)
+		}
+		for _, rwc := range clientBodies {
+			if _, err := rwc.Write(echo); err != nil {
+				log.Panicf("write: %v", err)
+			}
+		}
+		readBuf := bytes.Buffer{}
+		readBuf.Grow(bytes.MinRead)
+		for _, rwc := range clientBodies {
+			readBuf.Reset()
+			n, err := readBuf.ReadFrom(rwc)
+			if err != nil || n != int64(len(echo)) {
+				log.Panicf("read: want=4, got=%d, err=%v", n, err)
+			}
+			reply := readBuf.Bytes()
+			if !bytes.Equal(reply, echo) {
+				log.Panicf("echo: want=%q, got=%q", string(echo), string(reply))
+			}
+			rwc.Close()
+		}
+		os.Exit(0)
+		return
+	}
+
+	// Server mode
+	type waitingConn struct {
+		rwc   net.Conn
+		rwBuf *bufio.ReadWriter
+	}
+	waitingConns := make(chan waitingConn, b.N)
+	handlerDone := make(chan struct{})
+	var before runtime.MemStats
+	var after runtime.MemStats
+
+	cst := newClientServerTest(b, http1Mode, HandlerFunc(func(w ResponseWriter, r *Request) {
+		w.Header().Set("Connection", "Upgrade")
+		w.Header().Set("Upgrade", "someProto")
+		w.WriteHeader(StatusSwitchingProtocols)
+		rwc, rwBuf, err := w.(Hijacker).Hijack()
+		if err != nil {
+			b.Error(err)
+		}
+		waitingConns <- waitingConn{rwc, rwBuf}
+		handlerDone <- struct{}{}
+		// exit from handler, hold on to conn+buffer
+	}))
+
+	// Start client
+	cmd := testenv.Command(b, os.Args[0], "-test.run=^$", "-test.bench=^BenchmarkServerHijackMemoryUsage$")
+	cmd.Env = append([]string{
+		fmt.Sprintf("TEST_BENCH_CLIENT_N=%d", b.N),
+		fmt.Sprintf("TEST_BENCH_SERVER_URL=%s", cst.ts.URL),
+	}, os.Environ()...)
+	var outputBuf bytes.Buffer
+	outputBuf.Grow(4096)
+	cmd.Stdout = &outputBuf
+	cmd.Stderr = &outputBuf
+	if err := cmd.Start(); err != nil {
+		b.Fatal(err)
+	}
+
+	runtime.ReadMemStats(&before)
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	// accumulate connections
+	for i := 0; i < b.N; i++ {
+		<-handlerDone
+	}
+
+	runtime.GC()
+	runtime.ReadMemStats(&after)
+	b.ReportMetric(
+		float64(after.Alloc-before.Alloc)/float64(b.N), "retained_B/op",
+	)
+
+	close(waitingConns)
+	close(handlerDone)
+
+	// verify w+r
+	for wc := range waitingConns {
+		got, err := wc.rwBuf.Peek(len(echo))
+		if err != nil {
+			b.Fatal(err)
+		} else if !bytes.Equal(got, echo) {
+			b.Fatalf("echo=%q got=%q", string(echo), string(got))
+		}
+		wc.rwBuf.Write(got)
+		wc.rwBuf.Flush()
+		wc.rwc.Close()
+	}
+
+	// Check client result
+	if err := cmd.Wait(); err != nil {
+		b.Fatalf("client: %v, with output: %s", err, outputBuf.String())
+	}
+}
+
 // Issue 18319: test that the Server validates the request method.
 func TestServerValidatesMethod(t *testing.T) {
 	tests := []struct {

src/net/http/server.go

@@ -31,6 +31,7 @@ import (
 	"sync/atomic"
 	"time"
 	_ "unsafe" // for linkname
+	"weak"
 
 	"golang.org/x/net/http/httpguts"
 )
@@ -261,7 +262,7 @@ type conn struct {
 
 	// rwc is the underlying network connection.
 	// This is never wrapped by other types and is the value given out
-	// to CloseNotifier callers. It is usually of type *net.TCPConn or
+	// to [Hijacker] callers. It is usually of type *net.TCPConn or
 	// *tls.Conn.
 	rwc net.Conn
 
@@ -314,7 +315,7 @@ func (c *conn) hijacked() bool {
 }
 
 // c.mu must be held.
-func (c *conn) hijackLocked() (rwc net.Conn, buf *bufio.ReadWriter, err error) {
+func (c *conn) hijackLocked(w *response) (rwc net.Conn, buf *bufio.ReadWriter, err error) {
 	if c.hijackedv {
 		return nil, nil, ErrHijacked
 	}
@@ -324,12 +325,17 @@ func (c *conn) hijackLocked() (rwc net.Conn, buf *bufio.ReadWriter, err error) {
 	rwc = c.rwc
 	rwc.SetDeadline(time.Time{})
 
-	buf = bufio.NewReadWriter(c.bufr, bufio.NewWriter(rwc))
 	if c.r.hasByte {
 		if _, err := c.bufr.Peek(c.bufr.Buffered() + 1); err != nil {
 			return nil, nil, fmt.Errorf("unexpected Peek failure reading buffered byte: %v", err)
 		}
 	}
+
+	ccn := &connCloseNotify{rwc: rwc, weakResponse: weak.Make(w)}
+	ccn.attach(c.bufr)
+	c.bufw.Reset(rwc)
+	buf = bufio.NewReadWriter(c.bufr, c.bufw)
+
 	c.setState(rwc, StateHijacked, runHooks)
 	return
 }
@@ -473,6 +479,9 @@ type response struct {
 	// input from it.
 	requestBodyLimitHit bool
 
+	// closeNotifyTriggered tracks prior closeNotify calls.
+	closeNotifyTriggered bool
+
 	// trailers are the headers to be sent after the handler
 	// finishes writing the body. This field is initialized from
 	// the Trailer response header when the response header is
@@ -486,11 +495,10 @@ type response struct {
 	clenBuf   [10]byte
 	statusBuf [3]byte
 
+	// lazyCloseNotifyMu protects closeNotifyCh and closeNotifyTriggered.
+	lazyCloseNotifyMu sync.Mutex
 	// closeNotifyCh is the channel returned by CloseNotify.
-	// TODO(bradfitz): this is currently (for Go 1.8) always
-	// non-nil. Make this lazily-created again as it used to be?
-	closeNotifyCh  chan bool
-	didCloseNotify atomic.Bool // atomic (only false->true winner should send)
+	closeNotifyCh chan bool
 }
 
 func (c *response) SetReadDeadline(deadline time.Time) error {
@@ -646,6 +654,50 @@ type readResult struct {
 	b   byte // byte read, if n == 1
 }
 
+// connCloseNotify is a minimal version of [connReader] for use by [Hijacker]
+// to notify any [CloseNotifier] on read errors before the handler exists. It
+// allows the hijacked [conn], [response] and [Request] to get garbage
+// collected after the handler exits.
+type connCloseNotify struct {
+	// rwc is the underlying network connection.
+	rwc net.Conn
+	// weakResponse points at the [response] until the handler exits.
+	weakResponse weak.Pointer[response]
+	// previousBufferContents holds on to previously buffered bytes of the read buffer.
+	previousBufferContents []byte
+}
+
+func (c *connCloseNotify) Read(p []byte) (n int, err error) {
+	if c.previousBufferContents != nil {
+		n = copy(p, c.previousBufferContents)
+		c.previousBufferContents = nil
+		return
+	}
+	n, err = c.rwc.Read(p)
+	if err != nil {
+		if w := c.weakResponse.Value(); w != nil {
+			w.cancelCtx()
+			w.closeNotify()
+		}
+	}
+	return
+}
+
+// attach preserves the buffered bytes while attaching to the given reader.
+func (c *connCloseNotify) attach(r *bufio.Reader) (err error) {
+	n := r.Buffered()
+	if n != 0 {
+		if c.previousBufferContents, err = r.Peek(n); err != nil {
+			return
+		}
+	}
+	r.Reset(c)
+	if n != 0 {
+		_, err = r.Peek(n)
+	}
+	return
+}
+
 // connReader is the io.Reader wrapper used by *conn. It combines a
 // selectively-activated io.LimitedReader (to bound request header
 // read sizes) with support for selectively keeping an io.Reader.Read
@@ -762,8 +814,8 @@ func (cr *connReader) handleReadError(_ error) {
 // may be called from multiple goroutines.
 func (cr *connReader) closeNotify() {
 	res := cr.conn.curReq.Load()
-	if res != nil && !res.didCloseNotify.Swap(true) {
-		res.closeNotifyCh <- true
+	if res != nil {
+		res.closeNotify()
 	}
 }
 
@@ -1100,7 +1152,6 @@ func (c *conn) readRequest(ctx context.Context) (w *response, err error) {
 		reqBody:       req.Body,
 		handlerHeader: make(Header),
 		contentLength: -1,
-		closeNotifyCh: make(chan bool, 1),
 
 		// We populate these ahead of time so we're not
 		// reading from req.Header after their Handler starts
@@ -2241,7 +2292,7 @@ func (w *response) Hijack() (rwc net.Conn, buf *bufio.ReadWriter, err error) {
 
 	// Release the bufioWriter that writes to the chunk writer, it is not
 	// used after a connection has been hijacked.
-	rwc, buf, err = c.hijackLocked()
+	rwc, buf, err = c.hijackLocked(w)
 	if err == nil {
 		putBufioWriter(w.w)
 		w.w = nil
@@ -2250,12 +2301,32 @@ func (w *response) Hijack() (rwc net.Conn, buf *bufio.ReadWriter, err error) {
 }
 
 func (w *response) CloseNotify() <-chan bool {
+	w.lazyCloseNotifyMu.Lock()
+	defer w.lazyCloseNotifyMu.Unlock()
 	if w.handlerDone.Load() {
 		panic("net/http: CloseNotify called after ServeHTTP finished")
 	}
+	if w.closeNotifyCh == nil {
+		w.closeNotifyCh = make(chan bool, 1)
+		if w.closeNotifyTriggered {
+			w.closeNotifyCh <- true // action prior closeNotify call
+		}
+	}
 	return w.closeNotifyCh
 }
 
+func (w *response) closeNotify() {
+	w.lazyCloseNotifyMu.Lock()
+	defer w.lazyCloseNotifyMu.Unlock()
+	if w.closeNotifyTriggered {
+		return // already triggered
+	}
+	w.closeNotifyTriggered = true
+	if w.closeNotifyCh != nil {
+		w.closeNotifyCh <- true
+	}
+}
+
 func registerOnHitEOF(rc io.ReadCloser, fn func()) {
 	switch v := rc.(type) {
 	case *expectContinueReader: