go语言从零学起(四) -- 基于martini和gorilla实现的websocket聊天实例 ...

OStack程序员社区-中国程序员成长平台 › 门户 › 编程› Go›Go教程

原作者: [db:作者] 来自: [db:来源] 收藏邀请

如果只是想了解chat的实现方式，在gorilla和revel框架里面都有完整的chat实例可以提供参考。本篇讲解的是，如何基于martini实现websocket的聊天。

配置步骤:

1 已经安装了go命令，配置了GOPATH

2 安装gorilla/websocket

go get github.com/gorilla/websocket

3 安装martini

go get github.com/go-martini/martini

项目文件列表

$GOPATH/demo/home.html
$GOPATH/demo/martini_ws.go

home.html 代码

<!DOCTYPE html>
<html lang="en">
<head>
<title>Chat Example</title>
<script type="text/javascript">
window.onload = function () {
    var conn;
    var msg = document.getElementById("msg");
    var log = document.getElementById("log");

    function appendLog(item) {
        var doScroll = log.scrollTop === log.scrollHeight - log.clientHeight;
        log.appendChild(item);
        if (doScroll) {
            log.scrollTop = log.scrollHeight - log.clientHeight;
        }
    }

    document.getElementById("form").onsubmit = function () {
        if (!conn) {
            return false;
        }
        if (!msg.value) {
            return false;
        }
        conn.send(msg.value);
        msg.value = "";
        return false;
    };

    if (window["WebSocket"]) {
        conn = new WebSocket("ws://{{$}}/ws");
        conn.onclose = function (evt) {
            var item = document.createElement("div");
            item.innerHTML = "<b>Connection closed.</b>";
            appendLog(item);
        };
        conn.onmessage = function (evt) {
            var messages = evt.data.split('\n');
            for (var i = 0; i < messages.length; i++) {
                var item = document.createElement("div");
                item.innerText = messages[i];
                appendLog(item);
            }
        };
    } else {
        var item = document.createElement("div");
        item.innerHTML = "<b>Your browser does not support WebSockets.</b>";
        appendLog(item);
    }
};
</script>
<style type="text/css">
html {
    overflow: hidden;
}

body {
    overflow: hidden;
    padding: 0;
    margin: 0;
    width: 100%;
    height: 100%;
    background: gray;
}

#log {
    background: white;
    margin: 0;
    padding: 0.5em 0.5em 0.5em 0.5em;
    position: absolute;
    top: 0.5em;
    left: 0.5em;
    right: 0.5em;
    bottom: 3em;
    overflow: auto;
}

#form {
    padding: 0 0.5em 0 0.5em;
    margin: 0;
    position: absolute;
    bottom: 1em;
    left: 0px;
    width: 100%;
    overflow: hidden;
}

</style>
</head>
<body>
<div id="log"></div>
<form id="form">
    <input type="submit" value="Send" />
    <input type="text" id="msg" size="64"/>
</form>
</body>
</html>

martini_ws.go代码

package main

import (
    "fmt"
    "github.com/go-martini/martini"
    "github.com/gorilla/websocket"
    "log"
    "net/http"
    "text/template"
)

const (
    readBufferSize  = 1024
    writeBufferSize = 1024
)

type Client struct {
    conn     *websocket.Conn
    messages chan []byte
}

var clients map[*Client]bool // 存储所有的链接
var broadcast chan []byte    // 广播聊天的chan
var addClients chan *Client  // 新链接进来的chan

func (c *Client) readPump() {
    for {
        _, message, err := c.conn.ReadMessage()
        if err != nil {
            if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
                log.Printf("error: %v", err)
            }
            break
        }
        fmt.Printf("receive message is :%s\n", message)
        broadcast <- message
    }
}

func (c *Client) writePump() {
    for {
        select {
        case message := <-c.messages:
            fmt.Printf("send message is :%s\n", message)
            c.conn.WriteMessage(1, message)
        }
    }
}

func manager() {

    clients = make(map[*Client]bool)
    broadcast = make(chan []byte, 10)
    addClients = make(chan *Client)

    for {
        select {
        case message := <-broadcast:
            for client := range clients {
                select {
                case client.messages <- message:
                default:
                    close(client.messages)
                    delete(clients, client)
                }
            }
        case itemClient := <-addClients:
            clients[itemClient] = true
        }
    }
}

func main() {

    var homeTemplate = template.Must(template.ParseFiles("home.html"))

    m := martini.Classic()

    m.Get("/", func(res http.ResponseWriter, req *http.Request) {

        res.Header().Set("Content-Type", "text/html; charset=utf-8")
        homeTemplate.Execute(res, req.Host)
    })

    m.Get("/ws", func(res http.ResponseWriter, req *http.Request) { // res and req are injected by Martini
        conn, err := websocket.Upgrade(res, req, nil, readBufferSize, writeBufferSize)
        if err != nil {
            log.Println(err)
            return
        }
        client := &Client{conn: conn, messages: make(chan []byte, 5)}
        addClients <- client
        go client.writePump()
        client.readPump()
    })

    go manager()

    m.RunOnAddr(":3010")
    //m.Run()

}