Golang网络编程（一）Websocket编程

背景

​       在服务端与前端或者sdk端通过http进行数据交互的时候，常常会面临一些问题：比如在im场景下，前端需要主动push数据给服务端；比如在高并发场景下，sdk端请求服务端压力过大导致服务器文件句柄数不够前端无法继续请求、服务端开启goroutine过多内存消耗过大等问题；websocket是一种应用层通信协议，建立在TCP协议之上，具有数据格式轻量、通信双方均可以推送消息等特点。

使用

package main

import (
	"context"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/pkg/errors"
	"log"
	"net/http"
	"net/url"
	"sync"
	"time"
)

const (
	// client server服务地址
	addr = "127.0.0.1:1198"
	// 请求次数
	requestTotal = 20
	// websocket路由地址
	connectRoute = "/connect"
)

type Map struct {
	mutex *sync.RWMutex
	mp    map[string]chan string
}

func NewMap() Map {
	return Map{
		mutex: new(sync.RWMutex),
		mp:    make(map[string]chan string),
	}
}

type client struct {
	ctx    context.Context
	cancel context.CancelFunc

	addr string
	conn *websocket.Conn
}

func newClient(addr string) (c *client, err error) {
	u := url.URL{
		Scheme: "ws",
		Host:   addr,
		Path:   connectRoute,
	}
	var conn *websocket.Conn
	conn, _, err = websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		err = errors.Wrap(err, "failed to dial connect")
		return
	}
	ctx, cancel := context.WithCancel(context.Background())
	c = &client{
		ctx:    ctx,
		cancel: cancel,
		addr:   addr,
		conn:   conn,
	}
	return
}

func (c *client) close() {
	c.cancel()
}

func (c *client) handleRequest(reqMsg requestMsg) (resp responseMsg, err error) {
	if err = c.conn.WriteJSON(&reqMsg); err != nil {
		log.Println("failed to write request message to server websocket")
		return
	}
	if err = c.conn.ReadJSON(&resp); err != nil {
		log.Println("failed to read response message from server websocket")
		return
	}
	return
}

type requestMsg struct {
	Name string
}

type connReqMsg struct {
	ID string
	requestMsg
}

type responseMsg struct {
	Data interface{}
}

type connRespMsg struct {
	ID string
	responseMsg
}

type server struct {
	ctx    context.Context
	cancel context.CancelFunc

	addr   string
	engine *gin.Engine
}

func newServer(addr string) *server {
	ctx, cancel := context.WithCancel(context.Background())
	var s = &server{
		ctx:    ctx,
		cancel: cancel,
		addr:   addr,
		engine: gin.Default(),
	}
	gin.SetMode(gin.ReleaseMode)
	s.engine.GET(connectRoute, connectHandle)
	return s
}

func (s *server) close() {
	s.cancel()
}

func (s *server) run() error {
	return s.engine.Run(s.addr)
}

var (
	upGrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
)

func connectHandle(c *gin.Context) {
	// 建立websocket连接
	conn, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Println("connect failed")
		return
	}
	defer func() {
		conn.Close()
		log.Println("connect close")
	}()
	for {
		// 读取请求
		var req connReqMsg
		if err = conn.ReadJSON(&req); err != nil {
			log.Println(err.Error())
			break
		}
		// 处理请求
		var resp connRespMsg
		resp.Data = req.Name
		time.Sleep(2 * time.Second)
		fmt.Println("handle request ", req)
		// 返回结果
		if err = conn.WriteJSON(&resp); err != nil {
			log.Println(err.Error())
			break
		}
	}
}

func main() {
	var wg sync.WaitGroup

	s := newServer(addr)
	wg.Add(1)
	go func() {
		if err := s.run(); err != nil {
			fmt.Println(err)
		}
		wg.Done()
	}()
	defer s.close()

	c, err := newClient(addr)
	if err != nil {
		panic(err)
	}
	defer c.close()

	for i := 1; i <= requestTotal; i++ {
		var reqMsg = requestMsg{Name: fmt.Sprintf("name-%v", i)}
		var respMsg responseMsg
		respMsg, err = c.handleRequest(reqMsg)
		if err != nil {
			fmt.Println(reqMsg.Name, err)
			return
		}
		fmt.Println(respMsg.Data)
	}
	wg.Wait()
}

注意：在同一个websocket连接中，数据的发送存在并发问题，如果需要并发处理，那么需要在后台开启两个goroutine分别进行接收请求和返回结果的处理