具体可以参考GitHub上的代码
package websocket
import (
"bytes"
"github.com/gorilla/websocket"
"log"
"net/http"
"time"
)
var (
newline = []byte{'\n'}
space = []byte{' '}
// 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,
}
//定义-》 只有一个组里面的 client 才能互相聊天,
type WsClientGroup struct {
broadcast chan []byte //用于广播数据【广播协程监听这个通道】
clientEnter chan *WsClient
clientExit chan *WsClient
clients map[*WsClient]bool
}
func (g *WsClientGroup) HandleRun() {
//eventLoop:
log.Printf("handle run")
for {
select {
case client := <-g.clientEnter:
//注册用户
g.clients[client] = true
case client := <-g.clientExit:
//用户退出
if _, ok := g.clients[client]; ok {
delete(g.clients, client)
close(client.send)
}
case broadcastMsg := <-g.broadcast:
log.Printf("broadcastMsg, %s",broadcastMsg)
for cli := range g.clients {
select {
case cli.send <- broadcastMsg:
default:
//这种情况下,只能说明 cli.send == nil
close(cli.send)
delete(g.clients, cli)
}
}
}
}
}
func NewClientGroup() *WsClientGroup {
return &WsClientGroup{
broadcast: make(chan []byte),
clientEnter: make(chan *WsClient),
clientExit: make(chan *WsClient),
clients: make(map[*WsClient]bool),
}
}
type WsClient struct {
send chan []byte
conn *websocket.Conn
Groups *WsClientGroup
}
//开启读协程
func (cli *WsClient) ReadLoopGroup() {
defer func() {
cli.Groups.clientExit <- cli
// exit and close connection
cli.conn.Close()
}()
cli.conn.SetReadLimit(int64(maxMessageSize))
cli.conn.SetReadDeadline(time.Now().Add(pongWait))
cli.conn.SetPongHandler(func(s string) error {
//续期
cli.conn.SetReadDeadline(time.Now().Add(pongWait))
return nil
})
log.Printf("begin read loop ")
readEventLoop:
for {
_, msg, err := cli.conn.ReadMessage()
log.Printf("begin read msg")
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
log.Printf("websocket client error: %+v", err)
}
break readEventLoop
}
msg = bytes.TrimSpace(msg)
log.Printf("receive client msg = [%s]", msg)
// 这里可以广播给其他用户, 前端可以传个 type,后端 根据 type 判断是广播还是私聊
//假设这里用广播
cli.Groups.broadcast <- msg
}
}
func (cli *WsClient) WriteLoopGroup() {
ticker := time.NewTicker(pingPeriod)
defer func() {
log.Printf("exit writeLoopGroup")
ticker.Stop()
cli.conn.Close()
}()
writeEventLoop:
for {
select {
case msg, ok := <-cli.send:
if !ok {
cli.conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
w, err := cli.conn.NextWriter(websocket.TextMessage)
if err != nil {
break writeEventLoop
}
w.Write(msg)
n := len(cli.send)
for i := 0; i < n; i++ {
//继续发送
w.Write(newline)
w.Write(<-cli.send)
}
if err = w.Close(); err != nil {
break writeEventLoop
}
case <-ticker.C:
cli.conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := cli.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
//write errror
//超时 异常 ,主动退出
break writeEventLoop
}
}
}
}
func Register(groups *WsClientGroup, w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Printf("error info %+v", err)
return
}
client := &WsClient{
Groups: groups,
send: make(chan []byte, 256),
conn: conn,
}
groups.clientEnter <- client
log.Printf("enter client ")
//groups.HandleRun()
go client.ReadLoopGroup()
go client.WriteLoopGroup()
}
package main
import (
"log"
"net/http"
"websocket/websocket"
)
func serverHome(w http.ResponseWriter, r *http.Request) {
log.Println(r.URL)
if r.URL.Path != "/" {
http.Error(w, "Not found", http.StatusNotFound)
return
}
if r.Method != http.MethodGet {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
http.ServeFile(w, r, "index.html")
}
func main() {
http.HandleFunc("/", serverHome)
var groups = websocket.NewClientGroup()
go groups.HandleRun()
http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
//这里要区分不同的聊天室的话, 可能还需要加一个 hashMap, hashMap.put("chat_room_id",groups)
websocket.Register(groups, w, r)
})
http.ListenAndServe(":8080", nil)
}
DOCTYPE html>
<html lang="en">
<head>
<title>Chat Exampletitle>
<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 - 1;
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://" + document.location.host + "/ws");
conn.onclose = function (evt) {
var item = document.createElement("div");
item.innerHTML = "Connection closed.";
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 = "Your browser does not support WebSockets.";
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" autofocus />
form>
body>
html>