vendor/github.com/gorilla/websocket/examples/chat/client.go

   1 // Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style
   3 // license that can be found in the LICENSE file.
   4
   5 package main
   6
   7 import (
   8         "bytes"
   9         "log"
  10         "net/http"
  11         "time"
  12
  13         "github.com/gorilla/websocket"
  14 )
  15
  16 const (
  17         // Time allowed to write a message to the peer.
  18         writeWait = 10 * time.Second
  19
  20         // Time allowed to read the next pong message from the peer.
  21         pongWait = 60 * time.Second
  22
  23         // Send pings to peer with this period. Must be less than pongWait.
  24         pingPeriod = (pongWait * 9) / 10
  25
  26         // Maximum message size allowed from peer.
  27         maxMessageSize = 512
  28 )
  29
  30 var (
  31         newline = []byte{'\n'}
  32         space   = []byte{' '}
  33 )
  34
  35 var upgrader = websocket.Upgrader{
  36         ReadBufferSize:  1024,
  37         WriteBufferSize: 1024,
  38 }
  39
  40 // Client is a middleman between the websocket connection and the hub.
  41 type Client struct {
  42         hub *Hub
  43
  44         // The websocket connection.
  45         conn *websocket.Conn
  46
  47         // Buffered channel of outbound messages.
  48         send chan []byte
  49 }
  50
  51 // readPump pumps messages from the websocket connection to the hub.
  52 //
  53 // The application runs readPump in a per-connection goroutine. The application
  54 // ensures that there is at most one reader on a connection by executing all
  55 // reads from this goroutine.
  56 func (c *Client) readPump() {
  57         defer func() {
  58                 c.hub.unregister <- c
  59                 c.conn.Close()
  60         }()
  61         c.conn.SetReadLimit(maxMessageSize)
  62         c.conn.SetReadDeadline(time.Now().Add(pongWait))
  63         c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
  64         for {
  65                 _, message, err := c.conn.ReadMessage()
  66                 if err != nil {
  67                         if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
  68                                 log.Printf("error: %v", err)
  69                         }
  70                         break
  71                 }
  72                 message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
  73                 c.hub.broadcast <- message
  74         }
  75 }
  76
  77 // writePump pumps messages from the hub to the websocket connection.
  78 //
  79 // A goroutine running writePump is started for each connection. The
  80 // application ensures that there is at most one writer to a connection by
  81 // executing all writes from this goroutine.
  82 func (c *Client) writePump() {
  83         ticker := time.NewTicker(pingPeriod)
  84         defer func() {
  85                 ticker.Stop()
  86                 c.conn.Close()
  87         }()
  88         for {
  89                 select {
  90                 case message, ok := <-c.send:
  91                         c.conn.SetWriteDeadline(time.Now().Add(writeWait))
  92                         if !ok {
  93                                 // The hub closed the channel.
  94                                 c.conn.WriteMessage(websocket.CloseMessage, []byte{})
  95                                 return
  96                         }
  97
  98                         w, err := c.conn.NextWriter(websocket.TextMessage)
  99                         if err != nil {
 100                                 return
 101                         }
 102                         w.Write(message)
 103
 104                         // Add queued chat messages to the current websocket message.
 105                         n := len(c.send)
 106                         for i := 0; i < n; i++ {
 107                                 w.Write(newline)
 108                                 w.Write(<-c.send)
 109                         }
 110
 111                         if err := w.Close(); err != nil {
 112                                 return
 113                         }
 114                 case <-ticker.C:
 115                         c.conn.SetWriteDeadline(time.Now().Add(writeWait))
 116                         if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
 117                                 return
 118                         }
 119                 }
 120         }
 121 }
 122
 123 // serveWs handles websocket requests from the peer.
 124 func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
 125         conn, err := upgrader.Upgrade(w, r, nil)
 126         if err != nil {
 127                 log.Println(err)
 128                 return
 129         }
 130         client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)}
 131         client.hub.register <- client
 132
 133         // Allow collection of memory referenced by the caller by doing all work in
 134         // new goroutines.
 135         go client.writePump()
 136         go client.readPump()
 137 }