From 4afe952a37a495ae4ac0c1d4ce5f66e91058d149 Mon Sep 17 00:00:00 2001 From: Ben Tyler Date: Sat, 18 Feb 2017 23:52:39 +0100 Subject: [PATCH] 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 | 52 ++++++++++++++++++++++++++-------- gracehttp/http_test.go | 59 +++++++++++++++++++++++++++++++++++++-- gracehttp/testbin_test.go | 50 ++++++++++++++++++++++++++++++--- 3 files changed, 143 insertions(+), 18 deletions(-) diff --git a/gracehttp/http.go b/gracehttp/http.go index fa3ac88..eefe340 100644 --- a/gracehttp/http.go +++ b/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 - http *httpdown.HTTP - net *gracenet.Net - listeners []net.Listener - sds []httpdown.Server - errors chan error + servers []*http.Server + http *httpdown.HTTP + net *gracenet.Net + listeners []net.Listener + sds []httpdown.Server + 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 -// allowing for graceful termination (SIGTERM) or restart (SIGUSR2). -func Serve(servers ...*http.Server) error { - a := newApp(servers) - +func (a *app) run() error { // 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 diff --git a/gracehttp/http_test.go b/gracehttp/http_test.go index 4c9a70b..fa04a1c 100644 --- a/gracehttp/http_test.go +++ b/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, - newProcess: make(chan bool), + T: t, + 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() +} diff --git a/gracehttp/testbin_test.go b/gracehttp/testbin_test.go index 1b2b40b..4fa7dbd 100644 --- a/gracehttp/testbin_test.go +++ b/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) + } + } } }