with manager

manager.go

@@ -33,7 +33,9 @@ func NewManager() *Manager {
 type ComponentsAreReady struct {
 }
 
-func (m *Manager) RunNonBlocking(ctx context.Context, readyChan chan ComponentsAreReady) error {
+type CleanupFunc func(ctx context.Context) error
+
+func (m *Manager) RunNonBlocking(ctx context.Context, readyChan chan ComponentsAreReady) (CleanupFunc, error) {
 	go func() error {
 		m.initLifetime()
 		err := m.init(ctx)
@@ -53,13 +55,12 @@ func (m *Manager) RunNonBlocking(ctx context.Context, readyChan chan ComponentsA
 			return err
 		}
 
-		err = m.shutdown(ctx)
-		if err != nil {
-			return err
-		}
 		return nil
 	}()
-	return nil
+
+	return func(ctx context.Context) error {
+		return m.shutdown(ctx)
+	}, nil
 }
 
 func (m *Manager) Run(ctx context.Context) error {
@@ -162,8 +163,13 @@ func (m *Manager) initLifetime() {
 }
 
 func (m *Manager) wait(ctx context.Context) error {
-	exitCode := <-m.exitChan
-	m.exitCode = exitCode
+	select {
+	case exitCode := <-m.exitChan:
+		m.exitCode = exitCode
+		return nil
+	case <-ctx.Done():
+		return nil
+	}
 	return nil
 }