Graceful shutdowns (#1734)

* Initial graceful stop

* Fix dendritejs

* Use process context for outbound federation requests in destination queues

* Reduce logging

* Fix log level

1 files changed, 45 insertions, 0 deletions

diff --git a/setup/process/process.go b/setup/process/process.go
new file mode 100644
index 00000000..d55751d7
--- /dev/null
+++ b/setup/process/process.go
@@ -0,0 +1,45 @@
+package process
+
+import (
+	"context"
+	"sync"
+)
+
+type ProcessContext struct {
+	wg       *sync.WaitGroup    // used to wait for components to shutdown
+	ctx      context.Context    // cancelled when Stop is called
+	shutdown context.CancelFunc // shut down Dendrite
+}
+
+func NewProcessContext() *ProcessContext {
+	ctx, shutdown := context.WithCancel(context.Background())
+	return &ProcessContext{
+		ctx:      ctx,
+		shutdown: shutdown,
+		wg:       &sync.WaitGroup{},
+	}
+}
+
+func (b *ProcessContext) Context() context.Context {
+	return context.WithValue(b.ctx, "scope", "process") // nolint:staticcheck
+}
+
+func (b *ProcessContext) ComponentStarted() {
+	b.wg.Add(1)
+}
+
+func (b *ProcessContext) ComponentFinished() {
+	b.wg.Done()
+}
+
+func (b *ProcessContext) ShutdownDendrite() {
+	b.shutdown()
+}
+
+func (b *ProcessContext) WaitForShutdown() <-chan struct{} {
+	return b.ctx.Done()
+}
+
+func (b *ProcessContext) WaitForComponentsToFinish() {
+	b.wg.Wait()
+}