connect/websocket.go

package main

import (
	"flag"
	"fmt"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"sikey.com/websocket/config"
	"sikey.com/websocket/repositories"
	"sikey.com/websocket/server"
	"sikey.com/websocket/stackexchange"
	"sikey.com/websocket/utils/mysqlx"
	"sikey.com/websocket/utils/zlog"
)

var configFile = flag.String("f", "./etc/websocket.toml", "the config file")

func main() {
	flag.Parse()
	config.MustLoadConfig(*configFile)

	// Zaplog init
	zlog.WithZapLogger(zlog.NewLogger(config.MustLoadLogger()))

	app := newApp()
	app.Run(fmt.Sprintf(":%d", config.Config.Port))
}

func newApp() *gin.Engine {
	app := gin.Default()

	// 创建 Websocket 横向拓展应用
	exc := stackexchange.NewStackExchange(
		config.Kafka.Brokers,
		config.Kafka.Topic,
		config.Kafka.Partition,
		config.Kafka.MaxBytes,
	)

	srv := &server.Server{
		Upgrader: websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
		WriteWait: 10 * time.Second,
		ReadWait:  10 * time.Second,
		PingWait:  120 * time.Second,
		Hub: server.NewHub(server.HubConfig{
			ConnectSize:    1024,
			DisconnectSize: 1024,
			MessageSize:    125,
			StackExchange:  exc,
		}),
		Repositories: repositories.NewRepositories(mysqlx.ConnectMysql()),
	}
	app.GET("/websocket/endpoint", func(ctx *gin.Context) { server.WebsocketHandler(ctx, srv) })
	return app
}