浏览代码

Add closing handshake code to echo example client

Gary Burd 10 年之前
父节点
当前提交
844dd6d40e
共有 1 个文件被更改,包括 32 次插入6 次删除
  1. 32 6
      examples/echo/client.go

+ 32 - 6
examples/echo/client.go

@@ -10,6 +10,8 @@ import (
 	"flag"
 	"log"
 	"net/url"
+	"os"
+	"os/signal"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -21,6 +23,9 @@ func main() {
 	flag.Parse()
 	log.SetFlags(0)
 
+	interrupt := make(chan os.Signal, 1)
+	signal.Notify(interrupt, os.Interrupt)
+
 	u := url.URL{Scheme: "ws", Host: *addr, Path: "/echo"}
 	log.Printf("connecting to %s", u.String())
 
@@ -30,13 +35,16 @@ func main() {
 	}
 	defer c.Close()
 
+	done := make(chan struct{})
+
 	go func() {
 		defer c.Close()
+		defer close(done)
 		for {
 			_, message, err := c.ReadMessage()
 			if err != nil {
 				log.Println("read:", err)
-				break
+				return
 			}
 			log.Printf("recv: %s", message)
 		}
@@ -45,11 +53,29 @@ func main() {
 	ticker := time.NewTicker(time.Second)
 	defer ticker.Stop()
 
-	for t := range ticker.C {
-		err := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
-		if err != nil {
-			log.Println("write:", err)
-			break
+	for {
+		select {
+		case t := <-ticker.C:
+			err := c.WriteMessage(websocket.TextMessage, []byte(t.String()))
+			if err != nil {
+				log.Println("write:", err)
+				return
+			}
+		case <-interrupt:
+			log.Println("interrupt")
+			// To cleanly close a connection, a client should send a close
+			// frame and wait for the server to close the connection.
+			err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+			if err != nil {
+				log.Println("write close:", err)
+				return
+			}
+			select {
+			case <-done:
+			case <-time.After(time.Second):
+			}
+			c.Close()
+			return
 		}
 	}
 }