ws: add ability to cancel running task

- client is able to cancel task via close connection and cancel request

close #24
This commit is contained in:
criyle 2021-11-14 19:17:07 -08:00
parent 049d26a179
commit 1c3681ba39
2 changed files with 93 additions and 11 deletions

View File

@ -2,7 +2,9 @@ package wsexecutor
import (
"context"
"fmt"
"net/http"
"sync"
"time"
"github.com/criyle/go-judge/cmd/executorserver/model"
@ -44,6 +46,11 @@ type wsHandle struct {
logger *zap.Logger
}
type wsRequest struct {
model.Request
CancelRequestId string `json:"cancelRequestId"`
}
func (h *wsHandle) Register(r *gin.Engine) {
r.GET("/ws", h.handleWS)
}
@ -56,6 +63,48 @@ func (h *wsHandle) handleWS(c *gin.Context) {
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 {
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)
ret := <-h.worker.Submit(ctx, r)
h.logger.Sugar().Debugf("ws response: %+v", ret)
resp, err := model.ConvertResponse(ret, false)
if err != nil {
resultCh <- model.Response{
RequestID: r.RequestID,
ErrorMsg: resp.ErrorMsg,
}
return
}
resultCh <- resp
}()
return nil
}
// read request
go func() {
defer conn.Close()
@ -65,25 +114,19 @@ func (h *wsHandle) handleWS(c *gin.Context) {
return nil
})
ctx, cancel := context.WithCancel(context.TODO())
defer cancel()
baseCtx, baseCancel := context.WithCancel(context.TODO())
defer baseCancel()
for {
req := new(model.Request)
req := new(wsRequest)
if err := conn.ReadJSON(req); err != nil {
h.logger.Sugar().Warn("ws read error:", err)
return
}
r, err := model.ConvertRequest(req, h.srcPrefix)
if err != nil {
h.logger.Sugar().Warn("convert error: ", err)
if err := handleRequest(baseCtx, req); err != nil {
h.logger.Sugar().Warn(err.Error())
return
}
go func() {
ret := <-h.worker.Submit(ctx, r)
resp, _ := model.ConvertResponse(ret, false)
resultCh <- resp
}()
}
}()
@ -109,3 +152,36 @@ func (h *wsHandle) handleWS(c *gin.Context) {
}
}()
}
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()
}
}

View File

@ -249,6 +249,12 @@ func (w *worker) convertResult(result envexec.Result, cmd Cmd) (res Result) {
res.Files = make(map[string]*os.File)
res.FileIDs = make(map[string]string)
// Fix TLE due to context cancel
if res.Status == envexec.StatusTimeLimitExceeded && res.ExitStatus != 0 &&
res.Time < cmd.CPULimit && res.RunTime < cmd.ClockLimit {
res.Status = envexec.StatusSignalled
}
copyOutCachedSet := make(map[string]bool, len(cmd.CopyOutCached))
for _, f := range cmd.CopyOutCached {
copyOutCachedSet[f.Name] = true