From af95320fe1eff0661088ee5bc7c0f389c4dd88c3 Mon Sep 17 00:00:00 2001 From: jgheewala Date: Mon, 9 Nov 2020 09:56:07 -0800 Subject: [PATCH] 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 | 2 + web/servicestate.go | 119 +++++++++++++++++++++++++++++++++ web/servicestate_test.go | 139 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 260 insertions(+) create mode 100644 web/servicestate.go create mode 100644 web/servicestate_test.go diff --git a/web/closeheader.go b/web/closeheader.go index 141be75..6e32e75 100644 --- a/web/closeheader.go +++ b/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 diff --git a/web/servicestate.go b/web/servicestate.go new file mode 100644 index 0000000..b0dd25c --- /dev/null +++ b/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) + } + } + }) +} diff --git a/web/servicestate_test.go b/web/servicestate_test.go new file mode 100644 index 0000000..7610095 --- /dev/null +++ b/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) + }) +}