框架慢慢修改

controllers/ClientManager.go

@@ -0,0 +1,113 @@
+// Package ws is to define a websocket server and client connect.
+// Author: Arthur Zhang
+// Create Date: 20190101
+package controllers
+
+import (
+	"encoding/json"
+
+	"github.com/gorilla/websocket"
+)
+
+// ClientManager is a websocket manager
+type ClientManager struct {
+	Clients    map[*Client]bool
+	Broadcast  chan []byte
+	Register   chan *Client
+	Unregister chan *Client
+}
+
+// Client is a websocket client
+type Client struct {
+	ID     string
+	Socket *websocket.Conn
+	Send   chan []byte
+}
+
+// Message is an object for websocket message which is mapped to json type
+type Message struct {
+	Sender    string `json:"sender,omitempty"`
+	Recipient string `json:"recipient,omitempty"`
+	Content   string `json:"content,omitempty"`
+}
+
+// Manager define a ws server manager
+var Manager = ClientManager{
+	Broadcast:  make(chan []byte),
+	Register:   make(chan *Client),
+	Unregister: make(chan *Client),
+	Clients:    make(map[*Client]bool),
+}
+
+// Start is to start a ws server
+func (manager *ClientManager) Start() {
+	for {
+		select {
+		case conn := <-manager.Register:
+			manager.Clients[conn] = true
+			jsonMessage, _ := json.Marshal(&Message{Content: "/A new socket has connected."})
+			manager.Send(jsonMessage, conn)
+		case conn := <-manager.Unregister:
+			if _, ok := manager.Clients[conn]; ok {
+				close(conn.Send)
+				delete(manager.Clients, conn)
+				jsonMessage, _ := json.Marshal(&Message{Content: "/A socket has disconnected."})
+				manager.Send(jsonMessage, conn)
+			}
+		case message := <-manager.Broadcast:
+			for conn := range manager.Clients {
+				select {
+				case conn.Send <- message:
+				default:
+					close(conn.Send)
+					delete(manager.Clients, conn)
+				}
+			}
+		}
+	}
+}
+
+// Send is to send ws message to ws client
+func (manager *ClientManager) Send(message []byte, ignore *Client) {
+	for conn := range manager.Clients {
+		if conn != ignore {
+			conn.Send <- message
+		}
+	}
+}
+
+func (c *Client) Read() {
+	defer func() {
+		Manager.Unregister <- c
+		c.Socket.Close()
+	}()
+
+	for {
+		_, message, err := c.Socket.ReadMessage()
+		if err != nil {
+			Manager.Unregister <- c
+			c.Socket.Close()
+			break
+		}
+		jsonMessage, _ := json.Marshal(&Message{Sender: c.ID, Content: string(message)})
+		Manager.Broadcast <- jsonMessage
+	}
+}
+
+func (c *Client) Write() {
+	defer func() {
+		c.Socket.Close()
+	}()
+
+	for {
+		select {
+		case message, ok := <-c.Send:
+			if !ok {
+				c.Socket.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+
+			c.Socket.WriteMessage(websocket.TextMessage, message)
+		}
+	}
+}