From e827079604227c6e2473f82f2afcdf5b57120fa4 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Wed, 30 Dec 2020 18:32:58 +0000
Subject: [PATCH] Code maintenance: updater loop waitgroup

---
 internal/updater/loop.go | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/internal/updater/loop.go b/internal/updater/loop.go
index deaffe2e..0d938b81 100644
--- a/internal/updater/loop.go
+++ b/internal/updater/loop.go
@@ -94,10 +94,13 @@ func (l *looper) Run(ctx context.Context, wg *sync.WaitGroup) {
 
 	for ctx.Err() == nil {
 		updateCtx, updateCancel := context.WithCancel(ctx)
-		defer updateCancel()
+
 		serversCh := make(chan models.AllServers)
 		errorCh := make(chan error)
+		runWg := &sync.WaitGroup{}
+		runWg.Add(1)
 		go func() {
+			defer runWg.Done()
 			servers, err := l.updater.UpdateServers(updateCtx)
 			if err != nil {
 				if updateCtx.Err() == nil {
@@ -122,33 +125,37 @@ func (l *looper) Run(ctx context.Context, wg *sync.WaitGroup) {
 			case <-ctx.Done():
 				l.logger.Warn("context canceled: exiting loop")
 				updateCancel()
+				runWg.Wait()
 				close(errorCh)
 				return
 			case <-l.start:
 				l.logger.Info("starting")
 				updateCancel()
+				runWg.Wait()
 				stayHere = false
 			case <-l.stop:
 				l.logger.Info("stopping")
 				updateCancel()
+				runWg.Wait()
 				l.stopped <- struct{}{}
 			case servers := <-serversCh:
-				updateCancel()
 				l.setAllServers(servers)
 				if err := l.storage.FlushToFile(servers); err != nil {
 					l.logger.Error(err)
 				}
+				runWg.Wait()
 				l.state.setStatusWithLock(constants.Completed)
 				l.logger.Info("Updated servers information")
 			case err := <-errorCh:
-				updateCancel()
 				close(serversCh)
+				runWg.Wait()
 				l.state.setStatusWithLock(constants.Crashed)
 				l.logAndWait(ctx, err)
 				crashed = true
 				stayHere = false
 			}
 		}
+		updateCancel()
 		close(errorCh)
 	}
 }