diff --git a/sockjs/session.go b/sockjs/session.go index 2cf42e5..fe5b6c8 100644 --- a/sockjs/session.go +++ b/sockjs/session.go @@ -49,9 +49,10 @@ type session struct { closeFrame string // internal timer used to handle session expiration if no receiver is attached, or heartbeats if recevier is attached - sessionTimeoutInterval time.Duration - heartbeatInterval time.Duration - timer *time.Timer + timeoutInterval time.Duration + heartbeatInterval time.Duration + heartbeatHandler func() + timer *time.Timer // once the session timeouts this channel also closes closeCh chan struct{} } @@ -71,20 +72,21 @@ type receiver interface { } // Session is a central component that handles receiving and sending frames. It maintains internal state -func newSession(req *http.Request, sessionID string, sessionTimeoutInterval, heartbeatInterval time.Duration) *session { +func newSession(req *http.Request, sessionID string, timeoutInterval, heartbeatInterval time.Duration) *session { r, w := io.Pipe() s := &session{ - id: sessionID, - req: req, - msgReader: r, - msgWriter: w, - msgEncoder: gob.NewEncoder(w), - msgDecoder: gob.NewDecoder(r), - sessionTimeoutInterval: sessionTimeoutInterval, - heartbeatInterval: heartbeatInterval, - closeCh: make(chan struct{})} + id: sessionID, + req: req, + msgReader: r, + msgWriter: w, + msgEncoder: gob.NewEncoder(w), + msgDecoder: gob.NewDecoder(r), + timeoutInterval: timeoutInterval, + heartbeatInterval: heartbeatInterval, + heartbeatHandler: nil, + closeCh: make(chan struct{})} s.Lock() // "go test -race" complains if ommited, not sure why as no race can happen here - s.timer = time.AfterFunc(sessionTimeoutInterval, s.close) + s.timer = time.AfterFunc(s.timeoutInterval, s.close) s.Unlock() return s } @@ -140,16 +142,21 @@ func (s *session) detachReceiver() { s.Lock() defer s.Unlock() s.timer.Stop() - s.timer = time.AfterFunc(s.sessionTimeoutInterval, s.close) + s.timer = time.AfterFunc(s.timeoutInterval, s.close) s.recv = nil } func (s *session) heartbeat() { s.Lock() - defer s.Unlock() + var heartbeatHandler func() if s.recv != nil { // timer could have fired between Lock and timer.Stop in detachReceiver s.recv.sendFrame("h") s.timer = time.AfterFunc(s.heartbeatInterval, s.heartbeat) + heartbeatHandler = s.heartbeatHandler + } + s.Unlock() + if heartbeatHandler != nil { + heartbeatHandler() } } @@ -222,3 +229,9 @@ func (s *session) ID() string { return s.id } func (s *session) Request() *http.Request { return s.req } + +func (s *session) SetHeartbeatHandler(handler func()) { + s.Lock() + defer s.Unlock() + s.heartbeatHandler = handler +} diff --git a/sockjs/sockjs.go b/sockjs/sockjs.go index f4d6916..e5ca1ed 100644 --- a/sockjs/sockjs.go +++ b/sockjs/sockjs.go @@ -8,6 +8,8 @@ type Session interface { ID() string // Request returns the first http request Request() *http.Request + // SetHeartbeatHandler is called after a heartbeat is sent. + SetHeartbeatHandler(handler func()) // Recv reads one text frame from session Recv() (string, error) // Send sends one text frame to session