GO websocket 实现简易聊天室

架构图如下：

本项目由四个文件组成：

1. hub.go
2. client.go
3. main.go
4. home.html

Hub结构体实现：

• 拥有每一个Client的指针
• 一个boardcast管道接收任意Client的消息
• 接收用户注册的管道
• 接收用户注销的管道

1.  
   type Hub struct{
2.  
   broadcast chan string //broadcast管道里有数据时把它写入每一个Client的send管道中
3.  
   clients map[*Client]struct{} //Hub持有每个client的指针
4.  
   register chan *Client //注册管道
5.  
   unregister chan *Client //注销管道
6.  
   }

Hub构造函数：

1.  
   func NewHub()*Hub{
2.  
   return &Hub{broadcast: make(chan []byte),
3.  
   clients: make(map[*Client]struct{}),
4.  
   register: make(chan *Client), unregister: make(chan *Client)}
5.  
   }

Client结构体实现：

• 与前端的websocket连接
• hub的指针
• send管道传输信息
• name 字符串保存前端用户的姓名

1.  
   type Client struct{
2.  
   hub *Hub
3.  
   conn *websocket.Conn
4.  
   send chan []byte
5.  
   name []byte
6.  
   }

Hub工作实现：

• 若注册管道有输入则在map中注册
• 若注销管道有输入则在map中删除并将该client的send管道关闭
• 若boardcast管道有输入则对map里的每个client的send管道输入

1.  
   func (hub *Hub) Run(){
2.  
   for{
3.  
   select{
4.  
   case client := <-hub.register:
5.  
   hub.clients[client] = struct{}{}
6.  
   case client := <- hub.unregister:
7.  
   delete(hub.clients,client)
8.  
   close(client.send)
9.  
   case msg := <-hub.broadcast:
10.  
    for client := range hub.clients{
11.  
    select{
12.  
    case client.send <- msg://如果管道不能立即写入数据，就认为该client出故障了
13.  
    default:
14.  
    close(client.send)
15.  
    delete(hub.clients, client)
16.  
    }
17.  
    }
18.  
    }
19.  
     
20.  
    }
21.  
    }

Client从websocket读取内容：

• 善后工作：注销client，关闭websocket连接
• connection设置最大读入量和ping pong时间
• 死循环读取前端消息

1.  
   const (
2.  
   writeWait = 10 * time.Second //
3.  
   pongWait = 60 * time.Second // 每60秒向websocket发送一次pong
4.  
   pingPeriod = 9 * pongWait / 10 //连接不断时每隔54秒向client发送一次ping
5.  
   maxMsgSize = 512 //消息的长度不能超过512
6.  
   )
7.  
   // 从websocket读取数据
