connect/websocket.go

package main

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

	"github.com/DeanThompson/ginpprof"
	"github.com/denisbrodbeck/machineid"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"sikey.com/websocket/config"
	"sikey.com/websocket/repositories"
	"sikey.com/websocket/server"
	"sikey.com/websocket/utils/gid"
	"sikey.com/websocket/utils/mysqlx"
	"sikey.com/websocket/utils/redisx"
	"sikey.com/websocket/utils/zlog"
)

var configFile = flag.String("f", "./etc/websocket.toml", "the config file")
var nodeId = flag.Int64("n", 1, "the node id")

func main() {
	time.Local = time.UTC

	flag.Parse()
	config.MustLoadConfig(*configFile)
	gid.SetNodeId(*nodeId)

	// Zaplog init
	if config.Config.Environment == "release" {
		gin.SetMode(gin.ReleaseMode)
		zlog.WithZapLogger(zlog.NewSLSLogger(zlog.SLSWriterConfig{
			Endpoint:        config.SLS.Endpoint,
			AccessKeyId:     config.SLS.AccessKeyId,
			AccessKeySecret: config.SLS.AccessKeySecret,
			ProjectName:     config.SLS.ProjectName,
			LogStoreName:    config.SLS.LogStoreName,
			Topic:           config.SLS.Topic,
			SourceIP:        config.SLS.SourceIP,
			Environment:     config.Config.Environment,
		}))
	} else {
		zlog.WithZapLogger(zlog.NewLogger(config.MustLoadLogger()))
	}

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

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

	id := serverId()
	srv := &server.Server{
		ID: id,
		Upgrader: websocket.Upgrader{
			ReadBufferSize:  config.Websocket.ReadBufferSize,
			WriteBufferSize: config.Websocket.WriteBufferSize,
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
		Hub:          server.NewHub(id),
		Repositories: repositories.NewRepositories(mysqlx.Connect(), redisx.RedisConnect()),
	}

	app.GET("/websocket/endpoint", func(ctx *gin.Context) { srv.WebsocketHandler(ctx) })
	app.GET("/websocket/clients", func(ctx *gin.Context) {
		clients := srv.Hub.GetClients()
		ctx.JSON(http.StatusOK, gin.H{"clients": clients})
	})
	return app
}

func serverId() string {
	var id string
	id, err := machineid.ID()
	if err != nil {
		id = uuid.NewString()
	} else {
		id = strings.ToLower(id)
	}
	return id
}