From b87d70c71c704366d7b8145c99b4293749b97d7e Mon Sep 17 00:00:00 2001
From: James Shubin <james@shubin.ca>
Date: Mon, 7 Aug 2023 23:20:33 -0400
Subject: [PATCH] engine: resources: Move wait out of Close method

Watch should wait until everything is done before exiting.
---
 engine/resources/config_etcd.go | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/engine/resources/config_etcd.go b/engine/resources/config_etcd.go
index f3eb8a1b..e00bcbcb 100644
--- a/engine/resources/config_etcd.go
+++ b/engine/resources/config_etcd.go
@@ -57,7 +57,6 @@ type ConfigEtcdRes struct {
 	sizeFlag bool
 
 	interruptChan chan struct{}
-	wg            *sync.WaitGroup
 }
 
 // Default returns some sensible defaults for this resource.
@@ -83,25 +82,22 @@ func (obj *ConfigEtcdRes) Init(init *engine.Init) error {
 	obj.init = init // save for later
 
 	obj.interruptChan = make(chan struct{})
-	obj.wg = &sync.WaitGroup{}
 
 	return nil
 }
 
 // Close is run by the engine to clean up after the resource is done.
 func (obj *ConfigEtcdRes) Close() error {
-	obj.wg.Wait() // bonus
 	return nil
 }
 
 // Watch is the primary listener for this resource and it outputs events.
 func (obj *ConfigEtcdRes) Watch(ctx context.Context) error {
-	obj.wg.Add(1)
-	defer obj.wg.Done()
-	// FIXME: add timeout to context
-	ctx, cancel := context.WithCancel(ctx)
+	wg := &sync.WaitGroup{}
+	defer wg.Wait()
+	innerCtx, cancel := context.WithCancel(ctx)
 	defer cancel()
-	ch, err := obj.init.World.IdealClusterSizeWatch(util.CtxWithWg(ctx, obj.wg))
+	ch, err := obj.init.World.IdealClusterSizeWatch(util.CtxWithWg(innerCtx, wg))
 	if err != nil {
 		return errwrap.Wrapf(err, "could not watch ideal cluster size")
 	}