8.  
   func(client *Client)read(){
9.  
   defer func(){
10.  
    client.hub.unregister <- client //向hub发送注销
11.  
    fmt.Printf("%s offline\n", client.name)
12.  
    fmt.Printf("close connection to %s\n",client.conn.RemoteAddr().String())
13.  
    client.conn.Close() //关闭ws连接
14.  
    }()
15.  
     
16.  
    // conn细节设置
17.  
    client.conn.SetReadLimit(maxMsgSize)
18.  
    client.conn.SetReadDeadline(time.Now().Add(pongWait)) //设置最长可读时间
19.  
    client.conn.SetPongHandler(func(appData string) error {
20.  
    client.conn.SetReadDeadline(time.Now().Add(pongWait))//每次接收到ping后都将最长可读时间延后60秒
21.  
    return nil
22.  
    })
23.  
     
24.  
    for{
25.  
    _, p, err := client.conn.ReadMessage() //返回消息类型，消息，error
26.  
    if err != nil{
27.  
    //如果以意料之外的关闭状态关闭，就打印日志
28.  
    if websocket.IsUnexpectedCloseError(err, websocket.CloseAbnormalClosure, websocket.CloseGoingAway) {
29.  
    fmt.Printf("close websocket conn error: %v\n", err)
30.  
    }
31.  
    break //只要ReadMessage失败，就关闭websocket管道、注销client，退出
32.  
    }else{
33.  
    // trimspace:消去首尾空格， replace：将换行符换位空格，-1：全部转换
34.  
    message := bytes.TrimSpace(bytes.Replace(p,[]byte{'\n'}, []byte{' '}, -1))
35.  
    if len(client.name) == 0{//第一次输入的内容为自己的名字
36.  
    client.name = message
37.  
    }else{
38.  
    client.hub.broadcast<-bytes.Join([][]byte{client.name, message}, []byte(": "))
39.  
    }
40.  
    }
41.  
    }

Client 从hub中接收信息：

• 设置ticker，每pingPeriod时间向websocket发送ping并延长可写时间
• 善后工作：结束ticker，关闭连接
• 死循环：接收ticker和send管道的消息
• send管道：判断是否关闭，不是则创建一个writer写入message

1.  
   func(client *Client)write() {
2.  
   ticker := time.NewTicker(pingPeriod)
3.  
   defer func(){
4.  
   ticker.Stop() //ticker不用就stop，防止协程泄漏
5.  
   fmt.Printf("close connection to %s\n", client.conn.RemoteAddr().String())
6.  
   client.conn.Close() //给前端写数据失败，就可以关系连接了
7.  
   }()
8.  
   for{
9.  
   select{
10.  
    case msg, ok := <-client.send:
11.  
    if !ok{
12.  
    fmt.Println("管道已经关闭")
13.  
    client.conn.WriteMessage(websocket.CloseMessage, []byte{})
14.  
    return
15.  
    }
16.  
    client.conn.SetWriteDeadline(time.Now().Add(writeWait))10秒内必须把信息写给前端（写到websocket连接里去），否则就关闭连接
17.  
    if writer, err := client.conn.NextWriter(websocket.TextMessage); err != nil{
18.  
    return
19.  
    }else{
20.  
    writer.Write(msg)
21.  
    writer.Write([]byte{'\n'})
22.  
    // 有消息一次全写出去
23.  
    n := len(client.send)
24.  
    for i := 0; i < n; i {
25.  
    writer.Write(<-client.send)
26.  
    writer.Write([]byte{'\n'})
27.  
    }
28.  
    if err := writer.Close(); err != nil { //必须调close，否则下次调用client.conn.NextWriter时本条消息才会发送给浏览器
29.  
    return //结束一切
30.  
    }
31.  
    }
32.  
    case <-ticker.C:
33.  
    client.conn.SetWriteDeadline(time.Now().Add(writeWait))
34.  
    if err := client.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil{
35.  
    return
36.  
    }
37.  
    }
38.  
    }
39.  
    }

主页面路由设置:

1.  
   func serveHome(w http.ResponseWriter, r *http.Request) {
2.  
   if r.URL.Path != "/" { //只允许访问根路径
3.  
   http.Error(w, "Not Found", http.StatusNotFound)
4.  
   return
5.  
   }
6.  
   if r.Method != "GET" { //只允许GET请求
7.  
   http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
8.  
   return
9.  
   }
10.  
    http.ServeFile(w, r, "socket/chat_room/home.html") //请求根目录时直接返回一个html页面
11.  
    }

websocket服务：

1.  
   func ServeWS(hub *Hub, w http.ResponseWriter, r *http.Request){
2.  
   upgrader := websocket.Upgrader{
3.  
   HandshakeTimeout: 2 * time.Second, //握手超时时间
4.  
   ReadBufferSize: 1024, //读缓冲大小
5.  
   WriteBufferSize: 1024, //写缓冲大小
6.  
   CheckOrigin: func(r *http.Request) bool { return true },
7.  
   Error: func(w http.ResponseWriter, r *http.Request, status int, reason error) {},
8.  
   }
9.  
   conn, err := upgrader.Upgrade(w,r,nil)
10.  
    checkError(err)
11.  
    fmt.Printf("connect to client %s\n", conn.RemoteAddr().String())
12.  
    client := &Client{hub: hub, conn: conn, send: make(chan []byte,256)}
13.  
    hub.register <- client
14.  
    go client.read()
15.  
    go client.write()
16.  
    }

main函数：

1.  
   func main(){
2.  
   hub := NewHub()
3.  
   go hub.Run()
4.  
   http.HandleFunc("/", ServeHome)
5.  
   http.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) {
6.  
   ServeWS(hub, w, r)
7.  
   })
