Merge pull request #31 from krallin/context-for-exit

Use a context to exit after being signalled

cron/cron.go

@@ -121,7 +121,7 @@ func monitorJob(ctx context.Context, expression crontab.Expression, t0 time.Time
 	}
 }
 
-func StartJob(wg *sync.WaitGroup, cronCtx *crontab.Context, job *crontab.Job, exitChan chan interface{}, cronLogger *logrus.Entry) {
+func StartJob(wg *sync.WaitGroup, cronCtx *crontab.Context, job *crontab.Job, exitCtx context.Context, cronLogger *logrus.Entry) {
 	wg.Add(1)
 
 	go func() {
@@ -144,7 +144,7 @@ func StartJob(wg *sync.WaitGroup, cronCtx *crontab.Context, job *crontab.Job, ex
 			}
 
 			select {
-			case <-exitChan:
+			case <-exitCtx.Done():
 				cronLogger.Debug("shutting down")
 				return
 			case <-time.After(delay):

cron/cron_test.go

@@ -1,6 +1,7 @@
 package cron
 
 import (
+	"context"
 	"fmt"
 	"github.com/aptible/supercronic/crontab"
 	"github.com/sirupsen/logrus"
@@ -189,8 +190,10 @@ func TestStartJobExitsOnRequest(t *testing.T) {
 	logger, _ := newTestLogger()
 
 	var wg sync.WaitGroup
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
 
-	StartJob(&wg, &basicContext, &job, exitChan, logger)
+	StartJob(&wg, &basicContext, &job, ctx, logger)
 
 	wg.Wait()
 }
@@ -205,13 +208,12 @@ func TestStartJobRunsJob(t *testing.T) {
 		Position: 1,
 	}
 
-	exitChan := make(chan interface{}, 1)
-
 	var wg sync.WaitGroup
+	ctx, cancel := context.WithCancel(context.Background())
 
 	logger, channel := newTestLogger()
 
-	StartJob(&wg, &basicContext, &job, exitChan, logger)
+	StartJob(&wg, &basicContext, &job, ctx, logger)
 
 	select {
 	case entry := <-channel:
@@ -255,6 +257,6 @@ func TestStartJobRunsJob(t *testing.T) {
 		t.Fatalf("timed out waiting for second success")
 	}
 
-	exitChan <- nil
+	cancel()
 	wg.Wait()
 }

integration/exit.crontab

@@ -0,0 +1 @@
+* * * * * * * echo "$MSG_START" && sleep 2 && echo "$MSG_DONE"

integration/test.bats

@@ -5,6 +5,26 @@ function run_supercronic() {
     "${BATS_TEST_DIRNAME}/../supercronic" ${SUPERCRONIC_ARGS:-} "$crontab" 2>&1
 }
 
+setup () {
+  WORK_DIR="$(mktemp -d)"
+  export WORK_DIR
+}
+
+teardown() {
+  rm -r "$WORK_DIR"
+}
+
+wait_for() {
+  for i in $(seq 0 50); do
+    if "$@" > /dev/null 2>&1; then
+      return 0
+    fi
+    sleep 0.1
+  done
+
+  return 1
+}
+
 @test "it starts" {
   run_supercronic "${BATS_TEST_DIRNAME}/noop.crontab"
 }
@@ -42,3 +62,18 @@ function run_supercronic() {
 @test "it supports JSON logging " {
   SUPERCRONIC_ARGS="-json" run_supercronic "${BATS_TEST_DIRNAME}/noop.crontab" | grep -iE "^{"
 }
+
+@test "it waits for jobs to exit before terminating" {
+  ready="will start"
+  canary="all done"
+
+  out="${WORK_DIR}/out"
+
+  MSG_START="$ready" MSG_DONE="$canary" \
+    "${BATS_TEST_DIRNAME}/../supercronic" "${BATS_TEST_DIRNAME}/exit.crontab" >"$out" 2>&1 &
+
+  wait_for grep "$ready" "$out"
+  kill -TERM "$!"
+
+  wait_for grep "$canary" "$out"
+}

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"github.com/aptible/supercronic/cron"
@@ -48,22 +49,17 @@ func main() {
 		return
 	}
 
-	var (
-		wg        sync.WaitGroup
-		exitChans []chan interface{}
-	)
+	var wg sync.WaitGroup
+	exitCtx, notifyExit := context.WithCancel(context.Background())
 
 	for _, job := range tab.Jobs {
-		exitChan := make(chan interface{}, 1)
-		exitChans = append(exitChans, exitChan)
-
 		cronLogger := logrus.WithFields(logrus.Fields{
 			"job.schedule": job.Schedule,
 			"job.command":  job.Command,
 			"job.position": job.Position,
 		})
 
-		cron.StartJob(&wg, tab.Context, job, exitChan, cronLogger)
+		cron.StartJob(&wg, tab.Context, job, exitCtx, cronLogger)
 	}
 
 	termChan := make(chan os.Signal, 1)
@@ -72,9 +68,7 @@ func main() {
 	termSig := <-termChan
 
 	logrus.Infof("received %s, shutting down", termSig)
-	for _, c := range exitChans {
-		c <- true
-	}
+	notifyExit()
 
 	logrus.Info("waiting for jobs to finish")
 	wg.Wait()