session shutdown cleanup

master
Jordan Orelli 4 years ago
parent d81ded9a9f
commit ecdf1af964

@ -17,13 +17,15 @@ import (
) )
type Server struct { type Server struct {
sync.Mutex
*blammo.Log *blammo.Log
Host string Host string
Port int Port int
http *http.Server http *http.Server
sync.Mutex
lastSessionID int lastSessionID int
sessions map[int]*session sessions map[int]*session
waitOnSessions sync.WaitGroup
} }
func (s *Server) Start() error { func (s *Server) Start() error {
@ -77,23 +79,31 @@ func (s *Server) createSession(conn *websocket.Conn) *session {
sn := &session{ sn := &session{
Log: s.Log.Child("sessions").Child(strconv.Itoa(s.lastSessionID)), Log: s.Log.Child("sessions").Child(strconv.Itoa(s.lastSessionID)),
id: s.lastSessionID, id: s.lastSessionID,
start: time.Now(),
conn: conn, conn: conn,
outbox: make(chan wire.Response), outbox: make(chan wire.Response),
done: make(chan chan struct{}, 1), done: make(chan bool, 1),
} }
if s.sessions == nil { if s.sessions == nil {
s.sessions = make(map[int]*session) s.sessions = make(map[int]*session)
} }
s.waitOnSessions.Add(1)
s.sessions[sn.id] = sn s.sessions[sn.id] = sn
s.Info("created session %d, %d sessions active", sn.id, len(s.sessions))
return sn return sn
} }
// dropSession removes a session from the server. This should only be called as
// a result of the connection's read loop terminating
func (s *Server) dropSession(sn *session) { func (s *Server) dropSession(sn *session) {
s.Lock() s.Lock()
defer s.Unlock() defer s.Unlock()
close(sn.done) close(sn.done)
delete(s.sessions, sn.id) delete(s.sessions, sn.id)
s.waitOnSessions.Add(-1)
s.Info("dropped session %d after %v time connected, %d sessions active", sn.id, time.Since(sn.start), len(s.sessions))
} }
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@ -105,35 +115,48 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
return return
} }
defer func() {
s.Info("closing connection")
if err := conn.Close(); err != nil {
s.Error("error closing connection: %v", err)
}
}()
sn := s.createSession(conn) sn := s.createSession(conn)
defer s.dropSession(sn)
go sn.run() go sn.run()
sn.read() sn.read()
s.dropSession(sn)
sn.Info("closing connection")
if err := conn.Close(); err != nil {
s.Error("error closing connection: %v", err)
}
} }
func (s *Server) Shutdown() { func (s *Server) Shutdown() {
s.Info("shutting down") s.Info("starting shutdown procedure")
s.http.Shutdown(context.Background())
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
s.Info("shutting down http server")
if err := s.http.Shutdown(context.Background()); err != nil {
s.Error("error shutting down http server: %v", err)
} else {
s.Info("http server has shut down")
}
}()
go func() {
defer wg.Done()
s.Info("broadcasting shutdown to all active sessions")
s.Lock() s.Lock()
zzz := make([]chan struct{}, 0, len(s.sessions))
for id, sn := range s.sessions { for id, sn := range s.sessions {
s.Info("sending done signal to session: %d", id) s.Info("sending done signal to session: %d", id)
c := make(chan struct{}) sn.done <- true
zzz = append(zzz, c)
sn.done <- c
} }
s.Unlock() s.Unlock()
for _, c := range zzz {
<-c s.Info("waiting on connected sessions to shut down")
} s.waitOnSessions.Wait()
time.Sleep(time.Second) s.Info("all sessions have shut down")
}()
wg.Wait()
s.Info("shutdown procedure complete")
} }

@ -13,9 +13,10 @@ import (
type session struct { type session struct {
*blammo.Log *blammo.Log
id int id int
start time.Time
conn *websocket.Conn conn *websocket.Conn
outbox chan wire.Response outbox chan wire.Response
done chan chan struct{} done chan bool
} }
// run is the session run loop. // run is the session run loop.
@ -26,9 +27,9 @@ func (sn *session) run() {
if err := sn.sendResponse(res); err != nil { if err := sn.sendResponse(res); err != nil {
sn.Error(err.Error()) sn.Error(err.Error())
} }
case c, ok := <-sn.done: case sendCloseFrame := <-sn.done:
sn.Info("saw done signal: %t", ok) sn.Info("saw done signal")
if ok { if sendCloseFrame {
sn.Info("sending close frame") sn.Info("sending close frame")
msg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "") msg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
if err := sn.conn.WriteMessage(websocket.CloseMessage, msg); err != nil { if err := sn.conn.WriteMessage(websocket.CloseMessage, msg); err != nil {
@ -36,7 +37,6 @@ func (sn *session) run() {
} else { } else {
sn.Info("sent close frame") sn.Info("sent close frame")
} }
close(c)
} }
return return
} }

Loading…
Cancel
Save