webhooker/internal/server/shutdown_test.go

package server_test

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"
	"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.
func TestCleanShutdown_LeavesTailHookReserve(t *testing.T) {
	t.Parallel()

	entered := make(chan struct{})
	release := make(chan struct{})

	ts := httptest.NewServer(http.HandlerFunc(
		func(http.ResponseWriter, *http.Request) {
			close(entered)
			<-release
		},
	))
	// Cleanups run last first: the request is released before Close,
	// which waits for it.
	t.Cleanup(ts.Close)
	t.Cleanup(func() { close(release) })

	req, err := http.NewRequestWithContext(
		t.Context(), http.MethodGet, ts.URL, nil,
	)
	require.NoError(t, err)

	go func() {
		resp, err := ts.Client().Do(req)
		if err == nil {
			_ = resp.Body.Close()
		}
	}()

	<-entered

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

	server.CleanShutdownForTest(stopCtx, ts.Config)

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

// 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),
			)
		})
	}
}