Golang websocket推送

Golang websocket推送

在工做用主要使用的是Java，也作過IM（後端用的netty websocket）。最近想經過Golang重寫下，因而經過websocket擼了一個聊天室。

項目地址

Github

依賴

golang.org/x/net下的websocket。

因爲我使用的是golang版本是1.12，在國內訪問golang.org/x須要藉助代理，或者經過replace替換爲github下的鏡像。

module github.com/xuanbo/pusher

require golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3

replace (
        golang.org/x/crypto => github.com/golang/crypto v0.0.0-20190308221718-c2843e01d9a2
        golang.org/x/net => github.com/golang/net v0.0.0-20190404232315-eb5bcb51f2a3
        golang.org/x/sys => github.com/golang/sys v0.0.0-20190215142949-d0b11bdaac8a
        golang.org/x/text => github.com/golang/text v0.3.0
)

即工程下的go.mod.cn文件。

websocket用法

核心就是for循環下的處理收到的消息邏輯，而後對消息進行處理（轉發、廣播等）。

// websocket Handler
// usage: http.Handle("/websocket", websocket.Handler(pusher.Handler))
func Handler(conn *websocket.Conn) {
    // handle connected
    var userId string
    var err error
    if userId, err = doConnected(conn); err != nil {
        fmt.Println("Client connect error: ", err)
        return
    }

    fmt.Println("Client connected, userId: ", userId)

    for {
        msg := new(Message)

        if err := websocket.JSON.Receive(conn, msg); err != nil {
            fmt.Println("Can't receive, error: ", err)
            break
        }

        msg.UpdateAt = Timestamp()

        fmt.Println("Received from client: ", msg)

        // handle received message
        if err := doReceived(conn, msg); err != nil {
            fmt.Println("Received message error: ", err)
            break
        }
    }

    // handle disConnected
    if err := doDisConnected(userId, conn); err != nil {
        fmt.Println("Client disconnected error: ", err)
        return
    }

    fmt.Println("Client disconnected, userId: ", userId)
}

鏈接管理

在IM中比較重要的點就是管理客戶端鏈接，這樣咱們才能經過服務端轉發消息給對應的用戶。注意，下面沒有考慮集羣，只在單機中考慮。

// websocket connection manager
type ConnManager struct {
    // websocket connection number
    Online *int32
    // websocket connection
    connections *sync.Map
}

上面定義了一個鏈接管理結構體，Online爲在線的人數，connections爲客戶端的鏈接管理（key爲userId，value爲websocket connection）。

下面爲ConnManager添加一些方法來處理鏈接、斷開鏈接、發送消息、廣播等操做。

// add websocket connection
// online number + 1
func (m *ConnManager) Connected(k, v interface{}) {
    m.connections.Store(k, v)

    atomic.AddInt32(m.Online, 1)
}

// remove websocket connection by key
// online number - 1
func (m *ConnManager) DisConnected(k interface{}) {
    m.connections.Delete(k)

    atomic.AddInt32(m.Online, -1)
}

// get websocket connection by key
func (m *ConnManager) Get(k interface{}) (v interface{}, ok bool) {
    return m.connections.Load(k)
}

// iter websocket connections
func (m *ConnManager) Foreach(f func(k, v interface{})) {
    m.connections.Range(func(k, v interface{}) bool {
        f(k, v)
        return true
    })
}

// send message to one websocket connection
func (m *ConnManager) Send(k string, msg *Message) {
    v, ok := m.Get(k)
    if ok {
        if conn, ok := v.(*websocket.Conn); ok {
            if err := websocket.JSON.Send(conn, msg); err != nil {
                fmt.Println("Send msg error: ", err)
            }
        } else {
            fmt.Println("invalid type, expect *websocket.Conn")
        }
    } else {
        fmt.Println("connection not exist")
    }
}

// send message to multi websocket connections
func (m *ConnManager) SendMulti(keys []string, msg interface{}) {
    for _, k := range keys {
        v, ok := m.Get(k)
        if ok {
            if conn, ok := v.(*websocket.Conn); ok {
                if err := websocket.JSON.Send(conn, msg); err != nil {
                    fmt.Println("Send msg error: ", err)
                }
            } else {
                fmt.Println("invalid type, expect *websocket.Conn")
            }
        } else {
            fmt.Println("connection not exist")
        }
    }
}

// broadcast message to all websocket connections otherwise own connection
func (m *ConnManager) Broadcast(conn *websocket.Conn, msg *Message) {
    m.Foreach(func(k, v interface{}) {
        if c, ok := v.(*websocket.Conn); ok && c != conn {
            if err := websocket.JSON.Send(c, msg); err != nil {
                fmt.Println("Send msg error: ", err)
            }
        }
    })
}

消息類型、格式

消息類型（MessageType）主要有單聊、羣聊、系統通知等。

消息格式（MediaType）主要有文本格式、圖片、文件等。

type MessageType int
type MediaType int

const (
    Single MessageType = iota
    Group
    SysNotify
    OnlineNotify
    OfflineNotify
)

const (
    Text MediaType = iota
    Image
    File
)

// websocket message
type Message struct {
    MessageType MessageType `json:"messageType"`
    MediaType   MediaType   `json:"mediaType"`
    From        string      `json:"from"`
    To          string      `json:"to"`
    Content     string      `json:"content,omitempty"`
    FileId      string      `json:"fileId,omitempty"`
    Url         string      `json:"url,omitempty"`
    CreateAt    int64       `json:"createAt,omitempty"`
    UpdateAt    int64       `json:"updateAt,omitempty"`
}

上面定義了一個統一的消息（Message）。

效果

前端的代碼就不展現了，最終實現的聊天室效果以下：

補充

本例子沒有涉及到用戶認證、消息加密、idle、單聊、消息格式、消息持久化等等，只作了一個簡單的羣聊。

歡迎感興趣的道友，基於此擴展出本身的推送系統、IM等。

說明

Just for fun！