Websocket: Creating chat 'rooms'

Here's an uncompiled and untested modification to the Gorilla example's hub.go to support rooms:

type message struct {
    data []byte
    room string
}

type subscription struct {
    conn *connection
    room string
}

// hub maintains the set of active connections and broadcasts messages to the
// connections.
type hub struct {
    // Registered connections.
    rooms map[strng]map[*connection]bool

    // Inbound messages from the connections.
    broadcast chan message

    // Register requests from the connections.
    register chan subscription

    // Unregister requests from connections.
    unregister chan subscription
}

func (h *hub) run() {
    for {
        select {
        case s := <-h.register:
            connections := h.rooms[sub.rooom]
            if connections == nil {
                connections = make(map[*connection]bool)
                h.rooms[s.room] = connections
            }
            connections[s.conn] = true
        case s := <-h.unregister:
            connections := h.rooms[s.rooom]
            if connections != nil {
                if _, ok := connections[s.conn]; ok {
                    delete(connections, s.conn)
                    close(s.conn.send)
                    if len(connections) == 0 {
                        delete(h.rooms, s.room)
                    }
                }
            }
        case m := <-h.broadcast:
            connections := h.rooms[m.rooom]
            for c := range h.connections {
                select {
                case c.send <- m.data:
                default:
                    close(c.send)
                    delete(h.connections, c)
                    if len(connections) == 0 {
                        delete(h.rooms, m.room)
                    }
                }
            }
        }
    }
}

Modify the connection type to get the room name from a query string or a message and include that room name in all values sent to the hub's channels.

I don't have time to write and test a complete example. Hopefully this is enough to get you started.

Thank you, this helps a lot. Great to have an example to refer to.

For anyone trying to get this to work, I managed to finally get something functioning properly. Hope this helps anyone who is as stuck and frustrated as I was:

conn.go:

package main

import (
    "log"
    "net/http"
    "time"

    "github.com/gorilla/mux"
    "github.com/gorilla/websocket"
)

const (
    // Time allowed to write a message to the peer.
    writeWait = 10 * time.Second

    // 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

    // Maximum message size allowed from peer.
    maxMessageSize = 512
)

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
}

// connection is an middleman between the websocket connection and the hub.
type connection struct {
    // The websocket connection.
    ws *websocket.Conn

    // Buffered channel of outbound messages.
    send chan []byte
}

// readPump pumps messages from the websocket connection to the hub.
func (s subscription) readPump() {
    c := s.conn
    defer func() {
        h.unregister <- s
        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 {
            if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
                log.Printf("error: %v", err)
            }
            break
        }
        m := message{msg, s.room}
        h.broadcast <- m
    }
}

// write writes a message with the given message type and payload.
func (c *connection) 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 (s *subscription) writePump() {
    c := s.conn
    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 <-ticker.C:
            if err := c.write(websocket.PingMessage, []byte{}); err != nil {
                return
            }
        }
    }
}

// serveWs handles websocket requests from the peer.
func serveWs(w http.ResponseWriter, r *http.Request) {
    ws, err := upgrader.Upgrade(w, r, nil)
    vars := mux.Vars(r)
    log.Println(vars["room"])
    if err != nil {
        log.Println(err)
        return
    }
    c := &connection{send: make(chan []byte, 256), ws: ws}
    s := subscription{c, vars["room"]}
    h.register <- s
    go s.writePump()
    s.readPump()
}

hub.go:

package main

type message struct {
    data []byte
    room string
}

type subscription struct {
    conn *connection
    room string
}

// hub maintains the set of active connections and broadcasts messages to the
// connections.
type hub struct {
    // Registered connections.
    rooms map[string]map[*connection]bool

    // Inbound messages from the connections.
    broadcast chan message

    // Register requests from the connections.
    register chan subscription

    // Unregister requests from connections.
    unregister chan subscription
}

var h = hub{
    broadcast:  make(chan message),
    register:   make(chan subscription),
    unregister: make(chan subscription),
    rooms:      make(map[string]map[*connection]bool),
}

func (h *hub) run() {
    for {
        select {
        case s := <-h.register:
            connections := h.rooms[s.room]
            if connections == nil {
                connections = make(map[*connection]bool)
                h.rooms[s.room] = connections
            }
            h.rooms[s.room][s.conn] = true
        case s := <-h.unregister:
            connections := h.rooms[s.room]
            if connections != nil {
                if _, ok := connections[s.conn]; ok {
                    delete(connections, s.conn)
                    close(s.conn.send)
                    if len(connections) == 0 {
                        delete(h.rooms, s.room)
                    }
                }
            }
        case m := <-h.broadcast:
            connections := h.rooms[m.room]
            for c := range connections {
                select {
                case c.send <- m.data:
                default:
                    close(c.send)
                    delete(connections, c)
                    if len(connections) == 0 {
                        delete(h.rooms, m.room)
                    }
                }
            }
        }
    }
}