go-judge/cmd/executorserver/websocket.go
2020-03-27 21:29:37 -04:00

77 lines
1.5 KiB
Go

package main
import (
"net/http"
"time"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
const (
writeWait = 10 * time.Second
pongWait = 60 * time.Second
pingPeriod = 50 * time.Second
)
func handleWS(c *gin.Context) {
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
c.Error(err)
c.AbortWithStatusJSON(http.StatusBadRequest, err.Error())
return
}
resultCh := make(chan result, maxWaiting)
// read request
go func() {
defer conn.Close()
conn.SetReadDeadline(time.Now().Add(pongWait))
conn.SetPongHandler(func(string) error {
conn.SetReadDeadline(time.Now().Add(pongWait))
return nil
})
for {
req := new(request)
if err := conn.ReadJSON(req); err != nil {
printLog("ws read:", err)
return
}
workCh <- workRequest{
request: req,
resultCh: resultCh,
}
}
}()
// write result
go func() {
defer conn.Close()
ticker := time.NewTicker(pingPeriod)
defer ticker.Stop()
for {
select {
case r := <-resultCh:
if r.Error != nil {
r.ErrorMsg = r.Error.Error()
}
conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := conn.WriteJSON(r); err != nil {
printLog("ws write:", err)
return
}
case <-ticker.C:
conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}()
}