Bound shutdown hooks by their stop context (closes #102)

internal/database/retention.go

@@ -10,6 +10,7 @@ import (
 	"go.uber.org/fx"
 	"gorm.io/gorm"
 	"sneak.berlin/go/webhooker/internal/config"
+	"sneak.berlin/go/webhooker/internal/lifecycle"
 	"sneak.berlin/go/webhooker/internal/logger"
 )
 
@@ -62,8 +63,9 @@ func NewRetentionReaper(
 }
 
 // registerHooks wires the reaper's start and stop into the fx
-// lifecycle. The start hook's context is deliberately ignored: see
-// start for why the sweep loop must not inherit it.
+// lifecycle. The start hook's context is deliberately ignored (see
+// start for why the sweep loop must not inherit it); the stop hook's
+// context is honoured (see stop).
 func (r *RetentionReaper) registerHooks(lc fx.Lifecycle) {
 	lc.Append(fx.Hook{
 		//nolint:contextcheck // Not inheriting the hook context is
@@ -73,10 +75,8 @@ func (r *RetentionReaper) registerHooks(lc fx.Lifecycle) {
 
 			return nil
 		},
-		OnStop: func(_ context.Context) error {
-			r.stop()
-
-			return nil
+		OnStop: func(ctx context.Context) error {
+			return r.stop(ctx)
 		},
 	})
 }
@@ -105,15 +105,27 @@ func (r *RetentionReaper) start() {
 	)
 }
 
-func (r *RetentionReaper) stop() {
+// stop cancels the sweep loop's context and waits for it to
+// exit, bounded by the stop hook's context: a sweep wedged on a
+// locked database must not hang the process past fx's stop
+// timeout.
+func (r *RetentionReaper) stop(ctx context.Context) error {
 	r.log.Info("retention reaper stopping")
 
 	if r.cancel != nil {
 		r.cancel()
 	}
 
-	r.wg.Wait()
+	err := lifecycle.WaitForShutdown(
+		ctx, r.log, "retention reaper", &r.wg,
+	)
+	if err != nil {
+		return err
+	}
+
 	r.log.Info("retention reaper stopped")
+
+	return nil
 }
 
 func (r *RetentionReaper) run(ctx context.Context) {