1/trading-go/controller/chatcontroller.go

package controller

import (
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
	"strconv"
	"trading-go/model"
	"trading-go/response"
)

var UP = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

type client struct {
	conn *websocket.Conn
	pip  chan *model.Message
}

var Conns map[int64]*client

func init() {
	Conns = make(map[int64]*client)
}

// 发送消息
func send(conn *websocket.Conn, uid int64) {
	pip := Conns[uid].pip
	defer func() {
		close(pip)
		conn.Close()
		delete(Conns, uid)
	}()
	for {
		data := <-pip
		msg, err := json.Marshal(data)
		if err != nil {
			fmt.Println("link closed")
			break
		}
		err = conn.WriteMessage(1, msg)
		if err != nil {
			fmt.Println(err.Error())
			continue
		}
	}
}

// 接收消息
func reception(conn *websocket.Conn, uid int64) {
	pip := Conns[uid].pip
	defer func() {
		close(pip)
		conn.Close()
		delete(Conns, uid)
	}()
	for {
		var msg model.Message
		err := conn.ReadJSON(&msg)
		if err != nil {
			fmt.Println(err.Error())
			break
		}
		Conns[msg.To].pip <- &msg
	}
}

func Chat(w http.ResponseWriter, rq *http.Request, uid int64) {
	// 升级为websocket
	conn, err := UP.Upgrade(w, rq, nil)

	if err != nil {
		response.Fail(w, err.Error(), http.StatusUpgradeRequired)
		return
	}

	pip := make(chan *model.Message, 1024)
	client := client{
		conn: conn,
		pip:  pip,
	}
	Conns[uid] = &client

	go send(conn, uid)
	go reception(conn, uid)

	response.Success(w, "success", nil)
}

func LinkToServer(c *gin.Context) {
	uid, err := strconv.Atoi(c.Query("uid"))
	if err != nil {
		response.Fail(c.Writer, "failed", 500)
		return
	}
	Chat(c.Writer, c.Request, int64(uid))
}