以循环方式向客户端提供消息的 WebSocket 服务器

时间:2021-01-11 18:09:16

标签: go websocket

我在 Go 中有一个使用 Gorilla websocket 包的 websocket 服务器。在这个阶段,我将只有一台服务器为 5 个客户端提供服务。我从上游收到一些消息到 WebSocket 服务器。我的意图是不向连接的客户端广播所有消息。我只想以循环方式将消息的一个副本发送给连接的客户端。哪个客户得到它并不重要,只要只有一个客户得到它。

我尝试的解决方案 我有一个简单的 Go 服务器,创建了一个我正在接收的客户端池(websocket 连接)。但是,正如我上面提到的,我没有看到任何循环消息的选项。我所有的客户都收到了消息。如何只向连接的客户端发送一份消息副本,而不是向所有人广播。

Discalimer 我的代码取自在线资源并根据我的要求进行了修改。我对 Go 和 Websockets 比较陌生。 这是否可以使用 Websockets 实现?

main.go

package main

import (
    "fmt"
    "net/http"

    "github.com/realtime-chat-go-react/backend/pkg/websocket"
)

func serveWs(pool *websocket.Pool, w http.ResponseWriter, r *http.Request) {
    fmt.Println("WebSocket Endpoint Hit")
    conn, err := websocket.Upgrade(w, r)
    if err != nil {
        fmt.Fprintf(w, "%+v\n", err)
    }

    client := &websocket.Client{
        Conn: conn,
        Pool: pool,
    }

    pool.Register <- client
    client.Read()
}

func setupRoutes() {
    pool := websocket.NewPool()
    go pool.Start()

    http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
        serveWs(pool, w, r)
    })
}

func main() {
    setupRoutes()
    err := http.ListenAndServe(":8080",nil)

    if err != nil {
        fmt.Println(err)
    }
}

websocket.go

package websocket

import (
    "log"
    "net/http"

    "github.com/gorilla/websocket"
)

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

var wsList []*websocket.Conn

func Upgrade(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
    upgrader.CheckOrigin = func(r *http.Request) bool { return true }
    conn, err := upgrader.Upgrade(w, r, nil)
    wsList = append(wsList, conn) //Creating a list here to store all websocket clients.

    if err != nil {
        log.Println(err)
        return nil, err
    }

    return conn, nil
}

pool.go

package websocket

import "fmt"

type Pool struct {
    Register   chan *Client
    Unregister chan *Client
    Clients    map[*Client]bool
    Broadcast  chan Message
}

func NewPool() *Pool {
    return &Pool{
        Register:   make(chan *Client),
        Unregister: make(chan *Client),
        Clients:    make(map[*Client]bool),
        Broadcast:  make(chan Message),
    }
}

func (pool *Pool) Start() {
    for {
        select {
        case client := <-pool.Register:
            pool.Clients[client] = true
            fmt.Println("Size of Connection Pool: ", len(pool.Clients))
            for client, _ := range pool.Clients {
                fmt.Println(client)
                client.Conn.WriteJSON(Message{Type: 1, Body: "New User Joined..."})
            }
            break
        case client := <-pool.Unregister:
            delete(pool.Clients, client)
            fmt.Println("Size of Connection Pool: ", len(pool.Clients))
            for client, _ := range pool.Clients {
                client.Conn.WriteJSON(Message{Type: 1, Body: "User Disconnected..."})
            }
            break
        case message := <-pool.Broadcast:     //This is where I need to modify the code but not sure how
            fmt.Println("Sending message to all clients in Pool")
            for client, _ := range pool.Clients {
                if err := client.Conn.WriteJSON(message); err != nil {
                    fmt.Println(err)
                    return
                }
            }
        }
    }
}

client.go

package websocket

import (
    "fmt"
    "log"
    "sync"

    "github.com/gorilla/websocket"
)

type Client struct {
    ID   string
    Conn *websocket.Conn
    Pool *Pool
    mu   sync.Mutex
}

type Message struct {
    Type int    `json:"type"`
    Body string `json:"body"`
}

func (c *Client) Read() {
    defer func() {
        c.Pool.Unregister <- c
        c.Conn.Close()
    }()

    for {
        messageType, p, err := c.Conn.ReadMessage()
        if err != nil {
            log.Println(err)
            return
        }

        message := Message{Type: messageType, Body: string(p)}
        c.Pool.Broadcast <- message
        fmt.Printf("Message Received: %+v\n", message)

    }
}

1 个答案:

答案 0 :(得分:2)

修改池以将客户端存储在切片中而不是映射中。添加字段以记录之前使用过的客户端的索引。

type Pool struct {
    Register   chan *Client
    Unregister chan *Client
    Clients    []*Client
    Broadcast  chan Message
    PrevClientIndex int
}

循环而不是广播:

case message := <-pool.Broadcast: 
    if len(pool.Clients) == 0 {
        continue
    }
    pool.PrevClientIndex++
    if pool.PrevClientIndex >= len(pool.Clients) {
       pool.PrevClientIndex = 0
    }
    client := pool.Clients[pool.PrevClientIndex]
    if err := client.Conn.WriteJSON(message); err != nil {
        // handle error
        ...

注册附加到切片:

case client := <-pool.Register:
    pool.Clients = append(pool.Clients, client)
    ...

取消注册从切片中删除客户端:

case client := <-pool.Unregister:
    j := 0
    for _, c := range pool.Clients {
       if c != client {
          c.Clients[j] = c
          j++
       }
    }
    pool.Clients = pool.Clients[:j]
    ...