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

internal/delivery/export_test.go

@@ -216,8 +216,19 @@ func (e *Engine) ExportRegisterHooks(lc fx.Lifecycle) {
 }
 
 // ExportStop exposes stop for testing.
-func (e *Engine) ExportStop() {
-	e.stop()
+func (e *Engine) ExportStop(ctx context.Context) error {
+	return e.stop(ctx)
+}
+
+// ExportWedgeWorker adds a goroutine to the engine's WaitGroup
+// that never observes cancellation and returns only when release
+// is closed. It stands in for a worker stuck inside a delivery
+// target that never returns, which is the only way stop can be
+// made to outlast its context.
+func (e *Engine) ExportWedgeWorker(release <-chan struct{}) {
+	e.wg.Go(func() {
+		<-release
+	})
 }
 
 // ExportDeliveryCh returns the delivery channel.
@@ -518,8 +529,19 @@ func (s *ArchiveSweeper) ExportRegisterHooks(lc fx.Lifecycle) {
 }
 
 // ExportStop stops the sweeper's background loop for tests.
-func (s *ArchiveSweeper) ExportStop() {
-	s.stop()
+func (s *ArchiveSweeper) ExportStop(ctx context.Context) error {
+	return s.stop(ctx)
+}
+
+// ExportWedgeLoop adds a goroutine to the sweeper's WaitGroup
+// that never observes cancellation and returns only when release
+// is closed. It stands in for a prune stuck on a locked archive.
+func (s *ArchiveSweeper) ExportWedgeLoop(
+	release <-chan struct{},
+) {
+	s.wg.Go(func() {
+		<-release
+	})
 }
 
 // ExportSetInterval overrides the sweep interval for tests.