ArmedPolice/app/service/websocket.go

package service

import (
	"ArmedPolice/app/logic"
	"ArmedPolice/serve/logger"
	"github.com/gorilla/websocket"
)

type Websocket struct {
	ID   string
	conn *websocket.Conn
	send chan []byte
}

type Hub struct {
	clients    map[string]*Websocket
	broadcast  chan logic.INotice
	emit       chan *HubEmit
	register   chan *Websocket
	unregister chan *Websocket
}

type HubEmit struct {
	ID  string
	Msg logic.INotice
}

var HubMessage *Hub

func (this *Hub) Run() {
	for {
		select {
		case client := <-this.register:
			this.clients[client.ID] = client
			logger.InfoF("客户端【%s】发起链接", client.ID)
		case client := <-this.unregister:
			if _, ok := this.clients[client.ID]; ok {
				delete(this.clients, client.ID)
				close(client.send)
			}
		case message := <-this.broadcast:
			for _, client := range this.clients {
				select {
				case client.send <- message.ToBytes():
				default:
					close(client.send)
					delete(this.clients, client.ID)
				}
			}
		case iMsg := <-this.emit:
			client, has := this.clients[iMsg.ID]

			if has {
				client.send <- iMsg.Msg.ToBytes()
			}
		}
	}
}

func (this *Hub) EmitHandle(iMsge *HubEmit) {
	this.emit <- iMsge
}

func (this *Hub) BroadcastHandle(msg logic.INotice) {
	this.broadcast <- msg
}

func (this *Hub) RegisterHandle(ws *Websocket) {
	this.register <- ws
}

func (this *Hub) UnregisterHandle(ws *Websocket) {
	this.unregister <- ws
}

func NewHub() *Hub {
	HubMessage = &Hub{
		clients:    make(map[string]*Websocket),
		broadcast:  make(chan logic.INotice),
		emit:       make(chan *HubEmit),
		register:   make(chan *Websocket),
		unregister: make(chan *Websocket),
	}
	return HubMessage
}

func (this *Websocket) Write() {
	defer func() {
		this.conn.Close()
	}()
	for {
		select {
		case message, ok := <-this.send:
			if !ok {
				return
			}
			logger.InfoF("发送到客户端【%s】信息：%s", this.ID, string(message))
			_ = this.conn.WriteMessage(websocket.TextMessage, message)
		}
	}
}

func (this *Websocket) Read() {

}

func NewWebsocket(id string, conn *websocket.Conn) *Websocket {
	return &Websocket{
		ID:   id,
		conn: conn,
		send: make(chan []byte),
	}
}