使用 Go 實作「分散式聊天室」應用程式

前言

本文使用 WebSocket 作為前後端訊息交換的方式，並使用 Redis 的 Pub/Sub 來分享後端在不同節點之間的訊息。

做法

新增 main.go 檔。

package main

import (
	"context"
	"embed"
	"github.com/gorilla/mux"
	_ "github.com/joho/godotenv/autoload"
	"log"
	"net/http"
	"os"
	"os/signal"
	"time"
)

var (
	//go:embed index.html
	view embed.FS
)

func main() {
	quit := make(chan struct{})
	r := mux.NewRouter()
	r.Handle("/", http.FileServer(http.FS(view))) // 處理靜態資源
	r.HandleFunc("/", serveView).Methods(http.MethodGet) // 顯示前端頁面
	r.HandleFunc("/ws", serveWS).Methods(http.MethodGet) // 處理 WebSocket 請求
	srv := http.Server{
		Handler:      r,
		Addr:         ":80",
		WriteTimeout: 15 * time.Second,
		ReadTimeout:  15 * time.Second,
	}
	// 實現 graceful shutdown 機制 
	go func() {
		signals := make(chan os.Signal, 1)
		signal.Notify(signals, os.Interrupt)
		<-signals
		log.Println("Shutting down server...")
		if err := srv.Shutdown(context.Background()); err != nil {
			log.Println(err)
		}
		close(quit)
	}()
	if err := srv.ListenAndServe(); err != http.ErrServerClosed {
		log.Fatal(err)
	}
	<-quit
}

func serveView(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, "index.html")
}

func serveWS(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Fatal(err)
	}
	client := &Client{
		conn: conn,
	}
	go client.writePump()
	go client.readPump()
}

新增 rdb.go 檔，建立 Redis 連線。

package main

import (
	"context"
	"fmt"
	"github.com/go-redis/redis/v8"
	"log"
	"os"
)

var (
	RDB *redis.Client
)

func init() {
	RDB = redis.NewClient(&redis.Options{
		Addr:     fmt.Sprintf("%s:%s", os.Getenv("REDIS_HOST"), os.Getenv("REDIS_PORT")),
		Password: os.Getenv("REDIS_PASSWORD"),
	})
	if _, err := RDB.Ping(context.Background()).Result(); err != nil {
		log.Fatal(err)
	}
}

新增 client.go 檔，處理訊息交換。

package main

import (
	"context"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

const (
	writeWait      = 10 * time.Second
	pongWait       = 60 * time.Second
	pingPeriod     = pongWait * 9 / 10
	maxMessageSize = 512
)

const (
	channelGeneral = "general"
)

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

type Client struct {
	conn *websocket.Conn
}

// 建立一個讀取幫浦
func (c *Client) readPump() {
	defer func() {
		_ = c.conn.Close()
	}()
	c.conn.SetReadLimit(maxMessageSize)
	// 處理 pong 訊息
	c.conn.SetPongHandler(func(string) error { _ = c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	_ = c.conn.SetReadDeadline(time.Now().Add(pongWait))
	// 不斷讀取 WebSocket 訊息
	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Println(err)
			}
			break
		}
		// 向 Redis 頻道發布訊息
		if err := RDB.Publish(context.Background(), channelGeneral, message).Err(); err != nil {
			log.Println(err)
		}
	}
}

// 建立一個寫入幫浦
func (c *Client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	// 向 Redis 頻道訂閱訊息
	sub := RDB.Subscribe(context.Background(), channelGeneral)
	defer func() {
		ticker.Stop()
		_ = c.conn.Close()
		_ = sub.Close()
	}()
	for {
		select {
		// 不斷讀取 Redis 頻道訊息
		case msg := <-sub.Channel():
			_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			// 向 WebSocket 發送訊息
			if _, err := w.Write([]byte(msg.Payload)); err != nil {
				log.Println(err)
			}
			if err := w.Close(); err != nil {
				return
			}
		// 處理 ping 訊息
		case <-ticker.C:
			_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

參考資料

• gorilla/websocket