engine: resources: Move wait out of Close method

Watch should wait until everything is done before exiting.
This commit is contained in:
James Shubin
2023-08-07 23:20:33 -04:00
parent 3f403d34a4
commit b87d70c71c

View File

@@ -57,7 +57,6 @@ type ConfigEtcdRes struct {
sizeFlag bool sizeFlag bool
interruptChan chan struct{} interruptChan chan struct{}
wg *sync.WaitGroup
} }
// Default returns some sensible defaults for this resource. // 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.init = init // save for later
obj.interruptChan = make(chan struct{}) obj.interruptChan = make(chan struct{})
obj.wg = &sync.WaitGroup{}
return nil return nil
} }
// Close is run by the engine to clean up after the resource is done. // Close is run by the engine to clean up after the resource is done.
func (obj *ConfigEtcdRes) Close() error { func (obj *ConfigEtcdRes) Close() error {
obj.wg.Wait() // bonus
return nil return nil
} }
// Watch is the primary listener for this resource and it outputs events. // Watch is the primary listener for this resource and it outputs events.
func (obj *ConfigEtcdRes) Watch(ctx context.Context) error { func (obj *ConfigEtcdRes) Watch(ctx context.Context) error {
obj.wg.Add(1) wg := &sync.WaitGroup{}
defer obj.wg.Done() defer wg.Wait()
// FIXME: add timeout to context innerCtx, cancel := context.WithCancel(ctx)
ctx, cancel := context.WithCancel(ctx)
defer cancel() 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 { if err != nil {
return errwrap.Wrapf(err, "could not watch ideal cluster size") return errwrap.Wrapf(err, "could not watch ideal cluster size")
} }