use RunnerFunc, ShutdownFunc and implement Killed (#1)

* use RunnerFunc, ShutdownFunc and implement Killed

* update wording of docs

CHANGELOG.md

@@ -5,11 +5,20 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+[Unreleased]: https://github.com/mec07/rununtil/compare/v0.1.0...HEAD
+
+## [0.1.0] - 2019-06-12
+[0.1.0]: https://github.com/mec07/rununtil/compare/v0.0.1...v0.1.0
+### Changed
+- There are now RunnerFunc and ShutdownFunc function types to clarify the usage of this library (backwards incompatible change)
 
-## [0.0.1] - 2019-05-21
 ### Added
-- initial commit of rununtil library
+- Implemented the rununtil.Killed method, which allows you to test a function that uses rununtil.KillSignal
+
 
-[Unreleased]: https://github.com/mec07/rununtil/compare/v0.0.1...HEAD
+## [0.0.1] - 2019-06-05
 [0.0.1]: https://github.com/mec07/rununtil/releases/tag/v0.0.1
+### Added
+- initial commit of rununtil library
+
 

go.mod

@@ -1,3 +1,5 @@
 module github.com/mec07/rununtil
 
 go 1.12
+
+require github.com/pkg/errors v0.8.1

go.sum

@@ -0,0 +1,2 @@
+github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
+github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

rununtil.go

@@ -2,21 +2,21 @@
 
 Usage
 
-The main usage of rununtil is to run a webserver or other function until a kill signal has been received.
+The main usage of rununtil is to run a webserver or other function until a SIGINT or SIGTERM signal has been received.
 The runner function can do some setup but it should not run indefinitely, instead it should start go routines which can run in the background.
 The runner function should return a graceful shutdown function that will be called once the signal has been received.
 For example:
-	func Runner() func() {
+	func Runner() rununtil.ShutdownFunc {
 		r := chi.NewRouter()
 		r.Get("/healthz", healthzHandler)
 		httpServer := &http.Server{Addr: ":8080", Handler: r}
 		go runHTTPServer(httpServer)
 
-		return func() {
+		return rununtil.ShutdownFunc(func() {
 			if err := httpServer.Shutdown(context.Background()); err != nil {
 				log.Error().Err(err).Msg("error occurred while shutting down http server")
 			}
-		}
+		})
 	}
 
 	func runHTTPServer(srv *http.Server) {
@@ -30,19 +30,19 @@ For example:
 	}
 
 A nice pattern is to create a function that takes in the various depencies required, for example, a logger (but could be anything, e.g. configs, database, etc.), and returns a runner function:
-	func NewRunner(log *zerolog.Logger) func() func() {
-		return func() func() {
+	func NewRunner(log *zerolog.Logger) rununtil.RunnerFunc {
+		return rununtil.RunnerFunc(func() rununtil.ShutdownFunc {
 			r := chi.NewRouter()
 			r.Get("/healthz", healthzHandler)
 			httpServer := &http.Server{Addr: ":8080", Handler: r}
 			go runHTTPServer(httpServer, log)
 
-			return func() {
+			return rununtil.ShutdownFunc(func() {
 				if err := httpServer.Shutdown(context.Background()); err != nil {
 					log.Error().Err(err).Msg("error occurred while shutting down http server")
 				}
-			}
-		}
+			})
+		})
 	}
 
 	func runHTTPServer(srv *http.Server, log *zerolog.Logger) {
@@ -58,30 +58,82 @@ A nice pattern is to create a function that takes in the various depencies requi
 		}
 		rununtil.KillSignal(NewRunner(logger))
 	}
+
+It is of course possible to specify which signals you would like to use to kill your application using the `Signals` function, for example:
+	rununtil.Signals(NewRunner(logger), syscall.SIGKILL, syscall.SIGHUP, syscall.SIGINT)
+
+For testing purposes you may want to run your main function, which is using `rununtil.KillSignal`, and send it a kill signal when you're done with your tests. To aid with this you can use:
+	kill := rununtil.Killed(main)
+
+where `kill` is a function that sends a kill signal to the main function when executed (its type is context.CancelFunc).
 */
 package rununtil
 
 import (
+	"context"
+	"fmt"
 	"os"
 	"os/signal"
 	"syscall"
+
+	"github.com/pkg/errors"
 )
 
+// ShutdownFunc is a function that should be returned by a RunnerFunc which
+// gracefully shuts down whatever is being run.
+type ShutdownFunc func()
+
+// RunnerFunc is a function that sets off the worker go routines and returns
+// a function which can shutdown those worker go routines.
+type RunnerFunc func() ShutdownFunc
+
 // KillSignal runs the provided runner function until it receives a kill signal,
 // SIGINT or SIGTERM, at which point it executes the graceful shutdown function.
-func KillSignal(runner func() func()) {
+func KillSignal(runner RunnerFunc) {
 	Signals(runner, syscall.SIGINT, syscall.SIGTERM)
 }
 
 // Signals runs the provided runner function until the specified signals have
 // been recieved.
-func Signals(runner func() func(), signals ...os.Signal) {
+func Signals(runner RunnerFunc, signals ...os.Signal) {
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, signals...)
 
-	gracefulShutdown := runner()
-	defer gracefulShutdown()
+	shutdown := runner()
+	defer shutdown()
 
 	// Wait for a kill signal
 	<-c
 }
+
+// Killed is used for testing a function that is using rununtil.KillSignal.
+// It runs the function provided and sends a SIGINT signal to kill it when
+// the returned context.CancelFunc is executed. A sample usage of this could be:
+//	kill := rununtil.Killed(main)
+//	... do some stuff, e.g. send some requests to the webserver ...
+//	kill()
+//
+// where main is a function that is using rununtil.KillSignal.
+func Killed(main func()) context.CancelFunc {
+	ctx, cancel := context.WithCancel(context.Background())
+	go runMain(ctx, main)
+
+	return cancel
+}
+
+func runMain(ctx context.Context, main func()) {
+	p, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		fmt.Printf("ERROR: %+v\n", errors.Wrap(err, "trying to get PID"))
+	}
+	go killMainWhenDone(ctx, p)
+	main()
+}
+
+func killMainWhenDone(ctx context.Context, p *os.Process) {
+	<-ctx.Done()
+
+	if err := p.Signal(syscall.SIGINT); err != nil {
+		fmt.Printf("ERROR: %+v\n", errors.Wrap(err, "trying to kill main"))
+	}
+}

rununtil_test.go

@@ -17,8 +17,18 @@ func helperSendSignal(t *testing.T, p *os.Process, sent *bool, signal os.Signal,
 	*sent = true
 }
 
-func helperFakeRunner() func() {
-	return func() {}
+func helperMakeFakeRunner(hasBeenShutdown *bool) rununtil.RunnerFunc {
+	return rununtil.RunnerFunc(func() rununtil.ShutdownFunc {
+		return rununtil.ShutdownFunc(func() {
+			*hasBeenShutdown = true
+		})
+	})
+}
+
+func helperMakeMain(hasBeenKilled *bool) func() {
+	return func() {
+		rununtil.KillSignal(helperMakeFakeRunner(hasBeenKilled))
+	}
 }
 
 func TestRunUntilKillSignal(t *testing.T) {
@@ -38,16 +48,32 @@ func TestRunUntilKillSignal(t *testing.T) {
 	for _, test := range table {
 		t.Run(test.name, func(t *testing.T) {
 			var sentSignal bool
+			var hasBeenShutdown bool
 			p, err := os.FindProcess(os.Getpid())
 			if err != nil {
 				t.Fatalf("Unexpected error when finding process: %v", err)
 			}
 
 			go helperSendSignal(t, p, &sentSignal, test.signal, 1*time.Millisecond)
-			rununtil.KillSignal(helperFakeRunner)
+			rununtil.KillSignal(helperMakeFakeRunner(&hasBeenShutdown))
 			if !sentSignal {
 				t.Fatal("expected signal to have been sent")
 			}
+			if !hasBeenShutdown {
+				t.Fatal("expected the shutdown function to have been called")
+			}
 		})
 	}
 }
+
+func TestRunUntilKilled(t *testing.T) {
+	var hasBeenKilled bool
+	kill := rununtil.Killed(helperMakeMain(&hasBeenKilled))
+	kill()
+
+	// yield control back to scheduler so that killing can actually happen
+	time.Sleep(time.Millisecond)
+	if !hasBeenKilled {
+		t.Fatal("expected main to have been killed")
+	}
+}