web: adding service state support

currently web CloseHeader handler was only setting connection close if the service was
in graceful shutdown. If not then it was assuming the service to be healthy and
sending in http request downstream. This can potentially lead to leaking of incoming http
requests to the backed end service for states other than graceful shutdown.

Adding "ServiceState" support which will add more context if and when the service
is in graceful shutdown regardless of the request being GET or POST. This will also
include support for setting the state of the service as 503 (ServiceUnavailable).
This will allow the backend service to reject any incoming http request when ALB
still thinks the service is healthy due to the fact that refresh is not yet kicked in.

One place where this service state management can come in handy is when the backend service
memory utilization is very high (~80-90%). If your state is not set to 503 then ALB will keep
on sending requests to the backend service, which just adds to the memory pressure.
The only way to protect the service is by removing from ALB which can again take up to couple
of seconds and can potentially lead to SIGKILL. If a backend web server sets using "ServiceState"
sets the state to 503 then even if ALB has not yet removed the instance, "ServiceState" as part of
web middle layers will cut off all incoming http requests, giving enough time for the backend
web server to relieve memory pressure. It's better to protect the backend service by sending 503;
instead of OOM'ing continuously and loosing customer request.

Test: Added test cases to handle the different GET & POST scenarios for all 3 states of
the service

web/closeheader.go

@@ -6,6 +6,8 @@ import (
 	"sync/atomic"
 )
 
+// Note: CloseHeader is deprecated with advance handling provided by ServiceState
+
 // CloseHeader is used to control when connections should signal they should be closed
 type CloseHeader struct {
 	SetCloseHeader int32

web/servicestate.go

@@ -0,0 +1,119 @@
+package web
+
+import (
+	"net/http"
+	"sync/atomic"
+)
+
+const (
+	// healthy state of service
+	healthy int32 = iota
+	// gracefulShutdown state of service
+	gracefulShutdown
+	// serviceUnavailable right now
+	serviceUnavailable
+)
+
+const (
+	// RespOKStr should be used for response ok in string format
+	RespOKStr = "OK"
+	// Connection is used to set connection Header
+	Connection = "Connection"
+	// Close is used to set the value of Connection Header
+	Close = "Close"
+)
+
+var (
+	// RespOKByte should be used to send response ok in byte
+	RespOKByte                 = []byte(RespOKStr)
+	serviceUnavailableRespByte = []byte("service temporarily un-available")
+	gracefulShutdownRespByte   = []byte("graceful shutdown")
+)
+
+func writeGracefulShutdownResponse(rw http.ResponseWriter) {
+	rw.WriteHeader(http.StatusNotFound)
+	_, _ = rw.Write(gracefulShutdownRespByte)
+}
+
+func writeServiceUnavailableResponse(rw http.ResponseWriter) {
+	rw.WriteHeader(http.StatusServiceUnavailable)
+	_, _ = rw.Write(serviceUnavailableRespByte)
+}
+
+// ServiceState is used to control when connections should signal they should be closed
+type ServiceState struct {
+	state int32
+}
+
+// ServiceUnavailable sets service state as Unavailable
+func (s *ServiceState) ServiceUnavailable() {
+	atomic.StoreInt32(&s.state, serviceUnavailable)
+}
+
+// Healthy sets service state as Healthy
+func (s *ServiceState) Healthy() {
+	atomic.StoreInt32(&s.state, healthy)
+}
+
+// GracefulShutdown sets service state as gracefulShutdown
+func (s *ServiceState) GracefulShutdown() {
+	atomic.StoreInt32(&s.state, gracefulShutdown)
+}
+
+// State returns current service state
+func (s *ServiceState) State() int32 {
+	return atomic.LoadInt32(&s.state)
+}
+
+// IsHealthy returns true if service is healthy
+func (s *ServiceState) IsHealthy() bool {
+	return atomic.LoadInt32(&s.state) == healthy
+}
+
+// IsInShutdown returns true if service is in graceful shutdown stage
+func (s *ServiceState) IsInShutdown() bool {
+	return atomic.LoadInt32(&s.state) == gracefulShutdown
+}
+
+// IsUnavailable returns true if service is temporarily un-available
+func (s *ServiceState) IsUnavailable() bool {
+	return atomic.LoadInt32(&s.state) == serviceUnavailable
+}
+
+// Checker will generate response writer based on status
+// and is also used to check if service can accept incoming request..
+// This comes in handy when the ALB refresh is not
+// kicked in and we want to protect the system
+// Note: next can be set to nil during GET request and it should still work
+func (s *ServiceState) Checker(next http.Handler, incHealthCheck func()) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		switch atomic.LoadInt32(&s.state) {
+		case gracefulShutdown:
+			// we will optionally set header as Connection closed
+			rw.Header().Set(Connection, Close)
+			switch r.Method {
+			case http.MethodGet:
+				writeGracefulShutdownResponse(rw)
+			case http.MethodPost:
+				// this is incoming request during graceful shutdown...
+				// let this request come through as header for connection
+				// closed is already set
+				next.ServeHTTP(rw, r)
+			}
+		case serviceUnavailable:
+			// whatever the type of request is lets return 503
+			writeServiceUnavailableResponse(rw)
+		default:
+			switch r.Method {
+			case http.MethodGet:
+				// GET is used during health check so write resp
+				// and return
+				_, _ = rw.Write(RespOKByte)
+				incHealthCheck()
+			case http.MethodPost:
+				// if system is Healthy let the request come in
+				next.ServeHTTP(rw, r)
+			}
+		}
+	})
+}

