fix: bound native shutdown and close connections that become idle during drain

transport/native/shutdown_test.go

@@ -0,0 +1,101 @@
+package native
+
+import (
+	"bufio"
+	"context"
+	"io"
+	"net"
+	"net/http"
+	"sync"
+	"testing"
+	"time"
+)
+
+type pausedWriteListener struct {
+	net.Listener
+	writing chan struct{}
+	release chan struct{}
+}
+
+func (l *pausedWriteListener) Accept() (net.Conn, error) {
+	conn, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	return &pausedWriteConn{Conn: conn, writing: l.writing, release: l.release}, nil
+}
+
+type pausedWriteConn struct {
+	net.Conn
+	writing chan struct{}
+	release chan struct{}
+	once    sync.Once
+}
+
+func (c *pausedWriteConn) Write(data []byte) (int, error) {
+	c.once.Do(func() { close(c.writing); <-c.release })
+	return c.Conn.Write(data)
+}
+
+func TestShutdownDuringResponseWriteClosesNewlyIdleConnection(t *testing.T) {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	wrapped := &pausedWriteListener{Listener: listener, writing: make(chan struct{}), release: make(chan struct{})}
+	var once sync.Once
+	unblock := func() { once.Do(func() { close(wrapped.release) }) }
+	defer unblock()
+	defer listener.Close()
+	server := New(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = io.WriteString(w, "finished") }))
+	server.ShutdownTimeout = 200 * time.Millisecond
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	done := make(chan error, 1)
+	go func() { done <- server.Serve(ctx, wrapped) }()
+	client, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer client.Close()
+	client.SetDeadline(time.Now().Add(2 * time.Second))
+	if _, err := io.WriteString(client, "GET / HTTP/1.1\r\nHost: localhost\r\n\r\n"); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-wrapped.writing:
+	case <-time.After(time.Second):
+		t.Fatal("response did not reach write")
+	}
+	// The response has already chosen its Keep-Alive header. Shutdown must
+	// still close this connection when its write completes and it becomes idle.
+	cancel()
+	deadline := time.Now().Add(time.Second)
+	for !server.isClosing() {
+		if time.Now().After(deadline) {
+			t.Fatal("shutdown did not start")
+		}
+		time.Sleep(time.Millisecond)
+	}
+	unblock()
+	response, err := http.ReadResponse(bufio.NewReader(client), &http.Request{Method: "GET"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	body, err := io.ReadAll(response.Body)
+	response.Body.Close()
+	if err != nil || string(body) != "finished" {
+		t.Fatalf("response %q: %v", body, err)
+	}
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("completed request left a connection stuck in drain: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("shutdown did not complete")
+	}
+	if stats := server.Stats(); stats.ActiveConnections != 0 || stats.ActiveRequests != 0 {
+		t.Fatalf("resources retained: %+v", stats)
+	}
+}