Add an HTTP wrapper that does graceful shutdown.
diff --git a/accept_test.go b/accept_test.go
index 2a80b21..f9cfe75 100644
--- a/accept_test.go
+++ b/accept_test.go
@@ -25,6 +25,12 @@
[]string{"text/plain", "text/xml"})).Should(Equal("text/plain"))
})
+ It("No header, two choices 2", func() {
+ Expect(SelectMediaType(
+ makeRequest(""),
+ []string{"application/json", "text/plain"})).Should(Equal("application/json"))
+ })
+
It("One Header, two choices", func() {
Expect(SelectMediaType(
makeRequest("application/json"),
diff --git a/httpstop.go b/httpstop.go
new file mode 100644
index 0000000..49f014d
--- /dev/null
+++ b/httpstop.go
@@ -0,0 +1 @@
+package goscaffold
diff --git a/scaffold.go b/scaffold.go
new file mode 100644
index 0000000..0bfdc51
--- /dev/null
+++ b/scaffold.go
@@ -0,0 +1,138 @@
+package goscaffold
+
+import (
+ "encoding/json"
+ "net"
+ "net/http"
+ "time"
+)
+
+const (
+ // DefaultGraceTimeout is the default amount of time to wait for a request to complete
+ DefaultGraceTimeout = 30 * time.Second
+)
+
+/*
+An HTTPScaffold provides a set of features on top of a standard HTTP
+listener. It includes an HTTP handler that may be plugged in to any
+standard Go HTTP server. It is intended to be placed before any other
+handlers.
+*/
+type HTTPScaffold struct {
+ insecurePort int
+ tracker *requestTracker
+ insecureListener net.Listener
+}
+
+/*
+CreateHTTPScaffold makes a new scaffold. The default scaffold will
+do nothing.
+*/
+func CreateHTTPScaffold() *HTTPScaffold {
+ return &HTTPScaffold{}
+}
+
+/*
+SetInsecurePort sets the port number to listen on in regular "HTTP" mode.
+It may be set to zero, which indicates to listen on an ephemeral port.
+It must be called before "listen".
+*/
+func (s *HTTPScaffold) SetInsecurePort(ip int) {
+ s.insecurePort = ip
+}
+
+/*
+InsecureAddress returns the actual address (including the port if an
+ephemeral port was used) where we are listening. It must only be
+called after "Listen."
+*/
+func (s *HTTPScaffold) InsecureAddress() string {
+ return s.insecureListener.Addr().String()
+}
+
+/*
+Open opens up the port that was created when the scaffold was set up.
+*/
+func (s *HTTPScaffold) Open() error {
+ s.tracker = startRequestTracker(DefaultGraceTimeout)
+
+ il, err := net.ListenTCP("tcp", &net.TCPAddr{
+ Port: s.insecurePort,
+ })
+ if err != nil {
+ return err
+ }
+ s.insecureListener = il
+ return nil
+}
+
+/*
+Listen should be called instead of using the standard "http" and "net"
+libraries. It will open a port (or ports) and begin listening for
+HTTP traffic. It will block until the server is shut down by
+the various methods in this class.
+It will use the graceful shutdown logic to ensure that once marked down,
+the server will not exit until all the requests have completed,
+or until the shutdown timeout has expired.
+Like http.Serve, this function will block until we are done serving HTTP.
+If "SetInsecurePort" or "SetSecurePort" were not set, then it will listen on
+a dynamic port.
+Listen will block until the server is shutdown using "Shutdown" or one of
+the other shutdown mechanisms. It must not be called until after "Open"
+has been called.
+*/
+func (s *HTTPScaffold) Listen(baseHandler http.Handler) {
+ handler := &httpHandler{
+ s: s,
+ handler: baseHandler,
+ }
+ go http.Serve(s.insecureListener, handler)
+ <-s.tracker.C
+ s.insecureListener.Close()
+}
+
+/*
+Shutdown indicates that the server should stop handling incoming requests
+and exit from the "Serve" call. This may be called automatically by
+calling "CatchSignals," or automatically using this call.
+*/
+func (s *HTTPScaffold) Shutdown(reason error) {
+ s.tracker.shutdown(reason)
+}
+
+/*
+CatchSignals directs the scaffold to catch SIGINT and SIGTERM (the signals
+sent by "Control-C" and "kill" by default) to trigger the markdown
+logic. Using this logic, when these signals are caught, the server will
+catch
+*/
+func (s *HTTPScaffold) CatchSignals() {
+}
+
+type httpHandler struct {
+ s *HTTPScaffold
+ handler http.Handler
+}
+
+func (h *httpHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
+ startErr := h.s.tracker.start()
+ if startErr == nil {
+ h.handler.ServeHTTP(resp, req)
+ h.s.tracker.end()
+ } else {
+ mt := SelectMediaType(req, []string{"text/plain", "application/json"})
+ resp.Header().Set("Content-Type", mt)
+ resp.WriteHeader(http.StatusServiceUnavailable)
+ switch mt {
+ case "application/json":
+ re := map[string]string{
+ "error": "Stopping",
+ "message": startErr.Error(),
+ }
+ buf, _ := json.Marshal(&re)
+ resp.Write(buf)
+ default:
+ resp.Write([]byte(startErr.Error()))
+ }
+ }
+}
diff --git a/scaffold_test.go b/scaffold_test.go
new file mode 100644
index 0000000..9422a94
--- /dev/null
+++ b/scaffold_test.go
@@ -0,0 +1,109 @@
+package goscaffold
+
+import (
+ "errors"
+ "fmt"
+ "net/http"
+ "time"
+
+ . "github.com/onsi/ginkgo"
+ . "github.com/onsi/gomega"
+)
+
+var _ = Describe("Scaffold Tests", func() {
+ It("Validate framework", func() {
+ s := CreateHTTPScaffold()
+ stopChan := make(chan bool)
+ err := s.Open()
+ Expect(err).Should(Succeed())
+
+ go func() {
+ fmt.Fprintf(GinkgoWriter, "Gonna listen on %s\n", s.InsecureAddress())
+ s.Listen(&testHandler{})
+ fmt.Fprintf(GinkgoWriter, "Done listening\n")
+ stopChan <- true
+ }()
+
+ Eventually(func() bool {
+ return testGet(s, "")
+ }, 5*time.Second).Should(BeTrue())
+ resp, err := http.Get(fmt.Sprintf("http://%s", s.InsecureAddress()))
+ Expect(err).Should(Succeed())
+ Expect(resp.StatusCode).Should(Equal(200))
+ s.Shutdown(errors.New("Validate"))
+ Eventually(stopChan).Should(Receive(BeTrue()))
+ })
+
+ It("Shutdown", func() {
+ s := CreateHTTPScaffold()
+ stopChan := make(chan bool)
+ err := s.Open()
+ Expect(err).Should(Succeed())
+
+ go func() {
+ s.Listen(&testHandler{})
+ stopChan <- true
+ }()
+
+ go func() {
+ resp2, err2 := http.Get(fmt.Sprintf("http://%s?delay=1s", s.InsecureAddress()))
+ Expect(err2).Should(Succeed())
+ Expect(resp2.StatusCode).Should(Equal(200))
+ }()
+
+ // Just make sure server is listening
+ Eventually(func() bool {
+ return testGet(s, "")
+ }, 5*time.Second).Should(BeTrue())
+
+ // Previous call prevents server from exiting
+ Consistently(stopChan, 250*time.Millisecond).ShouldNot(Receive())
+
+ // Tell the server to try and exit
+ s.Shutdown(errors.New("Stop"))
+ // Should take one second -- in the meantime, calls should fail with 503
+ resp, err := http.Get(fmt.Sprintf("http://%s?", s.InsecureAddress()))
+ Expect(err).Should(Succeed())
+ Expect(resp.StatusCode).Should(Equal(503))
+ // But in less than two seconds, server should be down
+ Eventually(stopChan, 2*time.Second).Should(Receive(BeTrue()))
+ // Calls should now fail
+ Eventually(func() bool {
+ return testGet(s, "")
+ }, time.Second).Should(BeFalse())
+ })
+})
+
+func testGet(s *HTTPScaffold, path string) bool {
+ resp, err := http.Get(fmt.Sprintf("http://%s", s.InsecureAddress()))
+ if err != nil {
+ fmt.Fprintf(GinkgoWriter, "Get %s = %s\n", path, err)
+ return false
+ }
+ if resp.StatusCode != 200 {
+ fmt.Fprintf(GinkgoWriter, "Get %s = %d\n", path, resp.StatusCode)
+ return false
+ }
+ return true
+}
+
+type testHandler struct {
+}
+
+func (h *testHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
+ var err error
+ var delayTime time.Duration
+
+ delayStr := req.URL.Query().Get("delay")
+ if delayStr != "" {
+ delayTime, err = time.ParseDuration(delayStr)
+ if err != nil {
+ resp.WriteHeader(http.StatusBadRequest)
+ return
+ }
+ }
+
+ if delayTime > 0 {
+ time.Sleep(delayTime)
+ }
+}
diff --git a/tracker.go b/tracker.go
new file mode 100644
index 0000000..0b624ea
--- /dev/null
+++ b/tracker.go
@@ -0,0 +1,130 @@
+package goscaffold
+
+import (
+ "math"
+ "sync/atomic"
+ "time"
+)
+
+/*
+values for the command channel.
+*/
+const (
+ startRequest = iota
+ endRequest
+ shutdown
+)
+
+/*
+The requestTracker keeps track of HTTP requests. In normal operations it
+just counts. Once the server has been marked for shutdown, however, it
+counts down to zero and returns a shutdown indication when that
+happens.
+*/
+type requestTracker struct {
+ // A value will be delivered to this channel when the server can stop.
+ // If "shutdown" is never called then this will never happen.
+ C chan error
+ shutdownWait time.Duration
+ shuttingDown int32
+ shutdownReason *atomic.Value
+ commandChan chan int
+}
+
+/*
+startRequestTracker creates a new tracker. shutdownWait defines the
+maximum amount of time that we should wait for shutdown in case some
+do not complete in a timely way.
+*/
+func startRequestTracker(shutdownWait time.Duration) *requestTracker {
+ rt := &requestTracker{
+ C: make(chan error, 1),
+ commandChan: make(chan int, 100),
+ shutdownWait: shutdownWait,
+ shutdownReason: &atomic.Value{},
+ }
+ go rt.trackerLoop()
+ return rt
+}
+
+/*
+start indicates that a request started. It returns true if the request
+should proceed, and false if the request should fail because the server is
+shutting down.
+*/
+func (t *requestTracker) start() error {
+ sd := atomic.LoadInt32(&t.shuttingDown)
+ if sd != 0 {
+ reason := t.shutdownReason.Load().(*error)
+ if reason == nil {
+ return nil
+ }
+ return *reason
+ }
+ t.commandChan <- startRequest
+ return nil
+}
+
+/*
+end indicates that a request ended. In order for this thing to work, the
+caller needs to ensure that start and end are always paired.
+*/
+func (t *requestTracker) end() {
+ t.commandChan <- endRequest
+}
+
+/*
+shutdown indicates that the tracker should start counting down until
+the number of running requests reaches zero. The "reason" will be returned
+as the result of the "start" call.
+*/
+func (t *requestTracker) shutdown(reason error) {
+ t.shutdownReason.Store(&reason)
+ t.commandChan <- shutdown
+}
+
+func (t *requestTracker) sendStop(sent bool) bool {
+ if !sent {
+ reason := t.shutdownReason.Load().(*error)
+ if reason == nil {
+ return false
+ }
+ t.C <- *reason
+ }
+ return true
+}
+
+/*
+trackerLoop runs all day and manages stuff.
+*/
+func (t *requestTracker) trackerLoop() {
+ activeRequests := 0
+ stopping := false
+ sentStop := false
+ graceTimer := time.NewTimer(time.Duration(math.MaxInt64))
+
+ for !sentStop {
+ select {
+ case cmd := <-t.commandChan:
+ switch cmd {
+ case startRequest:
+ activeRequests++
+ case endRequest:
+ activeRequests--
+ if stopping && activeRequests == 0 {
+ sentStop = t.sendStop(sentStop)
+ }
+ case shutdown:
+ stopping = true
+ atomic.StoreInt32(&t.shuttingDown, 1)
+ if activeRequests <= 0 {
+ sentStop = t.sendStop(sentStop)
+ } else {
+ graceTimer.Reset(t.shutdownWait)
+ }
+ }
+ case <-graceTimer.C:
+ sentStop = t.sendStop(sentStop)
+ }
+ }
+}
diff --git a/tracker_test.go b/tracker_test.go
new file mode 100644
index 0000000..fc3ad0e
--- /dev/null
+++ b/tracker_test.go
@@ -0,0 +1,34 @@
+package goscaffold
+
+import (
+ "errors"
+ "time"
+
+ . "github.com/onsi/ginkgo"
+ . "github.com/onsi/gomega"
+)
+
+var _ = Describe("Tracker tests", func() {
+ It("Basic tracker", func() {
+ t := startRequestTracker(10 * time.Second)
+ t.start()
+ Consistently(t.C, 250*time.Millisecond).ShouldNot(Receive())
+ t.shutdown(errors.New("Basic"))
+ Consistently(t.C, 250*time.Millisecond).ShouldNot(Receive())
+ t.end()
+ Eventually(t.C).Should(Receive(MatchError("Basic")))
+ })
+
+ It("Tracker stop idle", func() {
+ t := startRequestTracker(10 * time.Second)
+ t.shutdown(errors.New("Stop"))
+ Eventually(t.C).Should(Receive(MatchError("Stop")))
+ })
+
+ It("Tracker grace timeout", func() {
+ t := startRequestTracker(time.Second)
+ t.start()
+ t.shutdown(errors.New("Stop"))
+ Eventually(t.C, 2*time.Second).Should(Receive(MatchError("Stop")))
+ })
+})
