空闲超时回收边缘连接

package main

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"runtime"
	"strconv"
	"sync/atomic"
	"time"
)

type connectionKey struct{}

func check(ok bool, message string) {
	if !ok {
		panic(message)
	}
}

func main() {
	const idle = 200 * time.Millisecond
	const headers = 500 * time.Millisecond
	const work = 350 * time.Millisecond
	var serial, handled atomic.Int64
	closed := make(chan struct{}, 8)
	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		handled.Add(1)
		if r.URL.Path == "/slow" {
			time.Sleep(work)
		}
		w.Header().Set("X-Connection-ID", strconv.FormatInt(r.Context().Value(connectionKey{}).(int64), 10))
		w.Header().Set("Content-Type", "text/plain")
		fmt.Fprintln(w, "ok")
	}))
	server.Config.IdleTimeout = idle
	server.Config.ReadHeaderTimeout = headers
	server.Config.ConnContext = func(ctx context.Context, _ net.Conn) context.Context {
		return context.WithValue(ctx, connectionKey{}, serial.Add(1))
	}
	server.Config.ConnState = func(_ net.Conn, state http.ConnState) {
		if state == http.StateClosed {
			closed <- struct{}{}
		}
	}
	server.Start()
	defer server.Close()
	transport := &http.Transport{MaxIdleConns: 1, MaxIdleConnsPerHost: 1}
	defer transport.CloseIdleConnections()
	client := &http.Client{Transport: transport, Timeout: 2 * time.Second}

	get := func(path string) int64 {
		response, err := client.Get(server.URL + path)
		check(err == nil, fmt.Sprint(err))
		body, readErr := io.ReadAll(response.Body)
		closeErr := response.Body.Close()
		check(readErr == nil && closeErr == nil, "response body did not finish cleanly")
		check(response.Proto == "HTTP/1.1" && response.StatusCode == 200, "unexpected protocol or status")
		check(string(body) == "ok\n", "unexpected response body")
		id, err := strconv.ParseInt(response.Header.Get("X-Connection-ID"), 10, 64)
		check(err == nil, "missing connection identity")
		return id
	}

	first := get("/ping")
	check(get("/ping") == first, "the fully consumed responses should allow reuse")
	started := time.Now()
	check(get("/slow") == first, "the active request should keep its connection")
	check(time.Since(started) >= work, "the handler did not complete its work")
	fmt.Println("two requests reuse one HTTP/1.1 connection; active work outlasts IdleTimeout")

	select {
	case <-closed:
	case <-time.After(2 * time.Second):
		panic("the idle connection was not closed")
	}
	check(get("/ping") != first, "the request after idle closure needs a new connection")
	fmt.Println("server reports idle closure; a later GET succeeds on a new connection")

	before := handled.Load()
	connection, err := net.DialTimeout("tcp", server.Listener.Addr().String(), time.Second)
	check(err == nil, fmt.Sprint(err))
	defer connection.Close()
	check(connection.SetDeadline(time.Now().Add(2*time.Second)) == nil, "could not bound raw probe")
	_, err = io.WriteString(connection, "GET /ping HTTP/1.1\r\nHost: probe\r\n")
	check(err == nil, fmt.Sprint(err))
	reply, err := io.ReadAll(connection)
	check(err == nil && len(reply) == 0, "partial headers should end with EOF in this server")
	check(handled.Load() == before, "incomplete headers must not invoke the handler")
	fmt.Println("incomplete request headers close without reaching the handler; ReadHeaderTimeout is separate")
	fmt.Println(runtime.Version())
}