dnswatcher/internal/notify/shutdown_test.go

package notify_test

import (
	"bytes"
	"context"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"go.uber.org/fx"

	"sneak.berlin/go/dnswatcher/internal/config"
	"sneak.berlin/go/dnswatcher/internal/globals"
	"sneak.berlin/go/dnswatcher/internal/logger"
	"sneak.berlin/go/dnswatcher/internal/notify"
)

// Timings used by the drain tests.  They stay in the same
// 10-100ms band as the retry tests so the suite never waits on
// a real backoff delay.
const (
	// inFlightHold is how long a delivery is kept mid-request
	// before the handler is released.
	inFlightHold = 30 * time.Millisecond

	// drainDeadline bounds a drain that is expected to time
	// out.
	drainDeadline = 50 * time.Millisecond

	// drainSlack is the upper bound on how long a bounded
	// drain may take; generous enough for a loaded CI box,
	// still far below the 20s test ceiling.
	drainSlack = 2 * time.Second

	// settleDelay is how long to wait before asserting that
	// something did *not* happen.
	settleDelay = 50 * time.Millisecond

	// idleDrainBound is the upper bound on a drain that has
	// nothing in flight.  It is deliberately far above the cost
	// of the goroutine hop through inFlight.Wait() — which
	// reached 57ms on a loaded box under -race with the package's
	// parallel tests — and far below drainSlack, the deadline
	// such a drain is given.  A drain that blocked until its
	// deadline instead of returning on the WaitGroup therefore
	// still fails this bound, but scheduling delay alone cannot.
	idleDrainBound = 500 * time.Millisecond
)

// syncBuffer is an io.Writer safe for concurrent use, so log
// output written from delivery goroutines can be inspected.
type syncBuffer struct {
	mu  sync.Mutex
	buf bytes.Buffer
}

func (sb *syncBuffer) Write(p []byte) (int, error) {
	sb.mu.Lock()
	defer sb.mu.Unlock()

	return sb.buf.Write(p) //nolint:wrapcheck // test helper
}

func (sb *syncBuffer) String() string {
	sb.mu.Lock()
	defer sb.mu.Unlock()

	return sb.buf.String()
}

// newLoggingService returns a Service writing JSON logs into
// the returned buffer.
func newLoggingService(
	transport http.RoundTripper,
) (*notify.Service, *syncBuffer) {
	logs := &syncBuffer{}
	handler := slog.NewJSONHandler(logs, nil)

	return notify.NewTestServiceWithLogger(transport, handler),
		logs
}

// blockingNtfyServer returns a server whose handler signals on
// entered, waits for release, and then responds 200.
func blockingNtfyServer(
	entered chan<- struct{},
	release <-chan struct{},
	served *atomic.Bool,
) *httptest.Server {
	var once sync.Once

	return httptest.NewServer(
		http.HandlerFunc(
			func(w http.ResponseWriter, _ *http.Request) {
				once.Do(func() { close(entered) })
				<-release

				served.Store(true)

				w.WriteHeader(http.StatusOK)
			}),
	)
}

// TestDrainWaitsForInFlightDelivery verifies that a delivery
// already under way when shutdown starts is allowed to finish.
func TestDrainWaitsForInFlightDelivery(t *testing.T) {
	t.Parallel()

	var served atomic.Bool

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

	srv := blockingNtfyServer(entered, release, &served)
	defer srv.Close()

	topicURL, _ := url.Parse(srv.URL)

	svc := notify.NewTestService(http.DefaultTransport)
	svc.SetNtfyURL(topicURL)

	svc.SendNotification(
		context.Background(), "t", "m", prioInfo,
	)

	// Make sure the delivery really is mid-request before the
	// drain begins.
	select {
	case <-entered:
	case <-time.After(drainSlack):
		t.Fatal("delivery never reached the endpoint")
	}

	// As in TestDrainBoundedByContextDeadline: start is captured
	// before the clock it is compared against, here the timer
	// holding the delivery open, so elapsed covers the whole hold
	// and the lower bound cannot come out short from scheduling
	// delay alone.
	start := time.Now()

	timer := time.AfterFunc(inFlightHold, func() {
		close(release)
	})
	defer timer.Stop()

	ctx, cancel := context.WithTimeout(
		context.Background(), drainSlack,
	)
	defer cancel()

	svc.Drain(ctx)

	elapsed := time.Since(start)

	if !served.Load() {
		t.Error(
			"drain returned before the in-flight delivery " +
				"completed",
		)
	}

	if elapsed < inFlightHold {
		t.Errorf(
			"drain took %v, want at least %v",
			elapsed, inFlightHold,
		)
	}

	if got := svc.OutstandingDeliveries(); got != 0 {
		t.Errorf("outstanding deliveries = %d, want 0", got)
	}
}

