golang实现简单的websocket【golang】【websocket】

具体可以参考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)
}