Websocket: Создание чат-комнат

Вот некомпилированная и непроверенная модификация хаба из примера Gorilla. 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)
                    }
                }
            }
        }
    }
}

Измените тип подключения, чтобы получить имя комнаты из строки запроса или сообщения, и включите это имя комнаты во все значения, отправляемые на каналы концентратора.

У меня нет времени писать и тестировать полный пример. Надеюсь, этого достаточно, чтобы вы начали.

Спасибо, это очень помогает. Приятно иметь пример, на который можно сослаться.

Для тех, кто пытается заставить это работать, мне, наконец, удалось заставить что-то функционировать должным образом. Надеюсь, это поможет любому, кто так же застрял и разочарован, как и я:

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)
                    }
                }
            }
        }
    }
}