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

internal/delivery/engine.go

@@ -13,6 +13,7 @@ import (
 	"go.uber.org/fx"
 	"gorm.io/gorm"
 	"sneak.berlin/go/webhooker/internal/database"
+	"sneak.berlin/go/webhooker/internal/lifecycle"
 	"sneak.berlin/go/webhooker/internal/logger"
 )
 
@@ -234,8 +235,9 @@ func (e *Engine) ScheduleRetry(
 }
 
 // registerHooks wires the engine's start and stop into the fx
-// lifecycle. The start hook's context is deliberately ignored:
-// see start for why the worker pool must not inherit it.
+// lifecycle. The start hook's context is deliberately ignored
+// (see start for why the worker pool must not inherit it); the
+// stop hook's context is honoured (see stop).
 func (e *Engine) registerHooks(lc fx.Lifecycle) {
 	lc.Append(fx.Hook{
 		//nolint:contextcheck // Not inheriting the hook context
@@ -245,10 +247,8 @@ func (e *Engine) registerHooks(lc fx.Lifecycle) {
 
 			return nil
 		},
-		OnStop: func(_ context.Context) error {
-			e.stop()
-
-			return nil
+		OnStop: func(ctx context.Context) error {
+			return e.stop(ctx)
 		},
 	})
 }
@@ -289,11 +289,26 @@ func (e *Engine) start() {
 	)
 }
 
-func (e *Engine) stop() {
+// stop cancels the worker pool's context and waits for the pool
+// to drain, bounded by the stop hook's context: a wedged worker
+// must not hang the process past fx's stop timeout.
+func (e *Engine) stop(ctx context.Context) error {
 	e.log.Info("delivery engine stopping")
-	e.cancel()
-	e.wg.Wait()
+
+	if e.cancel != nil {
+		e.cancel()
+	}
+
+	err := lifecycle.WaitForShutdown(
+		ctx, e.log, "delivery engine", &e.wg,
+	)
+	if err != nil {
+		return err
+	}
+
 	e.log.Info("delivery engine stopped")
+
+	return nil
 }
 
 func (e *Engine) worker(ctx context.Context) {