8.  
   if err := http.ListenAndServe(":5656",nil);err != nil{
9.  
   fmt.Printf("start http service error: %s\n", err)
10.  
    }
11.  
    }

前端实现：直接套模板 

1.  
   <!DOCTYPE html>
2.  
   <html lang="en">
3.  
    
4.  
   <head>
5.  
   <title>聊天室</title>
6.  
   <script type="text/javascript">
7.  
   window.onload = function () {//页面打开时执行以下初始化内容
8.  
   var conn;
9.  
   var msg = document.getElementById("msg");
10.  
    var log = document.getElementById("log");
11.  
     
12.  
    function appendLog(item) {
13.  
    var doScroll = log.scrollTop > log.scrollHeight - log.clientHeight - 1;
14.  
    log.appendChild(item);
15.  
    if (doScroll) {
16.  
    log.scrollTop = log.scrollHeight - log.clientHeight;
17.  
    }
18.  
    }
19.  
     
20.  
    document.getElementById("form").onsubmit = function () {
21.  
    if (!conn) {
22.  
    return false;
23.  
    }
24.  
    if (!msg.value) {
25.  
    return false;
26.  
    }
27.  
    conn.send(msg.value);
28.  
    msg.value = "";
29.  
    return false;
30.  
    };
31.  
     
32.  
    if (window["WebSocket"]) {//如果支持websockte就尝试连接
33.  
    //从浏览器的开发者工具里看一下ws的请求头
34.  
    conn = new WebSocket("ws://127.0.0.1:5656/chat");//请求跟websocket服务端建立连接(注意端口要一致)。关闭浏览器页面时会自动断开连接
35.  
    conn.onclose = function (evt) {
36.  
    var item = document.createElement("div")
37.  
    item.innerHTML = "<b>Connection closed.</b>";//连接关闭时打印一条信息
38.  
    appendLog(item);
39.  
    };
40.  
    conn.onmessage = function (evt) {//如果conn里有消息
41.  
    var messages = evt.data.split('\n');//用换行符分隔每条消息
42.  
    for (var i = 0; i < messages.length; i ) {
43.  
    var item = document.createElement("div");
44.  
    item.innerText = messages[i];//把消息逐条显示在屏幕上
45.  
    appendLog(item);
46.  
    }
47.  
    };
48.  
    } else {
49.  
    var item = document.createElement("div");
50.  
    item.innerHTML = "<b>Your browser does not support WebSockets.</b>";
51.  
    appendLog(item);
52.  
    }
53.  
    };
54.  
    </script>
55.  
    <style type="text/css">
56.  
    html {
57.  
    overflow: hidden;
58.  
    }
59.  
     
60.  
    body {
61.  
    overflow: hidden;
62.  
    padding: 0;
63.  
    margin: 0;
64.  
    width: 100%;
65.  
    height: 100%;
66.  
    background: gray;
67.  
    }
68.  
     
69.  
    #log {
70.  
    background: white;
71.  
    margin: 0;
72.  
    padding: 0.5em 0.5em 0.5em 0.5em;
73.  
    position: absolute;
74.  
    top: 0.5em;
75.  
    left: 0.5em;
76.  
    right: 0.5em;
77.  
    bottom: 3em;
78.  
    overflow: auto;
79.  
    }
80.  
     
81.  
    #form {
82.  
    padding: 0 0.5em 0 0.5em;
83.  
    margin: 0;
84.  
    position: absolute;
85.  
    bottom: 1em;
86.  
    left: 0px;
87.  
    width: 100%;
88.  
    overflow: hidden;
89.  
    }
90.  
    </style>
91.  
    </head>
92.  
     
93.  
    <body>
94.  
    <div id="log"></div>
95.  
    <form id="form">
96.  
    <input type="submit" value="发送" />
97.  
    <input type="text" id="msg" size="100" autofocus />
98.  
    </form>
99.  
    </body>
100.  
      
101.  
     </html>

起两个页面

这篇好文章是转载于：学新通技术网