rpc: send websocket ping when connection is idle

rpc/websocket.go

@@ -25,15 +25,18 @@ import (
 	"os"
 	"strings"
 	"sync"
+	"time"
 
 	mapset "github.com/deckarep/golang-set"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/gorilla/websocket"
 )
 
 const (
-	wsReadBuffer  = 1024
-	wsWriteBuffer = 1024
+	wsReadBuffer       = 1024
+	wsWriteBuffer      = 1024
+	wsPingInterval     = 60 * time.Second
+	wsPingWriteTimeout = 5 * time.Second
 )
 
 var wsBufferPool = new(sync.Pool)
@@ -168,7 +171,64 @@ func wsClientHeaders(endpoint, origin string) (string, http.Header, error) {
 	return endpointURL.String(), header, nil
 }
 
+type websocketCodec struct {
+	*jsonCodec
+	conn *websocket.Conn
+
+	wg        sync.WaitGroup
+	pingReset chan struct{}
+}
+
 func newWebsocketCodec(conn *websocket.Conn) ServerCodec {
 	conn.SetReadLimit(maxRequestContentLength)
-	return NewFuncCodec(conn, conn.WriteJSON, conn.ReadJSON)
+	wc := &websocketCodec{
+		jsonCodec: NewFuncCodec(conn, conn.WriteJSON, conn.ReadJSON).(*jsonCodec),
+		conn:      conn,
+		pingReset: make(chan struct{}),
+	}
+	wc.wg.Add(1)
+	go wc.pingLoop()
+	return wc
+}
+
+func (wc *websocketCodec) close() {
+	wc.jsonCodec.close()
+	wc.wg.Wait()
+}
+
+func (wc *websocketCodec) writeJSON(ctx context.Context, v interface{}) error {
+	err := wc.jsonCodec.writeJSON(ctx, v)
+	if err == nil {
+		// Delay the next idle ping.
+		select {
+		case wc.pingReset <- struct{}{}:
+		case <-wc.jsonCodec.closed():
+		}
+	}
+	return err
+}
+
+// pingLoop sends periodic ping frames when the connection is idle.
+func (wc *websocketCodec) pingLoop() {
+	var timer = time.NewTimer(wsPingInterval)
+	defer wc.wg.Done()
+	defer timer.Stop()
+
+	for {
+		select {
+		case <-wc.closed():
+			return
+		case <-wc.pingReset:
+			if !timer.Stop() {
+				<-timer.C
+			}
+			timer.Reset(wsPingInterval)
+		case <-timer.C:
+			wc.jsonCodec.encMu.Lock()
+			wc.conn.SetWriteDeadline(time.Now().Add(wsPingWriteTimeout))
+			wc.conn.WriteMessage(websocket.PingMessage, nil)
+			wc.jsonCodec.encMu.Unlock()
+			timer.Reset(wsPingInterval)
+		}
+	}
 }