From e237e9642d00d85e4734141a9f1d6edb57b00728 Mon Sep 17 00:00:00 2001
From: ChrisMcKenzie <chris@kbsurfer.com>
Date: Tue, 10 Nov 2015 21:28:24 -0800
Subject: [PATCH] fixed #12 app now properly cleans up.

---
 commands/agent.go  |  6 +++++-
 work/dispatcher.go | 16 ++++++++++------
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/commands/agent.go b/commands/agent.go
index 09a5028..2bb7937 100644
--- a/commands/agent.go
+++ b/commands/agent.go
@@ -4,6 +4,7 @@ import (
 	"log"
 	"os"
 	"os/signal"
+	"sync"
 
 	"github.com/ChrisMcKenzie/dropship/service"
 	"github.com/ChrisMcKenzie/dropship/work"
@@ -28,8 +29,10 @@ func agent(c *cobra.Command, args []string) {
 	t := work.NewRunner(len(services))
 	shutdownCh := make(chan struct{})
 
+	var wg sync.WaitGroup
+	wg.Add(len(services))
 	for _, s := range services {
-		_, err := work.NewDispatcher(s, t, shutdownCh)
+		_, err := work.NewDispatcher(s, t, &wg, shutdownCh)
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -39,6 +42,7 @@ func agent(c *cobra.Command, args []string) {
 	signal.Notify(sigs, os.Interrupt)
 	<-sigs
 	close(shutdownCh)
+	wg.Wait()
 
 	t.Shutdown()
 }
diff --git a/work/dispatcher.go b/work/dispatcher.go
index 13502de..1c92082 100644
--- a/work/dispatcher.go
+++ b/work/dispatcher.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"os/exec"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/ChrisMcKenzie/dropship/hook"
@@ -23,21 +24,24 @@ type Dispatcher struct {
 	ticker     *time.Ticker
 	task       *Runner
 	hash       string
+	duration   time.Duration
+	wg         *sync.WaitGroup
 	shutdownCh <-chan struct{}
 }
 
-func NewDispatcher(cfg service.Config, t *Runner, shutdownCh <-chan struct{}) (*Dispatcher, error) {
+func NewDispatcher(cfg service.Config, t *Runner, wg *sync.WaitGroup, shutdownCh <-chan struct{}) (*Dispatcher, error) {
 	w := Dispatcher{
 		config:     cfg,
 		task:       t,
 		shutdownCh: shutdownCh,
+		wg:         wg,
 	}
 
-	dur, err := time.ParseDuration(cfg.CheckInterval)
+	var err error
+	w.duration, err = time.ParseDuration(cfg.CheckInterval)
 	if err != nil {
 		return nil, err
 	}
-	w.ticker = time.NewTicker(dur)
 
 	go w.start()
 
@@ -47,14 +51,14 @@ func NewDispatcher(cfg service.Config, t *Runner, shutdownCh <-chan struct{}) (*
 func (w *Dispatcher) start() {
 	for {
 		select {
-		case <-w.ticker.C:
-			w.task.Do(w)
 		case _, ok := <-w.shutdownCh:
 			if !ok {
 				log.Printf("Shutting down dispatcher for %s", w.config.Name)
-				w.ticker.Stop()
+				w.wg.Done()
 				return
 			}
+		case <-time.After(w.duration):
+			w.task.Do(w)
 		}
 	}
 }