// neverFires returns a channel that never delivers, standing in
// for a long backoff sleep without actually sleeping.
func neverFires(_ time.Duration) <-chan time.Time {
	return make(chan time.Time)
}

// TestDrainBoundedByContextDeadline verifies that a delivery
// stuck retrying against a dead endpoint does not hold shutdown
// past the OnStop context deadline, and that the abandoned
// deliveries are logged at warn level rather than dropped
// silently.
func TestDrainBoundedByContextDeadline(t *testing.T) {
	t.Parallel()

	var requests atomic.Int64

	srv := httptest.NewServer(
		http.HandlerFunc(
			func(w http.ResponseWriter, _ *http.Request) {
				requests.Add(1)

				w.WriteHeader(http.StatusInternalServerError)
			}),
	)
	defer srv.Close()

	topicURL, _ := url.Parse(srv.URL)

	svc, logs := newLoggingService(http.DefaultTransport)
	svc.SetNtfyURL(topicURL)
	// Never let the backoff sleep complete: the delivery is
	// parked in its retry wait until shutdown releases it.
	svc.SetSleepFunc(neverFires)
	svc.SetRetryConfig(notify.RetryConfig{
		MaxRetries: 5,
		BaseDelay:  time.Hour,
		MaxDelay:   time.Hour,
	})

	svc.SendNotification(
		context.Background(), "t", "m", prioError,
	)

	waitForCondition(t, func() bool {
		return requests.Load() >= 1 &&
			svc.OutstandingDeliveries() == 1
	})

	// start must be captured *before* the deadline clock starts,
	// so that the measured interval is a superset of the deadline
	// interval.  Capturing it after context.WithTimeout would
	// make elapsed structurally smaller than drainDeadline and
	// the lower bound below unfalsifiable-by-luck: it would fail
	// whenever the two statements were separated by any
	// scheduling delay, and pass otherwise, regardless of what
	// the drain did.
	start := time.Now()

	ctx, cancel := context.WithTimeout(
		context.Background(), drainDeadline,
	)
	defer cancel()

	// The upper bound is enforced by a watchdog rather than by
	// measuring after the fact: a drain that is not bounded at
	// all never returns here (the delivery is parked in a backoff
	// that never fires), so an unbounded drain must fail this
	// test promptly instead of hanging the package until the test
	// binary's 30s timeout.
	returned := make(chan struct{})

	go func() {
		defer close(returned)

		svc.Drain(ctx)
	}()

	select {
	case <-returned:
	case <-time.After(drainSlack):
		t.Fatalf(
			"drain did not return within %v; its %v deadline "+
				"did not bound it",
			drainSlack, drainDeadline,
		)
	}

	// The lower bound is the real assertion: the drain must have
	// waited for its whole deadline rather than giving up on the
	// outstanding delivery early.  With start captured above, an
	// early return is the only thing that can make it fail.
	if elapsed := time.Since(start); elapsed < drainDeadline {
		t.Errorf(
			"drain returned after %v, before its %v deadline",
			elapsed, drainDeadline,
		)
	}

	assertAbandonLogged(t, logs.String())

	// The abandoned delivery must stop retrying rather than
	// outlive the drain.
	waitForCondition(t, func() bool {
		return svc.OutstandingDeliveries() == 0
	})
}

// assertAbandonLogged checks that the drain logged the
// abandoned deliveries at warn level with a count.
func assertAbandonLogged(t *testing.T, output string) {
	t.Helper()

	if !strings.Contains(output, `"level":"WARN"`) {
		t.Errorf(
			"abandoned deliveries not logged at warn level; "+
				"log output: %s",
			output,
		)
	}

	if !strings.Contains(output, `"abandoned":1`) {
		t.Errorf(
			"abandoned delivery count not logged; "+
				"log output: %s",
			output,
		)
	}
}

// TestDrainRefusesNewDeliveries verifies that notifications
// submitted after the drain has begun are refused and logged,
// so a stream of new work cannot extend shutdown indefinitely.
func TestDrainRefusesNewDeliveries(t *testing.T) {
	t.Parallel()

	var requests atomic.Int64

	srv := httptest.NewServer(
		http.HandlerFunc(
			func(w http.ResponseWriter, _ *http.Request) {
				requests.Add(1)

				w.WriteHeader(http.StatusOK)
			}),
	)
	defer srv.Close()

	target, _ := url.Parse(srv.URL)

	svc, logs := newLoggingService(http.DefaultTransport)
	svc.SetNtfyURL(target)
	svc.SetSlackWebhookURL(target)
	svc.SetMattermostWebhookURL(target)

	ctx, cancel := context.WithTimeout(
		context.Background(), drainSlack,
	)
	defer cancel()

	// Nothing is in flight, so this returns immediately and
	// leaves the service refusing further deliveries.
	svc.Drain(ctx)

	for range 3 {
		svc.SendNotification(
			context.Background(), "t", "m", prioInfo,
		)
	}

	time.Sleep(settleDelay)

	if got := requests.Load(); got != 0 {
		t.Errorf(
			"%d requests reached the endpoint after drain, "+
				"want 0",
			got,
		)
	}

	if got := svc.OutstandingDeliveries(); got != 0 {
		t.Errorf("outstanding deliveries = %d, want 0", got)
	}

	output := logs.String()
	if !strings.Contains(output, "shutdown in progress") {
		t.Errorf(
			"refused deliveries not logged; log output: %s",
			output,
		)
	}
}

