diff --git a/server/http/server.go b/server/http/server.go index 129e0b8..b2ae231 100644 --- a/server/http/server.go +++ b/server/http/server.go @@ -35,10 +35,9 @@ func constructor(srvcfg *chihaya.ServerConfig, tkr *tracker.Tracker) (server.Ser } type httpServer struct { - cfg *httpConfig - tkr *tracker.Tracker - grace *graceful.Server - stopping bool + cfg *httpConfig + tkr *tracker.Tracker + grace *graceful.Server } func (s *httpServer) Start() { @@ -49,9 +48,8 @@ func (s *httpServer) Start() { ReadTimeout: s.cfg.ReadTimeout, WriteTimeout: s.cfg.WriteTimeout, }, - Timeout: s.cfg.RequestTimeout, - NoSignalHandling: true, - ShutdownInitiated: func() { s.stopping = true }, + Timeout: s.cfg.RequestTimeout, + NoSignalHandling: true, ConnState: func(conn net.Conn, state http.ConnState) { switch state { case http.StateNew: @@ -76,7 +74,7 @@ func (s *httpServer) Start() { if err := s.grace.ListenAndServe(); err != nil { if opErr, ok := err.(*net.OpError); !ok || (ok && opErr.Op != "accept") { log.Printf("Failed to gracefully run HTTP server: %s", err.Error()) - return + panic(err) } } @@ -84,12 +82,8 @@ func (s *httpServer) Start() { } func (s *httpServer) Stop() { - if !s.stopping { - s.grace.Stop(s.grace.Timeout) - } - - s.grace = nil - s.stopping = false + s.grace.Stop(s.grace.Timeout) + <-s.grace.StopChan() } func (s *httpServer) routes() *httprouter.Router { diff --git a/server/prometheus/prometheus.go b/server/prometheus/prometheus.go index 50143a8..0bd9329 100644 --- a/server/prometheus/prometheus.go +++ b/server/prometheus/prometheus.go @@ -8,6 +8,8 @@ package prometheus import ( "errors" + "log" + "net" "net/http" "time" @@ -80,8 +82,18 @@ func (s *Server) Start() { Timeout: s.cfg.ShutdownTimeout, NoSignalHandling: true, } + + if err := s.grace.ListenAndServe(); err != nil { + if opErr, ok := err.(*net.OpError); !ok || (ok && opErr.Op != "accept") { + log.Printf("Failed to gracefully run Prometheus server: %s", err.Error()) + panic(err) + } + } + + log.Println("Prometheus server shut down cleanly") } func (s *Server) Stop() { s.grace.Stop(s.cfg.ShutdownTimeout) + <-s.grace.StopChan() } diff --git a/server/store/store.go b/server/store/store.go index b1f9cb9..e67ad18 100644 --- a/server/store/store.go +++ b/server/store/store.go @@ -111,6 +111,9 @@ type Store struct { } func (s *Store) Start() { + <-s.shutdown + s.wg.Wait() + log.Println("Store server shut down cleanly") } func (s *Store) Stop() {