gin结合vue封装WebSocket

Gin框架集成封装WebSocket

gin默认没有集成websocket，我们借用github.com/gorilla/websocket这个库来对gin进行封装

封装Handler

func WebSocketHandlerFunc(handler func(ctx *gin.Context, coon *websocket.Conn)) gin.HandlerFunc {
	upGrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	return func(context *gin.Context) {
		if coon, err := upGrader.Upgrade(context.Writer, context.Request, nil); err == nil {
			handler(context, coon)
		}
	}
}

使用装饰器模式将我们的处理函数封装成gin的HandlerFunc返回

直接使用

func handler(ctx *gin.Context, coon *websocket.Conn) {
    coon.Close()
}

engine.GET("/system/ws", WebSocketHandlerFunc(handler)

封装WebSocket

分析

1、一般我们用websocket都是用来接收信息和发送信息，但是接收信息的操作是阻塞的，当我们收不到消息时，收到消息下面的操作就无法执行，所以我们不要在一个协程里发送和读取信息。

2、收发消息如果分别用一个协程处理，那么我们应该要控制这两个协程什么时候结束。

3、收发消息都应该是一个无限循环的过程。

针对上面的分析，我想了一下，可以用下面的方式解决

type WebSocketHandler interface {
	Read(ws *WebSocketExecutor)
	Write(ws *WebSocketExecutor)
}

type WebSocketExecutor struct {
	Context context.Context
	Coon    *websocket.Conn
	cancel  context.CancelFunc
	Handler WebSocketHandler
}

// 结束当前的websocket连接
func (w *WebSocketExecutor) Done() {
	w.cancel()
}

func (w *WebSocketExecutor) ListenAndServe() {
    // 关闭连接
    defer w.Coon.Close()
    // 分别开启两个协程
	go w.Handler.Read(w)
	go w.Handler.Write(w)
    // 监听是否应该结束当前的连接
	w.listen()
}

func (w *WebSocketExecutor) listen() {
	select {
	case <-w.Context.Done():
		break
	}
}

func NewWebSocketExecutor(coon *websocket.Conn, handler WebSocketHandler) *WebSocketExecutor {
	ctx, cancel := context.WithCancel(context.Background())
	return &WebSocketExecutor{Context: ctx, Coon: coon, cancel: cancel, Handler: handler}
}

1、首先控制协程的关闭可以用go的官方库context来做，这里采用的是context.WithCancel

2、因为读写的操作是动态的，所以将读写的操作抽象成接口，让开发者自己处理

3、定义一个WebSocketExecutor的接口体，在NewWebSocketExecutor的构造方法中，只需要传入websocket的连接对象和自定义实现了WebSocketHandler接口对象即可。

4、最后调用WebSocketExecutor的ListenAndServe方法即可，它会调用初始化时传进来的WebSocketHandler的Read和Write方法，并且阻塞当前的主协程，开发者可以在自己实现的WebSocketHandler方法中自行关闭。

最终体验版

package main


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

func WebSocketHandlerFunc(handler func(ctx *gin.Context, coon *websocket.Conn)) gin.HandlerFunc {
	upGrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	return func(context *gin.Context) {
		if coon, err := upGrader.Upgrade(context.Writer, context.Request, nil); err == nil {
			handler(context, coon)
		}
	}
}

type WebSocketHandler interface {
	Read(ws *WebSocketExecutor)
	Write(ws *WebSocketExecutor)
}

type WebSocketExecutor struct {
	Context context.Context
	Coon    *websocket.Conn
	cancel  context.CancelFunc
	Handler WebSocketHandler
}

func (w *WebSocketExecutor) Done() {
	w.cancel()
}

func (w *WebSocketExecutor) ListenAndServe() {
	defer w.Coon.Close()
	go w.Handler.Read(w)
	go w.Handler.Write(w)
	w.listen()
}

func (w *WebSocketExecutor) listen() {
	select {
	case <-w.Context.Done():
		break
	}
}

func NewWebSocketExecutor(coon *websocket.Conn, handler WebSocketHandler) *WebSocketExecutor {
	ctx, cancel := context.WithCancel(context.Background())
	return &WebSocketExecutor{Context: ctx, Coon: coon, cancel: cancel, Handler: handler}
}



/*
自定测试的例子
*/
type systemWebSocketHandler struct{}

func (s systemWebSocketHandler) Read(ws *WebSocketExecutor) {
LOOP:
	for {
		if msgType, data, err := ws.Coon.ReadMessage(); msgType == websocket.CloseMessage || err != nil {
			ws.Done()
			break LOOP
		} else {
			log.Println(string(data))
		}
	}
}

func (s systemWebSocketHandler) Write(ws *WebSocketExecutor) {
LOOP:
	for {
		select {
		case <-ws.Context.Done():
			ws.Done()
			break LOOP
		default:
			if err := ws.Coon.WriteJSON("666"); err != nil {
				ws.Done()
				break LOOP
			}
			time.Sleep(time.Second)
		}
	}
}


// 处理函数
func SysTemWebSocket(ctx *gin.Context, coon *websocket.Conn) {
	handler := new(systemWebSocketHandler)
	executor := NewWebSocketExecutor(coon, handler)
	executor.ListenAndServe()
}



// 主函数
func main() {
    engine := gin.Default()
    engine.GET("/ws", WebSocketHandlerFunc(SysTemWebSocket))
    engine.Run()
}

结合VUE测试

<template>
  <div>

  </div>
</template>

<script>
export default {
  name: "System",
  data() {
    return {
      websocket: null
    }
  },
  methods: {
    initWebSocket() {
      const ws = "ws://127.0.0.1:8080/ws"
      this.websocket = new WebSocket(ws)
      this.websocket.onopen = this.onOpenWebsocket
      this.websocket.onerror = this.onErrorWebsocket
      this.websocket.onclose = this.onCloseWebsocket
      this.websocket.onmessage = this.onMessageWebsocket
    },
    onOpenWebsocket() {
      this.$message.success("连接成功!")
    },
    onErrorWebsocket() {
      this.$message.error("连接失败!")
    },
    onCloseWebsocket() {
      this.$message.info("连接关闭!")
    },
    onMessageWebsocket(message) {
      console.log(message)
    }
  },
  created() {
    this.initWebSocket()
  },
  destroyed() {
    console.log("close")
    this.websocket.close()
  }
}
</script>

<style scoped>

</style>