From d14d8b471525ed249398f84dbe56e8372c477f45 Mon Sep 17 00:00:00 2001
From: Gary Burd <gary@beagledreams.com>
Date: Sun, 18 Oct 2015 18:17:26 -0700
Subject: [PATCH] Add pinging to command example

---
 examples/command/main.go | 43 ++++++++++++++++++++++++++++++----------
 1 file changed, 33 insertions(+), 10 deletions(-)

diff --git a/examples/command/main.go b/examples/command/main.go
index ccd042d..73ac7e9 100644
--- a/examples/command/main.go
+++ b/examples/command/main.go
@@ -30,11 +30,19 @@ const (
 
 	// Maximum message size allowed from peer.
 	maxMessageSize = 8192
+
+	// 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
 )
 
 func pumpStdin(ws *websocket.Conn, w io.Writer) {
 	defer ws.Close()
 	ws.SetReadLimit(maxMessageSize)
+	ws.SetReadDeadline(time.Now().Add(pongWait))
+	ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
 	for {
 		_, message, err := ws.ReadMessage()
 		if err != nil {
@@ -45,11 +53,13 @@ func pumpStdin(ws *websocket.Conn, w io.Writer) {
 			break
 		}
 	}
-	log.Println("exit stdin pump")
 }
 
 func pumpStdout(ws *websocket.Conn, r io.Reader, done chan struct{}) {
-	defer ws.Close()
+	defer func() {
+		ws.Close()
+		close(done)
+	}()
 	s := bufio.NewScanner(r)
 	for s.Scan() {
 		ws.SetWriteDeadline(time.Now().Add(writeWait))
@@ -60,8 +70,21 @@ func pumpStdout(ws *websocket.Conn, r io.Reader, done chan struct{}) {
 	if s.Err() != nil {
 		log.Println("scan:", s.Err())
 	}
-	close(done)
-	log.Println("exit stdout pump")
+}
+
+func ping(ws *websocket.Conn, done chan struct{}) {
+	ticker := time.NewTicker(pingPeriod)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			if err := ws.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(writeWait)); err != nil {
+				log.Println("ping:", err)
+			}
+		case <-done:
+			return
+		}
+	}
 }
 
 func internalError(ws *websocket.Conn, msg string, err error) {
@@ -112,33 +135,33 @@ func serveWs(w http.ResponseWriter, r *http.Request) {
 	inr.Close()
 	outw.Close()
 
-	done := make(chan struct{})
-	go pumpStdout(ws, outr, done)
+	stdoutDone := make(chan struct{})
+	go pumpStdout(ws, outr, stdoutDone)
+	go ping(ws, stdoutDone)
 
 	pumpStdin(ws, inw)
 
 	// Some commands will exit when stdin is closed.
 	inw.Close()
 
-	// Other comamnds need a bonk on the head.
+	// Other commands need a bonk on the head.
 	if err := proc.Signal(os.Interrupt); err != nil {
 		log.Println("inter:", err)
 	}
 
 	select {
-	case <-done:
+	case <-stdoutDone:
 	case <-time.After(time.Second):
 		// A bigger bonk on the head.
 		if err := proc.Signal(os.Kill); err != nil {
 			log.Println("term:", err)
 		}
-		<-done
+		<-stdoutDone
 	}
 
 	if _, err := proc.Wait(); err != nil {
 		log.Println("wait:", err)
 	}
-	log.Println("exiting handler")
 }
 
 func serveHome(w http.ResponseWriter, r *http.Request) {