GO websocket 实现简易聊天室

架构图如下：

本项目由四个文件组成：

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

Hub结构体实现：

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

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

 

Hub构造函数：

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

Client结构体实现：

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

type Client struct{
	hub *Hub
	conn *websocket.Conn
	send chan []byte
	name []byte
}

Hub工作实现：

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

func (hub *Hub) Run(){
	for{
		select{
		case client := <-hub.register:
			hub.clients[client] = struct{}{}
		case client := <- hub.unregister:
			delete(hub.clients,client)
			close(client.send)
		case msg := <-hub.broadcast:
			for client :=  range hub.clients{
				select{
				case client.send <- msg://如果管道不能立即写入数据，就认为该client出故障了
				default:
					close(client.send)
					delete(hub.clients, client)
				}
			}
		}

	}
}

Client从websocket读取内容：

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

const (
	writeWait  = 10 * time.Second //
	pongWait   = 60 * time.Second // 每60秒向websocket发送一次pong
	pingPeriod = 9 * pongWait / 10 //连接不断时每隔54秒向client发送一次ping
	maxMsgSize = 512 //消息的长度不能超过512
)
// 从websocket读取数据
func(client *Client)read(){
	defer func(){
		client.hub.unregister <- client  //向hub发送注销
		fmt.Printf("%s offline\n", client.name)
		fmt.Printf("close connection to %s\n",client.conn.RemoteAddr().String())
		client.conn.Close() //关闭ws连接
	}()

	// conn细节设置
	client.conn.SetReadLimit(maxMsgSize)
	client.conn.SetReadDeadline(time.Now().Add(pongWait)) //设置最长可读时间
	client.conn.SetPongHandler(func(appData string) error {
		client.conn.SetReadDeadline(time.Now().Add(pongWait))//每次接收到ping后都将最长可读时间延后60秒
		return nil
	})		

	for{
		_, p, err := client.conn.ReadMessage() //返回消息类型，消息，error
		if err != nil{
			//如果以意料之外的关闭状态关闭，就打印日志
			if websocket.IsUnexpectedCloseError(err, websocket.CloseAbnormalClosure, websocket.CloseGoingAway) {
				fmt.Printf("close websocket conn error: %v\n", err)
			}
			break //只要ReadMessage失败，就关闭websocket管道、注销client，退出
		}else{
			// trimspace:消去首尾空格， replace：将换行符换位空格，-1：全部转换
			message := bytes.TrimSpace(bytes.Replace(p,[]byte{'\n'}, []byte{' '}, -1))
			if len(client.name) == 0{//第一次输入的内容为自己的名字
				client.name = message
			}else{
			client.hub.broadcast<-bytes.Join([][]byte{client.name, message}, []byte(": "))
		}
	}
}

Client 从hub中接收信息：

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

func(client *Client)write() {
	ticker := time.NewTicker(pingPeriod)
	defer func(){
		ticker.Stop() //ticker不用就stop，防止协程泄漏
		fmt.Printf("close connection to %s\n", client.conn.RemoteAddr().String())
		client.conn.Close() //给前端写数据失败，就可以关系连接了
	}()
	for{
		select{
		case msg, ok := <-client.send:
			if !ok{
				fmt.Println("管道已经关闭")
				client.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			client.conn.SetWriteDeadline(time.Now().Add(writeWait))10秒内必须把信息写给前端（写到websocket连接里去），否则就关闭连接
			if writer, err := client.conn.NextWriter(websocket.TextMessage); err != nil{
				return 
			}else{
				writer.Write(msg)
				writer.Write([]byte{'\n'})
				// 有消息一次全写出去
				n := len(client.send)
				for i := 0; i < n; i++ {
					writer.Write(<-client.send)
					writer.Write([]byte{'\n'})
				}
				if err := writer.Close(); err != nil { //必须调close，否则下次调用client.conn.NextWriter时本条消息才会发送给浏览器
					return //结束一切
				}
			}
		case <-ticker.C:
			client.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := client.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil{
				return
			}
		}
	}
}

主页面路由设置:

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

websocket服务：

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

main函数：

func main(){
	hub := NewHub()
	go hub.Run()
	http.HandleFunc("/", ServeHome)
	http.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) {
		ServeWS(hub, w, r)
	})
	if err := http.ListenAndServe(":5656",nil);err != nil{
		fmt.Printf("start http service error: %s\n", err)
	}
}

前端实现：直接套模板 

起两个页面