Golang服务端断线重连

断线重连的逻辑很简单，就是把用户存到服务器内存中，当客户端再次登录的时候，判断内存中是否有用户的值，有的话替换

package main

import (
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"sync"
	"time"
)

type Client struct {
	conn          *websocket.Conn
	id            string
	send          chan []byte
	lastHeartbeat time.Time
}

var (
	clients = make(map[string]*Client)
	mutex   sync.RWMutex
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		// 自定义的检查逻辑，可以根据需要进行修改
		// 这里简单地返回 true 允许所有来源
		return true
	},
}

func main() {
	http.HandleFunc("/ws", handleWebSocket)
	http.ListenAndServe(":8080", nil)
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Upgrade error:", err)
		return
	}

	id := r.URL.Query().Get("id")
	fmt.Println(id)
	mutex.RLock()
	//判断是否是断线重连
	if _, ok := clients[id]; ok {
		clients[id].conn = conn
		client := clients[id]
		go client.readPump()
		go client.writePump()
		fmt.Printf("断线连接成功%v", client)
	} else {
		client := &Client{
			conn:          conn,
			id:            id,
			send:          make(chan []byte),
			lastHeartbeat: time.Now(),
		}
		fmt.Printf("连接成功%v", client)
		clients[id] = client
		go client.readPump()
		go client.writePump()
	}
	mutex.RUnlock()
	// 检查心跳
	go checkHeartbeat(id)
}

func (client *Client) readPump() {
	defer func() {
		mutex.Lock()
		delete(clients, client.id)
		mutex.Unlock()
		client.conn.Close()
	}()

	for {
		_, message, err := client.conn.ReadMessage()
		if err != nil {
			if !websocket.IsCloseError(err, websocket.CloseNormalClosure) {
				log.Println("Read error:", err)
			}
			break
		}
		// 处理收到的消息
		log.Printf("Received message from client %s: %s\n", client.id, string(message))
		// 更新最后收到心跳的时间
		client.lastHeartbeat = time.Now()
	}
}

func (client *Client) writePump() {
	defer func() {
		client.conn.Close()
	}()

	for {
		select {
		case message, ok := <-client.send:
			if !ok {
				err := client.conn.WriteMessage(websocket.CloseMessage, []byte{})
				if err != nil {
					log.Println("Write error:", err)
				}
				return
			}
			err := client.conn.WriteMessage(websocket.TextMessage, message)
			if err != nil {
				log.Println("Write error:", err)
				return
			}
		}
	}
}

func checkHeartbeat(id string) {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if client, ok := clients[id]; ok {
				// 检查最后心跳时间，超过指定时间则踢掉客户端
				//if time.Since(client.lastHeartbeat) > 10*time.Second {
				//	client.send <- []byte("heartbeat timeout, disconnecting")
				//	close(client.send)
				//	return
				//}

				// 发送心跳消息
				client.send <- []byte("heartbeat")
			} else {
				return
			}
		}
	}
}