go websocket（copy直接可用）

主要依賴
 

"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"

 websocket.go 

package api

import (
	"errors"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/outakujo/utils"
	"log"
	"net/http"
	"sync"
)

type TextSocketSession struct {
	c  *websocket.Conn
	Id string
}

func (r *TextSocketSession) Send(data string) error {
	return r.c.WriteMessage(websocket.TextMessage, []byte(data))
}

func (r *TextSocketSession) SendJson(data interface{}) error {
	return r.c.WriteJSON(data)
}

type TextSocketSessionSet struct {
	mux sync.RWMutex
	set map[string]*TextSocketSession
}

var SocketSessionMiss = errors.New("socket session miss")

func (on *TextSocketSessionSet) add(s *TextSocketSession) {
	on.mux.Lock()
	on.set[s.Id] = s
	on.mux.Unlock()
}

func (on *TextSocketSessionSet) Get(id string) (s *TextSocketSession, err error) {
	on.mux.RLock()
	defer on.mux.RUnlock()
	s, ok := on.set[id]
	if !ok {
		err = SocketSessionMiss
	}
	return
}

func (on *TextSocketSessionSet) delete(id string) *TextSocketSession {
	s, err := on.Get(id)
	if err == nil {
		on.mux.Lock()
		delete(on.set, id)
		on.mux.Unlock()
		return s
	}
	return nil
}

func (on *TextSocketSessionSet) Counts() int {
	on.mux.RLock()
	defer on.mux.RUnlock()
	return len(on.set)
}

type OnReceive func(id string, message []byte)

type OnConnectOpen func(id string)

type OnConnectClose func(id string)


//gin config url path must be  /../xxx:id

func TextSocket(receive OnReceive, open OnConnectOpen, clos OnConnectClose) (hand gin.HandlerFunc, set *TextSocketSessionSet) {
	var up = websocket.Upgrader{CheckOrigin: func(r *http.Request) bool {
		return true
	}}
	set = &TextSocketSessionSet{set: map[string]*TextSocketSession{}}
	hand = func(c *gin.Context) {
		var mid = struct {
			Id string `uri:"id" binding:"required"`
		}{}
		e := c.ShouldBindUri(&mid)
		utils.PanicError(e)
		conn, e := up.Upgrade(c.Writer, c.Request, nil)
		utils.PanicError(e)
		session := &TextSocketSession{c: conn, Id: mid.Id}
		set.add(session)
		open(mid.Id)
		defer func() {
			if e := recover(); e != nil {
				log.Println(e)
			}
			e := conn.Close()
			if e != nil {
				log.Println(e)
			}
			set.delete(mid.Id)
			clos(mid.Id)
		}()
		for {
			_, message, err := conn.ReadMessage()
			utils.PanicError(err)
			receive(mid.Id, message)
		}
	}
	return
}

 使用

var soks *TextSocketSessionSet

func ApiFunc(r gin.IRouter) {
	var sochand gin.HandlerFunc
	sochand, soks = TextSocket(rece, open, clos)
	r.GET("socket/:id", sochand)
}

func rece(id string, mes []byte) {
	fmt.Println(id, " ", string(mes))
}

func open(id string) {
	fmt.Println(id, " open")
}

func clos(id string) {
	fmt.Println(id, " close")
}