diff --git a/examples/chat/conn.go b/examples/chat/conn.go index 6d7dd27..2fe2357 100644 --- a/examples/chat/conn.go +++ b/examples/chat/conn.go @@ -30,8 +30,8 @@ var upgrader = websocket.Upgrader{ WriteBufferSize: 1024, } -// connection is an middleman between the websocket connection and the hub. -type connection struct { +// Conn is an middleman between the websocket connection and the hub. +type Conn struct { // The websocket connection. ws *websocket.Conn @@ -40,51 +40,51 @@ type connection struct { } // readPump pumps messages from the websocket connection to the hub. -func (conn *connection) readPump() { +func (c *Conn) readPump() { defer func() { - mainHub.unregister <- conn - conn.ws.Close() + hub.unregister <- c + c.ws.Close() }() - conn.ws.SetReadLimit(maxMessageSize) - conn.ws.SetReadDeadline(time.Now().Add(pongWait)) - conn.ws.SetPongHandler(func(string) error { conn.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil }) + c.ws.SetReadLimit(maxMessageSize) + c.ws.SetReadDeadline(time.Now().Add(pongWait)) + c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil }) for { - _, message, err := conn.ws.ReadMessage() + _, message, err := c.ws.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) { log.Printf("error: %v", err) } break } - mainHub.broadcast <- message + hub.broadcast <- message } } // write writes a message with the given message type and payload. -func (conn *connection) write(mt int, payload []byte) error { - conn.ws.SetWriteDeadline(time.Now().Add(writeWait)) - return conn.ws.WriteMessage(mt, payload) +func (c *Conn) write(mt int, payload []byte) error { + c.ws.SetWriteDeadline(time.Now().Add(writeWait)) + return c.ws.WriteMessage(mt, payload) } // writePump pumps messages from the hub to the websocket connection. -func (conn *connection) writePump() { +func (c *Conn) writePump() { ticker := time.NewTicker(pingPeriod) defer func() { ticker.Stop() - conn.ws.Close() + c.ws.Close() }() for { select { - case message, ok := <-conn.send: + case message, ok := <-c.send: if !ok { - conn.write(websocket.CloseMessage, []byte{}) + c.write(websocket.CloseMessage, []byte{}) return } - if err := conn.write(websocket.TextMessage, message); err != nil { + if err := c.write(websocket.TextMessage, message); err != nil { return } case <-ticker.C: - if err := conn.write(websocket.PingMessage, []byte{}); err != nil { + if err := c.write(websocket.PingMessage, []byte{}); err != nil { return } } @@ -98,8 +98,8 @@ func serveWs(w http.ResponseWriter, r *http.Request) { log.Println(err) return } - conn := &connection{send: make(chan []byte, 256), ws: ws} - mainHub.register <- conn + conn := &Conn{send: make(chan []byte, 256), ws: ws} + hub.register <- conn go conn.writePump() conn.readPump() } diff --git a/examples/chat/hub.go b/examples/chat/hub.go index baa00e1..92da865 100644 --- a/examples/chat/hub.go +++ b/examples/chat/hub.go @@ -6,39 +6,39 @@ package main // hub maintains the set of active connections and broadcasts messages to the // connections. -type hub struct { +type Hub struct { // Registered connections. - connections map[*connection]bool + connections map[*Conn]bool // Inbound messages from the connections. broadcast chan []byte // Register requests from the connections. - register chan *connection + register chan *Conn // Unregister requests from connections. - unregister chan *connection + unregister chan *Conn } -var mainHub = hub{ +var hub = Hub{ broadcast: make(chan []byte), - register: make(chan *connection), - unregister: make(chan *connection), - connections: make(map[*connection]bool), + register: make(chan *Conn), + unregister: make(chan *Conn), + connections: make(map[*Conn]bool), } -func (hub *hub) run() { +func (h *Hub) run() { for { select { - case conn := <-hub.register: - hub.connections[conn] = true - case conn := <-hub.unregister: - if _, ok := hub.connections[conn]; ok { - delete(hub.connections, conn) + case conn := <-h.register: + h.connections[conn] = true + case conn := <-h.unregister: + if _, ok := h.connections[conn]; ok { + delete(h.connections, conn) close(conn.send) } - case message := <-hub.broadcast: - for conn := range hub.connections { + case message := <-h.broadcast: + for conn := range h.connections { select { case conn.send <- message: default: diff --git a/examples/chat/main.go b/examples/chat/main.go index 14b2e99..39943e6 100644 --- a/examples/chat/main.go +++ b/examples/chat/main.go @@ -29,7 +29,7 @@ func serveHome(w http.ResponseWriter, r *http.Request) { func main() { flag.Parse() - go mainHub.run() + go hub.run() http.HandleFunc("/", serveHome) http.HandleFunc("/ws", serveWs) err := http.ListenAndServe(*addr, nil)