package worker import ( "Wine-Server/utils" "Wine-Server/utils/tables" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" ) func workerConnected(wid string, conn *websocket.Conn) *tables.WorkerTable { table := &tables.WorkerTable{Id: wid} err := table.Get() if err != nil { utils.Logger.Println("worker query failed:", err) _ = conn.WriteJSON(utils.WsError("worker query failed")) return nil } table.Last = utils.TimeNow() err = table.Update(utils.JsonType{"last": table.Last}) if err != nil { utils.Logger.Println("update last active time failed:", err) _ = conn.WriteJSON(utils.WsError("Update last active time failed")) return nil } //keys := make([]string, 0, len(utils.WorkerWss)) //for k := range utils.WorkerWss { // keys = append(keys, k) //} //utils.Logger.Println(wid, keys) //if _, exist := utils.WorkerWss[wid]; exist { // _ = conn.WriteJSON(utils.WsEvent("onePageOnly", nil)) // return nil //} utils.WorkerLock.Lock() utils.WorkerWss[wid] = conn utils.WorkerLock.Unlock() utils.Logger.Printf("worker[%s] connected\n", wid) return table } func workerDisconnect(wid string) { utils.WorkerLock.Lock() delete(utils.WorkerWss, wid) utils.WorkerLock.Unlock() utils.Logger.Printf("worker[%s] disconnected\n", wid) } func loginSocket(conn *websocket.Conn) (string, error) { var wid string for { var msg utils.WsMsg err := conn.ReadJSON(&msg) if err != nil { _ = conn.Close() return "", err } switch msg.Event { case "pin": keepAlive(conn) break case "login": wid, err = login(conn, msg.Data) if err == nil { return wid, nil } break default: _ = conn.WriteJSON(utils.WsError("loginSocket event only")) break } } } func socketHandler(ctx *gin.Context) { conn, err := utils.UpgradeHttp2Ws.Upgrade(ctx.Writer, ctx.Request, nil) if err != nil { utils.Logger.Println("can't establish manager socket connect") ctx.JSON(utils.HttpError, utils.Fail("can't establish socket connect")) return } getSecrets(conn) var wid string token := ctx.Param("token") if token == "undefined" { // no token -> login wid, err = loginSocket(conn) } else { wid, err = utils.Redis.Get(utils.WxPayCli, token).Result() if err != nil { // expired -> login _ = conn.WriteJSON(utils.WsEvent("tokenExpired", nil)) wid, err = loginSocket(conn) } else { // update expire time utils.Redis.Expire(utils.WxPayCli, token, utils.Duration(7*24*60*60)) } } if err != nil { utils.Logger.Println(err) return } worker := workerConnected(wid, conn) if worker == nil { return } defer workerDisconnect(wid) info(conn, worker) workerConnected(wid, conn) defer workerDisconnect(wid) for { var msg utils.WsMsg err = conn.ReadJSON(&msg) if err != nil { return } switch msg.Event { case "pin": keepAlive(conn) break case "query": query(conn, wid, msg.Data) break case "detail": detail(conn, msg.Data) break case "ready": ready(conn, wid, msg.Data) break case "finish": finish(conn, worker, msg.Data) break default: _ = conn.WriteJSON(utils.WsError("unrecognized event")) break } } }