Merge pull request #11 from Fred07/feature/handle-signals

Handle signals

README.md

@@ -12,7 +12,7 @@ Gron provides a clear syntax for writing and deploying cron jobs.
 - Customizable Job Type.
 - Customizable Schedule.
 
-## Different to origin gron
+## Different to origin
 
 Most features and interfaces are as same as origin repository, but also introduces one new method to Cron, which `StopAfterJobDone` make cron has ability to hold the process, stop creating new child goroutine, and ensure sub-job is finished before main goroutine close. This benefit to the scenarios like handling `os.Signal` including `SIGINT`, `SIGTERM` (which handler does not including in library). After interrupt signal handling, stopping cron gracefully via `StopAfterJobDone`, prevent bundle of processing job stop inappropriately.
 

cron.go

@@ -1,6 +1,8 @@
 package gron
 
 import (
+	"os"
+	"os/signal"
 	"sort"
 	"sync"
 	"time"
@@ -59,16 +61,28 @@ type Cron struct {
 	add     chan *Entry
 	stop    chan struct{}
 	wg      sync.WaitGroup
+	sigChan chan os.Signal
 }
 
 // New instantiates new Cron instant c.
 func New() *Cron {
 	return &Cron{
-		stop: make(chan struct{}),
-		add:  make(chan *Entry),
+		stop:    make(chan struct{}),
+		add:     make(chan *Entry),
+		sigChan: make(chan os.Signal),
 	}
 }
 
+// HandleSignals to deside what signal should be aware
+func (c *Cron) HandleSignals(signals ...os.Signal) {
+	if len(signals) == 0 {
+		return
+	}
+
+	// sigChan would not receive signals without register
+	signal.Notify(c.sigChan, signals...)
+}
+
 // Start signals cron instant c to get up and running.
 func (c *Cron) Start() {
 	c.running = true
@@ -108,14 +122,19 @@ func (c *Cron) Stop() {
 	c.stop <- struct{}{}
 }
 
-// StopAfterJobDone halts cron after running jobs are finished
-func (c *Cron) StopAfterJobDone() {
+// GracefullyStop halts cron after running jobs are finished
+func (c *Cron) GracefullyStop() {
+
+	c.stopAndWait()
+	c.stop <- struct{}{}
+}
+
+func (c *Cron) stopAndWait() {
 	if !c.running {
 		return
 	}
 	c.running = false
 	c.wg.Wait()
-	c.stop <- struct{}{}
 }
 
 var after = time.After
@@ -154,6 +173,9 @@ func (c *Cron) run() {
 				c.wg.Add(1)
 				go entry.Job.Run(&c.wg)
 			}
+		case <-c.sigChan:
+			c.stopAndWait()
+			return
 		case e := <-c.add:
 			e.Next = e.Schedule.Next(time.Now())
 			c.entries = append(c.entries, e)

cron_test.go

@@ -6,6 +6,7 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"syscall"
 	"testing"
 	"time"
 
@@ -105,36 +106,62 @@ func TestJobsDontRunAfterStop(t *testing.T) {
 		// no job has run
 	case <-wait(wg):
 		t.FailNow()
-
 	}
 }
 
-func TestJobFinishedBeforeStop(t *testing.T) {
+func TestGracefullyStop(t *testing.T) {
 	failCh := make(chan bool, 1)
 	successCh := make(chan bool, 1)
 
-	go func(successCh chan bool, failCh chan bool) {
-	loop:
-		for {
-			select {
-			case <-successCh:
-				break loop
-			case <-failCh:
-				t.FailNow()
-			default:
-			}
+	go func() {
+		cron := New()
+		cron.Start()
+		cron.AddFunc(Every(2*time.Second), func() {
+			time.Sleep(time.Duration(1) * time.Second) // make job running while stopping
+			successCh <- true                          // success signal
+		})
+		time.Sleep(time.Duration(2) * time.Second) // prevent stop too fast
+		cron.GracefullyStop()                      // call stop
+		failCh <- true                             // fail signal
+	}()
+
+loop:
+	for {
+		select {
+		case <-successCh:
+			break loop
+		case <-failCh:
+			t.FailNow()
+		default:
 		}
-	}(successCh, failCh)
+	}
+}
 
+func TestInterruptSingal(t *testing.T) {
+	successCh := make(chan bool, 1)
 	cron := New()
-	cron.Start()
-	cron.AddFunc(Every(2*time.Second), func() {
-		time.Sleep(time.Duration(1) * time.Second) // make job running while stopping
-		successCh <- true                          // success signal
-	})
-	time.Sleep(time.Duration(2) * time.Second) // prevent stop too fast
-	cron.StopAfterJobDone()                    // call stop
-	failCh <- true                             // fail signal
+
+	go func() {
+		pid := syscall.Getpid()
+		cron.Start()
+		cron.HandleSignals(syscall.SIGINT, syscall.SIGTERM)
+		cron.AddFunc(Every(1*time.Second), func() {
+			time.Sleep(time.Duration(1) * time.Second)
+			successCh <- true
+		})
+		time.Sleep(time.Duration(1) * time.Second)
+		syscall.Kill(pid, syscall.SIGINT)
+	}()
+
+loop:
+	for {
+		select {
+		case <-successCh:
+			break loop
+		case <-time.After(time.Duration(5) * time.Second):
+			t.FailNow()
+		}
+	}
 
 }