From f96e784096f17a1a29d5834309b2fdf7825bd6bb Mon Sep 17 00:00:00 2001
From: Arne Bahlo <hallo@arne.me>
Date: Wed, 2 Sep 2015 10:31:06 +0200
Subject: [PATCH] Add GracefulServer

---
 http.go      | 79 +++++++++++++++++++++++++++++++++++++++++++
 http_test.go | 95 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 174 insertions(+)

diff --git a/http.go b/http.go
index f8615a2..af746dd 100644
--- a/http.go
+++ b/http.go
@@ -1,8 +1,14 @@
 package abutil
 
 import (
+	"net"
 	"net/http"
+	"strconv"
 	"strings"
+	"sync"
+	"time"
+
+	"github.com/tylerb/graceful"
 )
 
 // RemoteIP tries to get the remote ip and returns it or ""
@@ -24,3 +30,76 @@ func RemoteIP(r *http.Request) string {
 
 	return a
 }
+
+// GracefulServer is a wrapper around graceful.Server from
+// github.com/tylerb/graceful, but adds a running variable and a mutex for
+// controlling the access to it.
+type GracefulServer struct {
+	server *graceful.Server
+
+	// stopped determines if the server is stopped
+	stopped bool
+
+	// locker controls the access to running
+	locker sync.Locker
+}
+
+// NewGracefulServer creates a new GracefulServer with the given handler,
+// which listens on the given port. When Stop() ist called, it waits until
+// the timeout is finished or all connections are closed (whatever comes first)
+func NewGracefulServer(p int, h http.Handler, t time.Duration) *GracefulServer {
+	var m sync.Mutex
+	return &GracefulServer{
+		server: &graceful.Server{
+			Server: &http.Server{
+				Addr:    ":" + strconv.Itoa(p),
+				Handler: h,
+			},
+			NoSignalHandling: true,
+			Timeout:          t,
+		},
+		stopped: true,
+		locker:  &m,
+	}
+}
+
+// Stopped returns if the server is running
+func (g *GracefulServer) Stopped() bool {
+	g.locker.Lock()
+	defer g.locker.Unlock()
+
+	return g.stopped
+}
+
+func (g *GracefulServer) setStopped(r bool) {
+	g.locker.Lock()
+	g.stopped = r
+	g.locker.Unlock()
+}
+
+// Stop stops the server (this may last up to the server timeout)
+func (g *GracefulServer) Stop() {
+	g.setStopped(true)
+
+	g.server.Stop(g.server.Timeout)
+}
+
+// Serve is equivalent to http.Server.Serve, but the server is stoppable
+func (g *GracefulServer) Serve(l net.Listener) error {
+	g.setStopped(false)
+	return g.server.Serve(l)
+}
+
+// ListenAndServe is equivalent to http.Server.Serve, but the server is
+// stoppable
+func (g *GracefulServer) ListenAndServe() error {
+	g.setStopped(false)
+	return g.server.ListenAndServe()
+}
+
+// ListenAndServeTLS is equivalent to http.Server.Serve, but the server is
+// stoppable
+func (g *GracefulServer) ListenAndServeTLS(cf, kf string) error {
+	g.setStopped(false)
+	return g.server.ListenAndServeTLS(cf, kf)
+}
diff --git a/http_test.go b/http_test.go
index 3161217..ce74da6 100644
--- a/http_test.go
+++ b/http_test.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 )
 
 func mockRequestContext(t *testing.T, fn func(*http.Request)) {
@@ -91,3 +92,97 @@ func ExampleRemoteIP() {
 
 	// Output: New request from 123.456.7.8
 }
+
+func gracefulServerContext(t *testing.T, fn func(*GracefulServer)) {
+	p := 1337
+	h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("Foobar"))
+	})
+	to := 1 * time.Second
+
+	fn(NewGracefulServer(p, h, to))
+}
+
+func TestNewGracefulServer(t *testing.T) {
+	gracefulServerContext(t, func(s *GracefulServer) {
+		if s.server.NoSignalHandling != true {
+			t.Error("NoSignalHandling should be true")
+		}
+
+		if s.server.Addr != ":1337" {
+			t.Error("Didn't set the port correctly")
+		}
+	})
+}
+
+func TestGracefulServerStopped(t *testing.T) {
+	gracefulServerContext(t, func(s *GracefulServer) {
+		if !s.Stopped() {
+			t.Error("Stopped returned false, but shouldn't")
+		}
+
+		s.setStopped(false)
+
+		if s.Stopped() {
+			t.Error("Stopped returned true, but shouldn't")
+		}
+	})
+}
+
+func TestGracefulServerStop(t *testing.T) {
+	done := make(chan bool)
+
+	gracefulServerContext(t, func(s *GracefulServer) {
+		time.AfterFunc(10*time.Millisecond, func() {
+			s.Stop()
+
+			if !s.Stopped() {
+				t.Error("Expected stopped to be true")
+			}
+
+			done <- true
+		})
+
+		s.ListenAndServe()
+	})
+
+	<-done
+}
+
+func TestGracefulServerListenAndServe(t *testing.T) {
+	done := make(chan bool)
+
+	gracefulServerContext(t, func(s *GracefulServer) {
+		time.AfterFunc(10*time.Millisecond, func() {
+			if s.Stopped() {
+				t.Error("The server should not be stopped after ListenAndServe")
+			}
+
+			s.Stop()
+			done <- true
+		})
+
+		s.ListenAndServe()
+	})
+
+	<-done
+}
+
+func TestGracefulServerListenAndServeTLS(t *testing.T) {
+	done := make(chan bool)
+
+	gracefulServerContext(t, func(s *GracefulServer) {
+		time.AfterFunc(10*time.Millisecond, func() {
+			if s.Stopped() {
+				t.Error("The server should not be stopped after ListenAndServe")
+			}
+
+			s.Stop()
+			done <- true
+		})
+
+		s.ListenAndServeTLS("foo", "bar")
+	})
+
+	<-done
+}