golang之websocket

websocket是啥

websocket是socket连接和http协议的结合体，可以实现网页和服务端的长连接

通信流程

拉取websocket库

go get github.com/gorilla/websocket

简单示例：

package main

import (
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
)

var (
	Upgrader = websocket.Upgrader{
	//允许跨域
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
)

func Handle(ctx *gin.Context) {
	wsConn,_ := Upgrader.Upgrade(ctx.Writer,ctx.Request,nil)
	_,data,_ := wsConn.ReadMessage()
	wsConn.WriteMessage(websocket.TextMessage,data)
}

func main() {
	router := gin.Default()
	router.GET("/ws", Handle)
	router.Run(":8989")
}

测试：

socket或websocket连接的封装

在查阅他人的代码的时候，发现只要是socket连接，都进行了封装，一种说法是为了线程安全，一种说法是为了业务的拓展性（比如说在发之前过滤一些消息）
封装的代码如下：

package core

import "github.com/gorilla/websocket"

type Connection struct {
	Conn     *websocket.Conn
	inchan   chan []byte
	outchan  chan []byte
	exitchan chan bool
	isClose  bool
}

func NewConnection(conn *websocket.Conn) (*Connection, error) {
	wsConn := &Connection{
		Conn:     conn,
		inchan:   make(chan []byte, 1000),
		outchan:  make(chan []byte, 1000),
		exitchan: make(chan bool, 1),
	}
	return wsConn, nil
}

func (conn *Connection) Start() {
	// 不停的读长连接的消息，放到inchan中
	go conn.readLoop()
	// 不停的从outchan,读取消息，发到客户端
	go conn.writeLoop()
}

func (conn *Connection) ReadMessage() ([]byte, err error){
	if isClose {
		return nil,errors.New("connection is closed")
	}
	data = <-conn.inchan
	return data,nil
}
func (conn *Connection) WriteMessage(data []byte) error {
	if isClose {
		return errors.New("connection is closed")
	}
	conn.outchan <- data
	return nil
}

func (conn *Connection) Close() {
	if conn.isClose == true {
		return
	}
	conn.isClose = true
	// 线程安全，可重入的close
	conn.Conn.Close()
	close(conn.exitchan)
	close(conn.inchan)
	close(conn.outchan)
}

func (conn *Connection) readLoop() {
	defer conn.Close()
	for {
		_, data, err := conn.Conn.ReadMessage()
		if err != nil {
			conn.exitchan <- true
			conn.isClose = true
			return
		}
		conn.inchan <- data
	}
}

func (conn *Connection) writeLoop() {
	for {
		select {
		case data := <-conn.outchan:
			if err := conn.Conn.WriteMessage(websocket.TextMessage, data); err != nil {
				conn.exitchan <- true
				return
			}
		case <-conn.exitchan:
			return
		}
	}
}

测试：

package main

import (
	"awesomeProject1/wsocket/core"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
	"time"
)

var (
	Upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
)

func Handle(ctx *gin.Context) {
	conn,_ := Upgrader.Upgrade(ctx.Writer,ctx.Request,nil)
	wsConn,_ := core.NewConnection(conn)
	wsConn.Start()
	go func() {
		for {
			time.Sleep(time.Second)
			wsConn.WriteMessage([]byte("server"))
		}
	}()
	for {
		msg,_ := wsConn.ReadMessage()
		wsConn.WriteMessage(msg)
	}
}

func main() {
	router := gin.Default()
	router.GET("/ws", Handle)
	router.Run(":8989")
}