WebSocket Ping Pong Heartbeat

package main

import (
	"log"
	"time"

	"github.com/gorilla/websocket"
)

const (
	pongWait   = 60 * time.Second
	pingPeriod = 50 * time.Second // 必须小于 pongWait
)

type Peer struct {
	conn *websocket.Conn
}

func (p *Peer) readPump() {
	defer p.conn.Close()
	p.conn.SetReadLimit(512)
	p.conn.SetReadDeadline(time.Now().Add(pongWait))
	p.conn.SetPongHandler(func(string) error {
		return p.conn.SetReadDeadline(time.Now().Add(pongWait))
	})

	for {
		if _, _, err := p.conn.ReadMessage(); err != nil {
			log.Println("read closed:", err)
			return
		}
	}
}

func (p *Peer) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		p.conn.Close()
	}()

	for range ticker.C {
		if err := p.conn.WriteControl(
			websocket.PingMessage, nil, time.Now().Add(10*time.Second),
		); err != nil {
			log.Println("ping failed:", err)
			return
		}
		log.Println("heartbeat sent")
	}
}

func main() {
	// 客户端侧同样适用：拨号成功后启动两个泵
	// conn, _, err := websocket.DefaultDialer.Dial("wss://example.com/ws", nil)

	conn, _, _ := websocket.DefaultDialer.Dial("ws://localhost:8080/ws", nil)
	peer := &Peer{conn: conn}

	go peer.writePump()
	peer.readPump() // 阻塞直到任一方向断开
}