Add option for 'StartupHook' (#35)

* Add ServeWithOptions

This adds support for options to be added to 'Serve' and the app struct.
Options are implemented following the 'functional options' pattern
(https://dave.cheney.net/2014/10/17/functional-options-for-friendly-apis and
https://commandcenter.blogspot.co.uk/2014/01/self-referential-functions-and-design.html).

Future options can be added by creating an exported func that returns a closure
modifying the app struct, like the following:

    func HaltAndCatchFire(literallyCatchFire bool) option {
        return func(a *app) {
            a.haltAndCatchFire = literallyCatchFire
        }
    }

then in user code:

    gracehttp.ServeWithOptions(
        []*http.Server{ &myServer },
        gracehttp.HaltAndCatchFire(true),
    )

* Add 'StartupHook' option

This option attaches a callback to the application. This callback is triggered
directly before the new process is started during a graceful restart. This
allows the old process to release its hold on any resources that the new
process will need.

For example:

    gracehttp.ServeWithOptions(
        []*http.Server{ &myServer },
        gracehttp.StartupHook(func () error {
            // release port that new process will need to start up successfully
            return nil
        }
    )

* Rename 'StartupHook' to 'PreStartProcess'

This better indicates the timing of the callback by using terms already present
in the codebase. As part of the rename, the related constants in the tests were
fixed to follow the naming convention.

gracehttp/http.go

@@ -24,14 +24,17 @@ var (
 	ppid       = os.Getppid()
 )
 
+type option func(*app)
+
 // An app contains one or more servers and associated configuration.
 type app struct {
-	servers   []*http.Server
+	servers         []*http.Server
-	http      *httpdown.HTTP
+	http            *httpdown.HTTP
-	net       *gracenet.Net
+	net             *gracenet.Net
-	listeners []net.Listener
+	listeners       []net.Listener
-	sds       []httpdown.Server
+	sds             []httpdown.Server
-	errors    chan error
+	preStartProcess func() error
+	errors          chan error
 }
 
 func newApp(servers []*http.Server) *app {
@@ -42,6 +45,7 @@ func newApp(servers []*http.Server) *app {
 		listeners: make([]net.Listener, 0, len(servers)),
 		sds:       make([]httpdown.Server, 0, len(servers)),
 
+		preStartProcess: func() error { return nil },
 		// 2x num servers for possible Close or Stop errors + 1 for possible
 		// StartProcess error.
 		errors: make(chan error, 1+(len(servers)*2)),
@@ -108,6 +112,10 @@ func (a *app) signalHandler(wg *sync.WaitGroup) {
 			a.term(wg)
 			return
 		case syscall.SIGUSR2:
+			err := a.preStartProcess()
+			if err != nil {
+				a.errors <- err
+			}
 			// we only return here if there's an error, otherwise the new process
 			// will send us a TERM when it's ready to trigger the actual shutdown.
 			if _, err := a.net.StartProcess(); err != nil {
@@ -117,11 +125,7 @@ func (a *app) signalHandler(wg *sync.WaitGroup) {
 	}
 }
 
-// Serve will serve the given http.Servers and will monitor for signals
+func (a *app) run() error {
-// allowing for graceful termination (SIGTERM) or restart (SIGUSR2).
-func Serve(servers ...*http.Server) error {
-	a := newApp(servers)
-
 	// Acquire Listeners
 	if err := a.listen(); err != nil {
 		return err
@@ -172,6 +176,32 @@ func Serve(servers ...*http.Server) error {
 	}
 }
 
+// ServeWithOptions does the same as Serve, but takes a set of options to
+// configure the app struct.
+func ServeWithOptions(servers []*http.Server, options ...option) error {
+	a := newApp(servers)
+	for _, opt := range options {
+		opt(a)
+	}
+	return a.run()
+}
+
+// Serve will serve the given http.Servers and will monitor for signals
+// allowing for graceful termination (SIGTERM) or restart (SIGUSR2).
+func Serve(servers ...*http.Server) error {
+	a := newApp(servers)
+	return a.run()
+}
+
+// PreStartProcess configures a callback to trigger during graceful restart
+// directly before starting the successor process. This allows the current
+// process to release holds on resources that the new process will need.
+func PreStartProcess(hook func() error) option {
+	return func(a *app) {
+		a.preStartProcess = hook
+	}
+}
+
 // Used for pretty printing addresses.
 func pprintAddr(listeners []net.Listener) []byte {
 	var out bytes.Buffer

gracehttp/http_test.go

@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"os"
 	"os/exec"
+	"strconv"
 	"sync"
 	"syscall"
 	"testing"
@@ -19,6 +20,10 @@ import (
 	"github.com/facebookgo/freeport"
 )
 
+const (
+	testPreStartProcess = iota
+)
+
 // Debug logging.
 var debugLog = flag.Bool("debug", false, "enable debug logging")
 
@@ -39,6 +44,7 @@ type harness struct {
 	newProcess        chan bool      // A bool is sent on start/restart.
 	requestCount      int
 	requestCountMutex sync.Mutex
+	serveOption       int
 }
 
 // Find 3 free ports and setup addresses.
@@ -60,7 +66,7 @@ func (h *harness) setupAddr() {
 // Start a fresh server and wait for pid updates on restart.
 func (h *harness) Start() {
 	h.setupAddr()
-	cmd := exec.Command(os.Args[0], "-http", h.httpAddr, "-https", h.httpsAddr)
+	cmd := exec.Command(os.Args[0], "-http", h.httpAddr, "-https", h.httpsAddr, "-testOption", strconv.Itoa(h.serveOption))
 	stderr, err := cmd.StderrPipe()
 	if err != nil {
 		h.T.Fatal(err)
@@ -209,8 +215,9 @@ func (h *harness) Wait() {
 
 func newHarness(t *testing.T) *harness {
 	return &harness{
-		T:          t,
+		T:           t,
-		newProcess: make(chan bool),
+		newProcess:  make(chan bool),
+		serveOption: -1,
 	}
 }
 
@@ -258,3 +265,49 @@ func TestComplexAgain(t *testing.T) {
 	debug("Waiting")
 	h.Wait()
 }
+
+func TestPreStartProcess(t *testing.T) {
+	t.Parallel()
+	debug("Started TestPreStartProcess")
+	h := newHarness(t)
+	h.serveOption = testPreStartProcess
+	debug("Initial Start")
+	h.Start()
+	debug("Send Request 1")
+	h.SendRequest()
+	debug("Restart 1")
+	h.Restart()
+	debug("Send Request 2")
+	h.SendRequest()
+	debug("Restart 2")
+	h.Restart()
+	debug("Send Request 3")
+	h.SendRequest()
+	debug("Stopping")
+	h.Stop()
+	debug("Waiting")
+	h.Wait()
+}
+
+func TestPreStartProcessAgain(t *testing.T) {
+	t.Parallel()
+	debug("Started TestPreStartProcessAgain")
+	h := newHarness(t)
+	h.serveOption = testPreStartProcess
+	debug("Initial Start")
+	h.Start()
+	debug("Send Request 1")
+	h.SendRequest()
+	debug("Restart 1")
+	h.Restart()
+	debug("Send Request 2")
+	h.SendRequest()
+	debug("Restart 2")
+	h.Restart()
+	debug("Send Request 3")
+	h.SendRequest()
+	debug("Stopping")
+	h.Stop()
+	debug("Waiting")
+	h.Wait()
+}

gracehttp/testbin_test.go

@@ -16,6 +16,8 @@ import (
 	"github.com/facebookgo/grace/gracehttp"
 )
 
+const preStartProcessEnv = "GRACEHTTP_PRE_START_PROCESS"
+
 func TestMain(m *testing.M) {
 	const (
 		testbinKey   = "GRACEHTTP_TEST_BIN"
@@ -101,8 +103,10 @@ func httpsServer(addr string) *http.Server {
 
 func testbinMain() {
 	var httpAddr, httpsAddr string
+	var testOption int
 	flag.StringVar(&httpAddr, "http", ":48560", "http address to bind to")
 	flag.StringVar(&httpsAddr, "https", ":48561", "https address to bind to")
+	flag.IntVar(&testOption, "testOption", -1, "which option to test on ServeWithOptions")
 	flag.Parse()
 
 	// we have self signed certs
@@ -128,12 +132,50 @@ func testbinMain() {
 		}
 	}()
 
-	err := gracehttp.Serve(
+	servers := []*http.Server{
 		&http.Server{Addr: httpAddr, Handler: newHandler()},
 		httpsServer(httpsAddr),
-	)
+	}
-	if err != nil {
+
-		log.Fatalf("Error in gracehttp.Serve: %s", err)
+	if testOption == -1 {
+		err := gracehttp.Serve(servers...)
+		if err != nil {
+			log.Fatalf("Error in gracehttp.Serve: %s", err)
+		}
+	} else {
+		if testOption == testPreStartProcess {
+			switch os.Getenv(preStartProcessEnv) {
+			case "":
+				err := os.Setenv(preStartProcessEnv, "READY")
+				if err != nil {
+					log.Fatalf("testbin (first incarnation) could not set %v to 'ready': %v", preStartProcessEnv, err)
+				}
+			case "FIRED":
+				// all good, reset for next round
+				err := os.Setenv(preStartProcessEnv, "READY")
+				if err != nil {
+					log.Fatalf("testbin (second incarnation) could not reset %v to 'ready': %v", preStartProcessEnv, err)
+				}
+			case "READY":
+				log.Fatalf("failure to update startup hook before new process started")
+			default:
+				log.Fatalf("something strange happened with %v: it ended up as %v, which is not '', 'FIRED', or 'READY'", preStartProcessEnv, os.Getenv(preStartProcessEnv))
+			}
+
+			err := gracehttp.ServeWithOptions(
+				servers,
+				gracehttp.PreStartProcess(func() error {
+					err := os.Setenv(preStartProcessEnv, "FIRED")
+					if err != nil {
+						log.Fatalf("startup hook could not set %v to 'fired': %v", preStartProcessEnv, err)
+					}
+					return nil
+				}),
+			)
+			if err != nil {
+				log.Fatalf("Error in gracehttp.Serve: %s", err)
+			}
+		}
 	}
 }