add HandlePong

.travis.yml

@@ -3,6 +3,7 @@ sudo: required
 go:
 - 1.4
 - 1.5
+- 1.6
 before_install:
   - go get github.com/axw/gocov/gocov
   - go get github.com/mattn/goveralls

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 2016-05-09
+
+* Add method `HandlePong` to melody instance.
+
 ## 2015-10-07
 
 * Add broadcast methods for binary messages.

README.md

@@ -71,6 +71,7 @@ func main() {
 
 * Ola Holmström (@olahol)
 * Shogo Iwano (@shiwano)
+* Matt Caldwell (@mattcaldwell)
 
 ## FAQ
 

melody.go

@@ -18,6 +18,7 @@ type Melody struct {
 	errorHandler         handleErrorFunc
 	connectHandler       handleSessionFunc
 	disconnectHandler    handleSessionFunc
+	pongHandler          handleSessionFunc
 	hub                  *hub
 }
 
@@ -40,6 +41,7 @@ func New() *Melody {
 		errorHandler:         func(*Session, error) {},
 		connectHandler:       func(*Session) {},
 		disconnectHandler:    func(*Session) {},
+		pongHandler:          func(*Session) {},
 		hub:                  hub,
 	}
 }
@@ -54,6 +56,11 @@ func (m *Melody) HandleDisconnect(fn func(*Session)) {
 	m.disconnectHandler = fn
 }
 
+// Fires fn when a pong is received from a session.
+func (m *Melody) HandlePong(fn func(*Session)) {
+	m.pongHandler = fn
+}
+
 // Callback when a text message comes in.
 func (m *Melody) HandleMessage(fn func(*Session, []byte)) {
 	m.messageHandler = fn

melody_test.go

@@ -512,3 +512,33 @@ func TestSmallMessageBuffer(t *testing.T) {
 
 	conn.WriteMessage(websocket.TextMessage, []byte("12345"))
 }
+
+func TestPong(t *testing.T) {
+	echo := NewTestServerHandler(func(session *Session, msg []byte) {
+		session.Write(msg)
+	})
+	echo.m.Config.PongWait = time.Second
+	echo.m.Config.PingPeriod = time.Second * 9 / 10
+	server := httptest.NewServer(echo)
+	defer server.Close()
+
+	conn, err := NewDialer(server.URL)
+	defer conn.Close()
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	fired := false
+	echo.m.HandlePong(func(s *Session) {
+		fired = true
+	})
+
+	conn.WriteMessage(websocket.PongMessage, nil)
+
+	time.Sleep(time.Millisecond)
+
+	if !fired {
+		t.Error("should have fired pong handler")
+	}
+}

session.go

@@ -4,7 +4,6 @@ import (
 	"errors"
 	"github.com/gorilla/websocket"
 	"net/http"
-	"sync"
 	"time"
 )
 
@@ -12,7 +11,6 @@ import (
 type Session struct {
 	Request *http.Request
 	conn    *websocket.Conn
-	mutex   sync.Mutex
 	output  chan *envelope
 	melody  *Melody
 }
@@ -76,31 +74,17 @@ loop:
 	}
 }
 
-func (s *Session) SetPongHandler(f func() error) {
-	s.mutex.Lock()
-	{
-		s.conn.SetPongHandler(func(string) error {
-			s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
-			return f()
-		})
-	}
-	s.mutex.Unlock()
-}
-
 func (s *Session) readPump() {
 	defer s.conn.Close()
 
 	s.conn.SetReadLimit(s.melody.Config.MaxMessageSize)
 	s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
 
-	s.mutex.Lock()
-	{
-		s.conn.SetPongHandler(func(string) error {
-			s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
-			return nil
-		})
-	}
-	s.mutex.Unlock()
+	s.conn.SetPongHandler(func(string) error {
+		s.conn.SetReadDeadline(time.Now().Add(s.melody.Config.PongWait))
+		s.melody.pongHandler(s)
+		return nil
+	})
 
 	for {
 		t, message, err := s.conn.ReadMessage()