diff --git a/.travis.yml b/.travis.yml index c7dde77..bf528aa 100644 --- a/.travis.yml +++ b/.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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 74b4888..804f071 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 2016-05-09 + +* Add method `HandlePong` to melody instance. + ## 2015-10-07 * Add broadcast methods for binary messages. diff --git a/README.md b/README.md index e61f486..13f95c0 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,7 @@ func main() { * Ola Holmström (@olahol) * Shogo Iwano (@shiwano) +* Matt Caldwell (@mattcaldwell) ## FAQ diff --git a/melody.go b/melody.go index acf8c45..849fc3b 100644 --- a/melody.go +++ b/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 diff --git a/melody_test.go b/melody_test.go index 3d5fca5..1b47059 100644 --- a/melody_test.go +++ b/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") + } +} diff --git a/session.go b/session.go index 24fb950..f262364 100644 --- a/session.go +++ b/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()