// recordingLifecycle is a minimal fx.Lifecycle that records the
// hooks appended to it, so the wiring done by notify.New can be
// inspected without standing up a whole fx application.
type recordingLifecycle struct {
	hooks []fx.Hook
}

func (l *recordingLifecycle) Append(hook fx.Hook) {
	l.hooks = append(l.hooks, hook)
}

// newNotifyService builds a Service through the real
// constructor, wired to the given lifecycle.
func newNotifyService(
	t *testing.T,
	lifecycle fx.Lifecycle,
	ntfyTopic string,
) *notify.Service {
	t.Helper()

	g, err := globals.New(nil)
	if err != nil {
		t.Fatalf("globals.New: %v", err)
	}

	log, err := logger.New(nil, logger.Params{Globals: g})
	if err != nil {
		t.Fatalf("logger.New: %v", err)
	}

	svc, err := notify.New(lifecycle, notify.Params{
		Logger: log,
		Config: &config.Config{NtfyTopic: ntfyTopic},
	})
	if err != nil {
		t.Fatalf("notify.New: %v", err)
	}

	return svc
}

// TestNewRegistersDrainingStopHook verifies that notify.New
// wires an OnStop hook into the fx lifecycle and that the hook
// waits for in-flight deliveries.
func TestNewRegistersDrainingStopHook(t *testing.T) {
	t.Parallel()

	var served atomic.Bool

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

	srv := blockingNtfyServer(entered, release, &served)
	defer srv.Close()

	lifecycle := &recordingLifecycle{}
	svc := newNotifyService(t, lifecycle, srv.URL)

	if len(lifecycle.hooks) != 1 {
		t.Fatalf(
			"appended %d lifecycle hooks, want 1",
			len(lifecycle.hooks),
		)
	}

	stop := lifecycle.hooks[0].OnStop
	if stop == nil {
		t.Fatal("lifecycle hook has no OnStop function")
	}

	svc.SendNotification(
		context.Background(), "t", "m", prioInfo,
	)

	select {
	case <-entered:
	case <-time.After(drainSlack):
		t.Fatal("delivery never reached the endpoint")
	}

	timer := time.AfterFunc(inFlightHold, func() {
		close(release)
	})
	defer timer.Stop()

	ctx, cancel := context.WithTimeout(
		context.Background(), drainSlack,
	)
	defer cancel()

	err := stop(ctx)
	if err != nil {
		t.Fatalf("OnStop returned error: %v", err)
	}

	if !served.Load() {
		t.Error(
			"OnStop returned before the in-flight delivery " +
				"completed",
		)
	}
}

// TestDrainWithoutDeliveriesReturnsImmediately verifies the
// common case: nothing in flight, shutdown is not delayed.
func TestDrainWithoutDeliveriesReturnsImmediately(t *testing.T) {
	t.Parallel()

	svc := notify.NewTestService(http.DefaultTransport)

	// Captured before the deadline clock, as elsewhere in this
	// file; for an upper bound that is the conservative
	// direction, since the measured interval can then only be
	// longer than the drain itself.
	start := time.Now()

	ctx, cancel := context.WithTimeout(
		context.Background(), drainSlack,
	)
	defer cancel()

	svc.Drain(ctx)

	if elapsed := time.Since(start); elapsed > idleDrainBound {
		t.Errorf(
			"drain of an idle service took %v, want well "+
				"under its %v deadline",
			elapsed, drainSlack,
		)
	}
}

// TestDrainWithCancelledContextDoesNotWarn verifies that an
// OnStop context that is already dead on entry does not produce
// an "abandoning them" warning when there was nothing in flight
// to abandon.  The expired context wins the select immediately,
// so only the outstanding count can tell the difference between
// a genuine timeout and a shutdown that had simply already run
// out of time with no work left.
func TestDrainWithCancelledContextDoesNotWarn(t *testing.T) {
	t.Parallel()

	svc, logs := newLoggingService(http.DefaultTransport)

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	svc.Drain(ctx)

	if output := logs.String(); strings.Contains(
		output, `"level":"WARN"`,
	) {
		t.Errorf(
			"drain with nothing in flight warned about "+
				"abandoned deliveries; log output: %s",
			output,
		)
	}
}