web/servicestate_test.go

@@ -0,0 +1,139 @@
+package web
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+func TestServiceState_Checker(t *testing.T) {
+	type testInfo struct {
+		name           string
+		method         string
+		next           http.HandlerFunc
+		state          int32
+		health         int
+		expectedHealth int
+		statusCode     int
+		resp           []byte
+		expectedHeader string
+	}
+	for _, tC := range []testInfo{
+		{
+			name:           "Healthy status",
+			method:         http.MethodGet,
+			next:           nil,
+			state:          healthy,
+			statusCode:     http.StatusOK,
+			expectedHealth: 1,
+			resp:           RespOKByte,
+		},
+		{
+			name:   "Healthy status",
+			method: http.MethodPost,
+			next: http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+				rw.Write(RespOKByte)
+				rw.WriteHeader(http.StatusOK)
+			}),
+			state:          healthy,
+			statusCode:     http.StatusOK,
+			expectedHealth: 0,
+			resp:           RespOKByte,
+		},
+		{
+			name:   "service unavailable status",
+			method: http.MethodPost,
+			next: http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+				panic(req)
+			}),
+			state:          serviceUnavailable,
+			statusCode:     http.StatusServiceUnavailable,
+			expectedHealth: 0,
+			resp:           serviceUnavailableRespByte,
+		},
+		{
+			name:   "service unavailable status",
+			method: http.MethodGet,
+			next: http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+				panic(req)
+			}),
+			state:          serviceUnavailable,
+			statusCode:     http.StatusServiceUnavailable,
+			expectedHealth: 0,
+			resp:           serviceUnavailableRespByte,
+		},
+		{
+			name:           "graceful shutdown status",
+			method:         http.MethodGet,
+			next:           nil,
+			state:          gracefulShutdown,
+			statusCode:     http.StatusNotFound,
+			expectedHealth: 0,
+			resp:           gracefulShutdownRespByte,
+			expectedHeader: Close,
+		},
+		{
+			name:   "graceful shutdown status",
+			method: http.MethodPost,
+			next: http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+				rw.Write(RespOKByte)
+				rw.WriteHeader(http.StatusOK)
+			}),
+			state:          gracefulShutdown,
+			statusCode:     http.StatusOK,
+			expectedHealth: 0,
+			resp:           RespOKByte,
+			expectedHeader: Close,
+		},
+	} {
+		tC := tC
+		t.Run(fmt.Sprintf("testing:%s and method:%s", tC.name, tC.method), func(t *testing.T) {
+			Convey(fmt.Sprintf("testing:%s and method:%s", tC.name, tC.method), t, func() {
+				rw := httptest.NewRecorder()
+				req, _ := http.NewRequestWithContext(context.Background(), tC.method, "", nil)
+				s := &ServiceState{
+					state: tC.state,
+				}
+				h := s.Checker(tC.next, func() { tC.health++ })
+				h.ServeHTTP(rw, req)
+				So(rw.Code, ShouldEqual, tC.statusCode)
+				So(tC.health, ShouldEqual, tC.expectedHealth)
+				So(rw.Body.Bytes(), ShouldResemble, tC.resp)
+				if tC.expectedHeader != "" {
+					So(rw.Header().Get(Connection), ShouldEqual, tC.expectedHeader)
+				}
+			})
+		})
+	}
+}
+
+func TestServiceState_ServiceUnavailable(t *testing.T) {
+	Convey("setting service as unavailable should work", t, func() {
+		s := &ServiceState{}
+		s.ServiceUnavailable()
+		So(s.State(), ShouldEqual, serviceUnavailable)
+		So(s.IsUnavailable(), ShouldBeTrue)
+	})
+}
+
+func TestServiceState_Healthy(t *testing.T) {
+	Convey("setting service healthy should work", t, func() {
+		s := &ServiceState{serviceUnavailable}
+		s.Healthy()
+		So(s.State(), ShouldEqual, healthy)
+		So(s.IsHealthy(), ShouldBeTrue)
+	})
+}
+
+func TestServiceState_GracefulShutdown(t *testing.T) {
+	Convey("setting service gracefulShutdown should work", t, func() {
+		s := &ServiceState{}
+		s.GracefulShutdown()
+		So(s.State(), ShouldEqual, gracefulShutdown)
+		So(s.IsInShutdown(), ShouldBeTrue)
+	})
+}