Added SetHeartbeatHandler to Session interface

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
+}

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