go-judge/cmd/executorserver/ws_executor/websocket.go

211 lines
4.5 KiB
Go

package wsexecutor
import (
"context"
"fmt"
"net/http"
"sync"
"time"
"github.com/criyle/go-judge/cmd/executorserver/model"
"github.com/criyle/go-judge/worker"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"go.uber.org/zap"
)
// Register registers web socket handle /ws
type Register interface {
Register(*gin.Engine)
}
// New creates new websocket handle
func New(worker worker.Worker, srcPrefix string, logger *zap.Logger) Register {
return &wsHandle{
worker: worker,
srcPrefix: srcPrefix,
logger: logger,
}
}
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
)
type wsHandle struct {
worker worker.Worker
srcPrefix string
logger *zap.Logger
}
type wsRequest struct {
model.Request
CancelRequestId string `json:"cancelRequestId"`
}
func (h *wsHandle) Register(r *gin.Engine) {
r.GET("/ws", h.handleWS)
}
func (h *wsHandle) 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 model.Response, 128)
cm := newContextMap()
handleRequest := func(baseCtx context.Context, req *wsRequest) error {
if req.CancelRequestId != "" {
h.logger.Sugar().Debugf("ws cancel: %s", req.CancelRequestId)
cm.Remove(req.CancelRequestId)
return nil
}
r, err := model.ConvertRequest(&req.Request, h.srcPrefix)
if err != nil {
return fmt.Errorf("ws convert error: %v", err)
}
ctx, cancel := context.WithCancel(baseCtx)
if err := cm.Add(r.RequestID, cancel); err != nil {
select {
case <-baseCtx.Done():
case resultCh <- model.Response{
RequestID: req.RequestID,
ErrorMsg: err.Error(),
}:
}
cancel()
h.logger.Sugar().Debugf("ws request error: %v", err)
return nil
}
go func() {
defer cm.Remove(r.RequestID)
h.logger.Sugar().Debugf("ws request: %+v", r)
retCh, started := h.worker.Submit(ctx, r)
var ret worker.Response
select {
case <-baseCtx.Done(): // if connection lost
return
case <-ctx.Done(): // if context cancelled by cancelling request
select {
case <-started: // if started, wait for result
ret = <-retCh
default: // not started
ret = worker.Response{
RequestID: r.RequestID,
Error: fmt.Errorf("request cancelled before execute"),
}
}
case ret = <-retCh:
}
h.logger.Sugar().Debugf("ws response: %+v", ret)
resp, err := model.ConvertResponse(ret, false)
if err != nil {
resp = model.Response{
RequestID: r.RequestID,
ErrorMsg: resp.ErrorMsg,
}
}
select {
case <-baseCtx.Done():
case resultCh <- resp:
}
}()
return nil
}
// 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
})
baseCtx, baseCancel := context.WithCancel(context.TODO())
defer baseCancel()
for {
req := new(wsRequest)
if err := conn.ReadJSON(req); err != nil {
h.logger.Sugar().Warn("ws read error:", err)
return
}
if err := handleRequest(baseCtx, req); err != nil {
h.logger.Sugar().Warn(err.Error())
return
}
}
}()
// write result
go func() {
defer conn.Close()
ticker := time.NewTicker(pingPeriod)
defer ticker.Stop()
for {
select {
case r := <-resultCh:
conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := conn.WriteJSON(r); err != nil {
h.logger.Sugar().Warn("ws write error:", err)
return
}
case <-ticker.C:
conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}()
}
type contextMap struct {
m map[string]context.CancelFunc
mu sync.Mutex
}
func newContextMap() *contextMap {
return &contextMap{m: make(map[string]context.CancelFunc)}
}
func (c *contextMap) Add(reqId string, cancel context.CancelFunc) error {
if reqId == "" {
return fmt.Errorf("empty request id")
}
c.mu.Lock()
defer c.mu.Unlock()
if _, exist := c.m[reqId]; exist {
return fmt.Errorf("duplicated request id: %v", reqId)
}
c.m[reqId] = cancel
return nil
}
func (c *contextMap) Remove(reqId string) {
c.mu.Lock()
defer c.mu.Unlock()
if cancel, exist := c.m[reqId]; exist {
delete(c.m, reqId)
cancel()
}
}