notify: drain in-flight deliveries at shutdown (closes #106)

Shutdown waits for deliveries already in flight instead of dropping them.

Reviewed and green; squashed to next by the dispatcher, dnswatcher having no manager.

Model: opus-5

internal/notify/shutdown_test.go

@@ -0,0 +1,531 @@
+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,
+		)
+	}
+}