package main import ( "fmt" "time" "github.com/gorilla/websocket" ) const ( writeWait = 10 * time.Second pongWait = 60 * time.Second pingPeriod = (pongWait * 9) / 10 maxMessageSize = 1024 * 1024 ) var upgrader = websocket.Upgrader{ ReadBufferSize: maxMessageSize, WriteBufferSize: maxMessageSize, } func (c *wsClient) writeListener() { ticker := time.NewTicker(pingPeriod) defer func() { ticker.Stop() c.ws.Close() }() for { select { case message, ok := <-c.send: if !ok { c.write(websocket.CloseMessage, []byte{}) return } if err := c.write(websocket.TextMessage, message); err != nil { return } case data, ok := <-c.sendJSON: if !ok { c.write(websocket.CloseMessage, []byte{}) return } if err := c.writeJSON(data); err != nil { return } case <-ticker.C: if err := c.write(websocket.PingMessage, []byte{}); err != nil { return } } } } func (c *wsClient) readListener(b *Bot) { defer func() { b.Hub.unregister <- c c.ws.Close() }() 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 { _, msg, err := c.ws.ReadMessage() if err != nil { break } fmt.Println("received message!") fmt.Printf("got message: %s\n", msg) b.Hub.broadcast <- string(msg) } } func (c *wsClient) write(mt int, message []byte) error { c.ws.SetWriteDeadline(time.Now().Add(writeWait)) return c.ws.WriteMessage(mt, message) } func (c *wsClient) writeJSON(data interface{}) error { c.ws.SetWriteDeadline(time.Now().Add(writeWait)) return c.ws.WriteJSON(data) }