diff --git a/.gitignore b/.gitignore index 66fd13c..839a56d 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,4 @@ *.out # Dependency directories (remove the comment below to include it) -# vendor/ +vendor/ diff --git a/README.md b/README.md index dd20804..d1be737 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,45 @@ # graceful Package graceful is a Go 1.8+ package enabling graceful shutdown of http.Handler servers. + +## Usage + +Simply use `ListenAndServe` to create a http server: + +```go +package main + +import ( + "log" + + "github.com/ETZhangSX/graceful" + "github.com/gin-gonic/gin" +) + +func main() { + r := gin.Default() + if err := graceful.ListenAndServe(":8080", r.Handler()); err != nil { + log.Fatal(err) + } +} +``` + +The default timeout of shutdown is 5 seconds. You can configure it by using `WithShutdownTimeout` + +```go +func main() { + ... + ... + + if err := graceful.ListenAndServe( + ":8080", + handler, + graceful.WithShutdownTimeout(10*time.Second), + // add customer function before shutting down + graceful.WithShutdownFunc(func() { + log.Info("Http Server is shutting down...") + } + ); err != nil { + log.Fatal(err) + } +} +``` diff --git a/doc.go b/doc.go new file mode 100644 index 0000000..d91518b --- /dev/null +++ b/doc.go @@ -0,0 +1,22 @@ +/* +Package graceful is a Go 1.8+ package enabling graceful shutdown of http.Handler servers. + +Usage: + + package main + + import ( + "log" + + "github.com/ETZhangSX/graceful" + "github.com/gin-gonic/gin" + ) + + func main() { + r := gin.Default() + if err := graceful.ListenAndServe(":8080", r.Handler()); err != nil { + log.Fatal(err) + } + } +*/ +package graceful diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..b19d824 --- /dev/null +++ b/go.mod @@ -0,0 +1,11 @@ +module github.com/ETZhangSX/graceful + +go 1.18 + +require github.com/stretchr/testify v1.7.2 + +require ( + github.com/davecgh/go-spew v1.1.0 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..834baff --- /dev/null +++ b/go.sum @@ -0,0 +1,11 @@ +github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.7.2 h1:4jaiDzPyXQvSd7D0EjG45355tLlV3VOECpq10pLC+8s= +github.com/stretchr/testify v1.7.2/go.mod h1:R6va5+xMeoiuVRoj+gSkQ7d3FALtqAAGI1FQKckRals= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= \ No newline at end of file diff --git a/graceful.go b/graceful.go new file mode 100644 index 0000000..60225d9 --- /dev/null +++ b/graceful.go @@ -0,0 +1,114 @@ +package graceful + +import ( + "context" + "fmt" + "net/http" + "os" + "os/signal" + "syscall" + "time" +) + +const DefaultShutdownTimeout = 5 * time.Second + +// A Server defines parameters for gracefully running an HTTP server. +// The zero value for Server is a valid configuration. +type Server struct { + *http.Server + + // ShutdownTimeout is the maximum duration for shutting down the server. + // A zero or negative value means there will be no timeout. + ShutdownTimeout time.Duration + + errChan chan error +} + +// init server +func (s *Server) init() { + s.errChan = make(chan error, 1) +} + +func (s *Server) load(opts []Option) { + for _, opt := range opts { + opt.apply(s) + } +} + +// ListenAndServe listens on the TCP network address s.Addr and then +// calls s.Server.ListenAndServe to handle requests on incoming connections. +func (s *Server) ListenAndServe(opts ...Option) error { + s.init() + s.load(opts) + go func() { + if err := s.Server.ListenAndServe(); err != nil && err != http.ErrServerClosed { + s.errChan <- fmt.Errorf("[graceful] %w", err) + } + }() + return s.waitForShutdown() +} + +// ListenAndServeTLS listens on the TCP network address srv.Addr and +// then calls ServeTLS to handle requests on incoming TLS connections. +// Accepted connections are configured to enable TCP keep-alives. +func (s *Server) ListenAndServeTLS(certFile, keyFile string, opts ...Option) error { + s.init() + s.load(opts) + go func() { + if err := s.Server.ListenAndServeTLS(certFile, keyFile); err != nil && err != http.ErrServerClosed { + s.errChan <- fmt.Errorf("[graceful] %w", err) + } + }() + return s.waitForShutdown() +} + +// waiting for shutdown or error occur. +func (s *Server) waitForShutdown() error { + quit := make(chan os.Signal, 1) + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + select { + case err := <-s.errChan: + return err + case <-quit: + } + ctx, cancel := context.WithTimeout(context.Background(), s.ShutdownTimeout) + defer cancel() + if err := s.Server.Shutdown(ctx); err != nil { + return err + } + return nil +} + +// ListenAndServe listens on the TCP network address addr and then calls +// Serve with handler to handle requests on incoming connections. +// Accepted connections are configured to enable TCP keep-alives. +// +// The handler is typically nil, in which case the DefaultServeMux is used. +// ShutdownTimeout defaults 5 seconds. +func ListenAndServe(addr string, handler http.Handler, opts ...Option) error { + server := &Server{ + Server: &http.Server{ + Addr: addr, + Handler: handler, + }, + ShutdownTimeout: DefaultShutdownTimeout, + } + return server.ListenAndServe(opts...) +} + +// ListenAndServeTLS acts identically to ListenAndServe, except that it +// expects HTTPS connections. Additionally, files containing a certificate and +// matching private key for the server must be provided. If the certificate +// is signed by a certificate authority, the certFile should be the concatenation +// of the server's certificate, any intermediates, and the CA's certificate. +// ShutdownTimeout defaults 5 seconds. +func ListenAndServeTLS(addr, certFile, keyFile string, handler http.Handler, opts ...Option) error { + server := &Server{ + Server: &http.Server{ + Addr: addr, + Handler: handler, + }, + ShutdownTimeout: DefaultShutdownTimeout, + } + return server.ListenAndServeTLS(certFile, keyFile, opts...) +} diff --git a/graceful_test.go b/graceful_test.go new file mode 100644 index 0000000..dbb9fbe --- /dev/null +++ b/graceful_test.go @@ -0,0 +1,13 @@ +package graceful + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestListenAndServe(t *testing.T) { + err := ListenAndServe("bbb", nil) + + assert.Equal(t, "[graceful] listen tcp: address bbb: missing port in address", err.Error()) +} diff --git a/option.go b/option.go new file mode 100644 index 0000000..8a4f089 --- /dev/null +++ b/option.go @@ -0,0 +1,41 @@ +package graceful + +import ( + "time" +) + +// Option interface +type Option interface { + // apply option to server + apply(*Server) +} + +// option wraps a function that modifies Server into an +// implementation of the Option interface. +type option struct { + f func(*Server) +} + +func (o *option) apply(s *Server) { + o.f(s) +} + +func newOption(f func(s *Server)) *option { + return &option{f: f} +} + +// WithShutdownTimeout set timeout for shutting down server +func WithShutdownTimeout(timeout time.Duration) Option { + return newOption(func(s *Server) { + s.ShutdownTimeout = timeout + }) +} + +// WithShutdownFunc registers function(s) running while shutting down by calling s.RegisterOnShutdown. +func WithShutdownFunc(fs ...func()) Option { + return newOption(func(s *Server) { + for _, f := range fs { + s.RegisterOnShutdown(f) + } + }) +}