webhooker/internal/server/shutdown_test.go

package server_test

import (
	"context"
	"net"
	"net/http"
	"testing"
	"testing/synctest"
	"time"

	"github.com/stretchr/testify/require"
	"sneak.berlin/go/webhooker/internal/server"
)

// TestDrainBudget covers the clamp that keeps the HTTP drain from
// spending the tail hooks' share of the fx stop budget when the hooks
// before the server have already used part of it.
func TestDrainBudget(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name      string
		remaining time.Duration
		want      time.Duration
	}{
		{
			name:      "only the reserve is left",
			remaining: server.TailHookReserve,
			want:      0,
		},
		{
			name:      "earlier hooks spent part of the budget",
			remaining: server.TailHookReserve + time.Second,
			want:      time.Second,
		},
		{
			name:      "capped at the nominal timeout",
			remaining: time.Hour,
			want:      server.ShutdownTimeout,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			require.Equal(t, tt.want, server.DrainBudget(tt.remaining))
		})
	}
}

// TestCleanShutdown_LeavesTailHookReserve stops the server with a
// request still in flight, after the hooks before it have spent all
// of the stop budget but TailHookReserve. The drain must give up at
// once rather than wait for the request: what is left belongs to the
// hooks after the server, the database close among them. A drain
// bounded only by ShutdownTimeout waits until the stop context
// expires, and fx then skips those hooks.
//
// The test runs in a synctest bubble, whose clock moves only while
// every goroutine in it is blocked, so a drain that gives up at once
// leaves the stop context unexpired however slow the host is. The
// request travels over net.Pipe because a goroutine waiting on a
// real socket would stop that clock from moving at all.
func TestCleanShutdown_LeavesTailHookReserve(t *testing.T) {
	t.Parallel()

	synctest.Test(t, func(t *testing.T) {
		entered := make(chan struct{})
		release := make(chan struct{})

		hs := &http.Server{
			Handler: http.HandlerFunc(
				func(http.ResponseWriter, *http.Request) {
					close(entered)
					<-release
				},
			),
			ReadHeaderTimeout: time.Second,
		}

		srvConn, cliConn := net.Pipe()

		listener := pipeListener{
			conns:  make(chan net.Conn, 1),
			closed: make(chan struct{}),
		}
		listener.conns <- srvConn

		go func() { _ = hs.Serve(listener) }()

		// Cleanups run last first: the handler returns, then closing
		// the client end ends the server's write of the response.
		t.Cleanup(func() { _ = cliConn.Close() })
		t.Cleanup(func() { close(release) })

		_, err := cliConn.Write(
			[]byte("GET / HTTP/1.1\r\nHost: webhooker.test\r\n\r\n"),
		)
		require.NoError(t, err)

		<-entered

		stopCtx, cancel := context.WithTimeout(
			t.Context(), server.TailHookReserve,
		)
		defer cancel()

		server.CleanShutdownForTest(stopCtx, hs)

		require.NoError(
			t, stopCtx.Err(), "the drain spent the tail hooks' reserve",
		)
	})
}

// pipeListener is the net.Listener http.Server.Serve needs to serve
// the server end of a net.Pipe: Accept returns that one connection,
// then waits until Close, as a real listener with no more clients
// does.
type pipeListener struct {
	conns  chan net.Conn
	closed chan struct{}
}

func (l pipeListener) Accept() (net.Conn, error) {
	select {
	case conn := <-l.conns:
		return conn, nil
	case <-l.closed:
		return nil, net.ErrClosed
	}
}

func (l pipeListener) Close() error {
	close(l.closed)

	return nil
}

// Addr is never called by http.Server.Serve.
func (pipeListener) Addr() net.Addr {
	return nil
}

// TestSentryFlushBudget covers the clamp that keeps the Sentry flush
// from spending the tail hooks' share of the fx stop budget.
// sentry.Flush ignores the stop context, so without the clamp a
// stalled flush adds its whole timeout on top of the HTTP drain.
func TestSentryFlushBudget(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name      string
		remaining time.Duration
		want      time.Duration
	}{
		{
			name:      "full drain leaves only the reserve",
			remaining: server.TailHookReserve,
			want:      0,
		},
		{
			name:      "expired budget",
			remaining: -time.Second,
			want:      0,
		},
		{
			name:      "sliver above the reserve is not worth it",
			remaining: server.TailHookReserve + 10*time.Millisecond,
			want:      0,
		},
		{
			name:      "partial flush when some room is left",
			remaining: server.TailHookReserve + time.Second,
			want:      time.Second,
		},
		{
			name:      "capped at the nominal timeout",
			remaining: time.Hour,
			want:      2 * time.Second,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			require.Equal(
				t, tt.want, server.SentryFlushBudget(tt.remaining),
			)
		})
	}
}