105 lines
1.9 KiB
Go
105 lines
1.9 KiB
Go
// Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
|
|
|
|
package graceful
|
|
|
|
import (
|
|
"net"
|
|
"net/http"
|
|
"os"
|
|
"sync"
|
|
"time"
|
|
|
|
middleware "github.com/justinas/alice"
|
|
"gopkg.in/tylerb/graceful.v1"
|
|
)
|
|
|
|
const timeout = 5 * time.Second
|
|
|
|
type HTTPServer struct {
|
|
sync.Mutex
|
|
|
|
network string
|
|
router *http.ServeMux
|
|
server *graceful.Server
|
|
err error
|
|
}
|
|
|
|
func recovery(handler http.Handler) http.Handler {
|
|
f := func(w http.ResponseWriter, r *http.Request) {
|
|
defer func() {
|
|
if recover() != nil {
|
|
http.Error(w, "internal error, check logs for details", http.StatusInternalServerError)
|
|
}
|
|
}()
|
|
handler.ServeHTTP(w, r)
|
|
}
|
|
return http.HandlerFunc(f)
|
|
}
|
|
|
|
func NewHTTPServer(net, addr string, mw ...middleware.Constructor) *HTTPServer {
|
|
r := http.NewServeMux()
|
|
|
|
return &HTTPServer{
|
|
network: net,
|
|
router: r,
|
|
server: &graceful.Server{
|
|
Timeout: timeout,
|
|
Server: &http.Server{
|
|
Addr: addr,
|
|
Handler: middleware.New(recovery).Append(mw...).Then(r),
|
|
ReadTimeout: timeout,
|
|
WriteTimeout: timeout,
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func (s *HTTPServer) Handle(method, route string, handler http.HandlerFunc) {
|
|
f := func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != method {
|
|
http.NotFound(w, r)
|
|
return
|
|
}
|
|
handler.ServeHTTP(w, r)
|
|
}
|
|
s.router.HandleFunc(route, f)
|
|
}
|
|
|
|
func (s *HTTPServer) Serve() <-chan struct{} {
|
|
if s.network == "unix" {
|
|
os.Remove(s.server.Addr)
|
|
}
|
|
l, err := net.Listen(s.network, s.server.Addr)
|
|
if err != nil {
|
|
s.Lock()
|
|
s.err = err
|
|
s.Unlock()
|
|
c := make(chan struct{})
|
|
close(c)
|
|
return c
|
|
}
|
|
|
|
c := s.server.StopChan()
|
|
go func() {
|
|
s.Lock()
|
|
defer s.Unlock()
|
|
|
|
err = s.server.Serve(l)
|
|
if e, ok := err.(*net.OpError); !ok || (ok && e.Op != "accept") {
|
|
s.err = err
|
|
}
|
|
}()
|
|
return c
|
|
}
|
|
|
|
func (s *HTTPServer) Stop() {
|
|
s.server.Stop(timeout)
|
|
}
|
|
|
|
func (s *HTTPServer) Error() error {
|
|
s.Lock()
|
|
defer s.Unlock()
|
|
|
|
return s.err
|
|
}
|