From fa75a59f10231735bd3a1886df809f76b3d80272 Mon Sep 17 00:00:00 2001
From: willsky <hdu_willsky@foxmail.com>
Date: Fri, 2 Oct 2015 22:20:47 +0800
Subject: [PATCH] add websocket ping/pong feature

---
 session/sessions.go | 38 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/session/sessions.go b/session/sessions.go
index 49daa8c..b93378a 100644
--- a/session/sessions.go
+++ b/session/sessions.go
@@ -25,6 +25,7 @@ package session
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"math/rand"
 	"net/http"
 	"os"
@@ -189,6 +190,17 @@ func (f userReports) Len() int           { return len(f) }
 func (f userReports) Swap(i, j int)      { f[i], f[j] = f[j], f[i] }
 func (f userReports) Less(i, j int) bool { return f[i].processCnt > f[j].processCnt }
 
+const (
+	// Time allowed to write a message to the peer.
+	writeWait = 10 * time.Second
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 60 * time.Second
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+)
+
 // WSHandler handles request of creating session channel.
 //
 // When a channel closed, releases all resources associated with it.
@@ -225,13 +237,34 @@ func WSHandler(w http.ResponseWriter, r *http.Request) {
 	logger.Tracef("Open a new [Session Channel] with session [%s], %d", sid, len(SessionWS))
 
 	input := map[string]interface{}{}
+	
+	wsChan.Conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+	wsChan.Conn.SetPongHandler(func(string) error { wsChan.Conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
+	ticker := time.NewTicker(pingPeriod)
+	defer func() {
+		WideSessions.Remove(sid)
+		ticker.Stop()
+		wsChan.Close()
+	}()
+
+	// send websocket ping message.
+	go func(t *time.Ticker, channel util.WSChannel) {
+		for {
+			select {
+			case <-t.C:
+				if err := channel.Conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
+					return
+				}
+			}
+		}
+
+	}(ticker, wsChan)
 
 	for {
+
 		if err := wsChan.ReadJSON(&input); err != nil {
 			logger.Tracef("[Session Channel] of session [%s] disconnected, releases all resources with it, user [%s]", sid, wSession.Username)
 
-			WideSessions.Remove(sid)
-
 			return
 		}
 
@@ -245,6 +278,7 @@ func WSHandler(w http.ResponseWriter, r *http.Request) {
 
 		wsChan.Time = time.Now()
 	}
+
 }
 
 // SaveContentHandler handles request of session content string.