102 lines
2.7 KiB
Go
102 lines
2.7 KiB
Go
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)
|
|
}
|
|
}
|