HTTP/2 多路复用消除排队

package main

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/http/httptrace"
	"runtime"
	"sync/atomic"
	"time"
)

type connectionKey struct{}

func must(err error) {
	if err != nil {
		panic(err)
	}
}

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

func main() {
	cancelObserved, releaseSurvivor, stop := make(chan struct{}), make(chan struct{}), make(chan struct{})
	var connections, handled atomic.Int64
	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		handled.Add(1)
		w.Header().Set("X-Connection-ID", fmt.Sprint(r.Context().Value(connectionKey{})))
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusOK)
		w.(http.Flusher).Flush()
		switch r.URL.Path {
		case "/cancel":
			select {
			case <-r.Context().Done():
				close(cancelObserved)
			case <-stop:
			}
		case "/survive":
			select {
			case <-releaseSurvivor:
				fmt.Fprint(w, "survived")
			case <-r.Context().Done():
			case <-stop:
			}
		case "/after":
			fmt.Fprint(w, "ready")
		}
	}))
	server.Config.ConnContext = func(ctx context.Context, _ net.Conn) context.Context {
		return context.WithValue(ctx, connectionKey{}, connections.Add(1))
	}
	server.EnableHTTP2 = true
	server.TLS = &tls.Config{MinVersion: tls.VersionTLS13, MaxVersion: tls.VersionTLS13}
	server.StartTLS()
	defer server.Close()
	defer close(stop)
	config := server.Client().Transport.(*http.Transport).TLSClientConfig.Clone()
	check(!config.InsecureSkipVerify, "certificate verification must stay enabled")
	transport := &http.Transport{TLSClientConfig: config, ForceAttemptHTTP2: true, MaxConnsPerHost: 1}
	defer transport.CloseIdleConnections()
	client := &http.Client{Transport: transport, Timeout: 2 * time.Second}
	open := func(ctx context.Context, route string) *http.Response {
		request, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL+route, nil)
		must(err)
		response, err := client.Do(request)
		must(err)
		check(response.StatusCode == 200 && response.Proto == "HTTP/2.0", "HTTP/2 response required")
		check(response.Header.Get("X-Connection-ID") == "1", "streams must share the first TCP connection")
		check(response.TLS.Version == tls.VersionTLS13 && response.TLS.NegotiatedProtocol == "h2", "TLS 1.3 and h2 required")
		check(len(response.TLS.VerifiedChains) > 0, "verified server certificate required")
		return response
	}
	firstContext, cancel := context.WithCancel(context.Background())
	defer cancel()
	first := open(firstContext, "/cancel")
	survivor := open(context.Background(), "/survive")
	cancel()
	body, readError := io.ReadAll(first.Body)
	closeError := first.Body.Close()
	must(closeError)
	check(errors.Is(readError, context.Canceled) && len(body) == 0, "first stream must report cancellation")
	select {
	case <-cancelObserved:
	case <-time.After(2 * time.Second):
		panic("server did not observe stream cancellation")
	}
	close(releaseSurvivor)
	body, readError = io.ReadAll(survivor.Body)
	closeError = survivor.Body.Close()
	must(readError)
	must(closeError)
	check(string(body) == "survived", "sibling stream must complete")
	var reused atomic.Bool
	var handshakes atomic.Int64
	trace := &httptrace.ClientTrace{
		GotConn:           func(info httptrace.GotConnInfo) { reused.Store(info.Reused) },
		TLSHandshakeStart: func() { handshakes.Add(1) },
	}
	after := open(httptrace.WithClientTrace(context.Background(), trace), "/after")
	body, readError = io.ReadAll(after.Body)
	closeError = after.Body.Close()
	must(readError)
	must(closeError)
	check(string(body) == "ready" && reused.Load() && handshakes.Load() == 0, "connection must remain reusable")
	check(connections.Load() == 1 && handled.Load() == 3, "three streams must use one TCP connection")
	fmt.Println("cancel: context canceled; server observes cancellation")
	fmt.Println("survive: body complete on connection 1")
	fmt.Println("after: connection 1 reused, no new TLS handshake")
	fmt.Println(runtime.Version())
}