go-omni/transport/http2/lifecycle_test.go

package http2

import (
	"bytes"
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"testing"
	"time"

	xhttp2 "golang.org/x/net/http2"
	"golang.org/x/net/http2/hpack"
)

// A raw peer makes GOAWAY and RST_STREAM observable without client retries.
type framePeer struct {
	net.Conn
	framer *xhttp2.Framer
}

func newFramePeer(t *testing.T, address string) *framePeer {
	t.Helper()
	c, err := net.DialTimeout("tcp", address, time.Second)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { c.Close() })
	c.SetDeadline(time.Now().Add(5 * time.Second))
	p := &framePeer{c, xhttp2.NewFramer(c, c)}
	if _, err = io.WriteString(c, xhttp2.ClientPreface); err != nil {
		t.Fatal(err)
	}
	if err = p.framer.WriteSettings(); err != nil {
		t.Fatal(err)
	}
	return p
}
func (p *framePeer) get(t *testing.T, id uint32, path string) {
	t.Helper()
	var b bytes.Buffer
	e := hpack.NewEncoder(&b)
	for _, field := range []hpack.HeaderField{{Name: ":method", Value: "GET"}, {Name: ":scheme", Value: "http"}, {Name: ":authority", Value: "localhost"}, {Name: ":path", Value: path}} {
		if err := e.WriteField(field); err != nil {
			t.Fatal(err)
		}
	}
	if err := p.framer.WriteHeaders(xhttp2.HeadersFrameParam{StreamID: id, EndStream: true, EndHeaders: true, BlockFragment: b.Bytes()}); err != nil {
		t.Fatal(err)
	}
}
func (p *framePeer) read(t *testing.T) xhttp2.Frame {
	t.Helper()
	f, err := p.framer.ReadFrame()
	if err != nil {
		t.Fatal(err)
	}
	if s, ok := f.(*xhttp2.SettingsFrame); ok && !s.IsAck() {
		if err := p.framer.WriteSettingsAck(); err != nil {
			t.Fatal(err)
		}
	}
	return f
}
func (p *framePeer) body(t *testing.T, id uint32) string {
	t.Helper()
	var b bytes.Buffer
	for {
		switch f := p.read(t).(type) {
		case *xhttp2.DataFrame:
			if f.StreamID != id {
				t.Fatalf("unexpected stream %d", f.StreamID)
			}
			b.Write(f.Data())
			if f.StreamEnded() {
				return b.String()
			}
		case *xhttp2.HeadersFrame:
			if f.StreamID == id && f.StreamEnded() {
				return b.String()
			}
		case *xhttp2.RSTStreamFrame:
			t.Fatalf("unexpected RST_STREAM: %v", f)
		}
	}
}
func waitSignal(t *testing.T, signal <-chan struct{}) {
	t.Helper()
	select {
	case <-signal:
	case <-time.After(3 * time.Second):
		t.Fatal("timed out waiting for event")
	}
}
func launchServer(t *testing.T, s *Server, ctx context.Context) (string, <-chan error) {
	t.Helper()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() { done <- s.Serve(ctx, l) }()
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		if err := s.Shutdown(ctx); err != nil {
			t.Errorf("cleanup shutdown: %v", err)
		}
		l.Close()
	})
	return l.Addr().String(), done
}
func serveResult(t *testing.T, done <-chan error) error {
	t.Helper()
	select {
	case err := <-done:
		return err
	case <-time.After(3 * time.Second):
		t.Fatal("Serve did not finish")
		return nil
	}
}
func TestShutdownSendsGoAwayAndDrains(t *testing.T) {
	started := make(chan struct{})
	release := make(chan struct{})
	defer close(release)
	requestCtx := make(chan context.Context, 1)
	s := New(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/hold" {
			t.Error("new stream reached handler during drain")
			return
		}
		requestCtx <- r.Context()
		close(started)
		select {
		case <-release:
		case <-r.Context().Done():
			return
		}
		io.WriteString(w, "completed")
	}))
	address, done := launchServer(t, s, context.Background())
	p := newFramePeer(t, address)
	p.get(t, 1, "/hold")
	waitSignal(t, started)
	rctx := <-requestCtx
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	stopped := make(chan error, 1)
	go func() { stopped <- s.Shutdown(ctx) }()
	for {
		if f, ok := p.read(t).(*xhttp2.GoAwayFrame); ok {
			if f.ErrCode != xhttp2.ErrCodeNo || f.LastStreamID != 1 {
				t.Fatalf("GOAWAY=%+v", f)
			}
			break
		}
	}
	if err := rctx.Err(); err != nil {
		t.Fatalf("graceful shutdown canceled active request: %v", err)
	}
	select {
	case err := <-done:
		t.Fatalf("Serve returned before drain: %v", err)
	default:
	}
	p.get(t, 3, "/late")
	release <- struct{}{}
	if got := p.body(t, 1); got != "completed" {
		t.Fatalf("body=%q", got)
	}
	p.Close()
	if err := serveResult(t, stopped); err != nil {
		t.Fatal(err)
	}
	if err := serveResult(t, done); err != nil {
		t.Fatal(err)
	}
	if stats := s.Stats(); stats.ActiveConnections != 0 || stats.ActiveRequests != 0 {
		t.Fatalf("not drained: %+v", stats)
	}
}
func TestRSTStreamCancelsOnlyOneRequest(t *testing.T) {
	started := make(chan struct{}, 2)
	canceled := make(chan string, 2)
	release := make(chan struct{})
	defer close(release)
	s := New(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		started <- struct{}{}
		select {
		case <-r.Context().Done():
			canceled <- r.URL.Path
			return
		case <-release:
			io.WriteString(w, r.URL.Path)
		}
	}))
	s.MaxConcurrentRequests = 2
	address, _ := launchServer(t, s, context.Background())
	p := newFramePeer(t, address)
	p.get(t, 1, "/cancel")
	p.get(t, 3, "/survive")
	waitSignal(t, started)
	waitSignal(t, started)
	if err := p.framer.WriteRSTStream(1, xhttp2.ErrCodeCancel); err != nil {
		t.Fatal(err)
	}
	select {
	case path := <-canceled:
		if path != "/cancel" {
			t.Fatalf("canceled %s", path)
		}
	case <-time.After(time.Second):
		t.Fatal("RST did not cancel request")
	}
	release <- struct{}{}
	if got := p.body(t, 3); got != "/survive" {
		t.Fatalf("sibling body=%q", got)
	}
	p.get(t, 5, "/reused")
	waitSignal(t, started)
	release <- struct{}{}
	if got := p.body(t, 5); got != "/reused" {
		t.Fatalf("connection reuse body=%q", got)
	}
}
func TestCanceledServeContextWaitsForShutdownDeadline(t *testing.T) {
	started, canceled := make(chan struct{}), make(chan struct{})
	s := New(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { close(started); <-r.Context().Done(); close(canceled) }))
	s.ShutdownTimeout = 100 * time.Millisecond
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	address, done := launchServer(t, s, ctx)
	p := newFramePeer(t, address)
	p.get(t, 1, "/hold")
	waitSignal(t, started)
	cancel()
	if err := serveResult(t, done); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("Serve should report drain timeout, got %v", err)
	}
	waitSignal(t, canceled)
}