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) + } + } } }