mirror of
https://github.com/criyle/go-judge.git
synced 2025-11-04 14:50:02 +08:00
Handle cancellation
This commit is contained in:
parent
891ed8165b
commit
f61a319cd6
@ -25,7 +25,7 @@ func handleRun(c *gin.Context) {
|
||||
c.AbortWithStatusJSON(http.StatusBadRequest, err.Error)
|
||||
return
|
||||
}
|
||||
rt := <-work.Submit(r)
|
||||
rt := <-work.Submit(c.Request.Context(), r)
|
||||
execObserve(rt)
|
||||
if rt.Error != nil {
|
||||
c.Error(rt.Error)
|
||||
|
||||
@ -34,7 +34,7 @@ func (e *execServer) Exec(ctx context.Context, req *pb.Request) (*pb.Response, e
|
||||
if len(si) > 0 || len(so) > 0 {
|
||||
return nil, fmt.Errorf("Stream in / out are not avaliable for exec request")
|
||||
}
|
||||
rt := <-work.Submit(r)
|
||||
rt := <-work.Submit(ctx, r)
|
||||
execObserve(rt)
|
||||
if rt.Error != nil {
|
||||
return nil, err
|
||||
@ -90,7 +90,7 @@ func (e *execServer) ExecStream(es pb.Executor_ExecStreamServer) error {
|
||||
}
|
||||
}
|
||||
|
||||
rtCh := work.Submit(rq)
|
||||
rtCh := work.Execute(es.Context(), rq)
|
||||
for {
|
||||
select {
|
||||
case err := <-errCh:
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
@ -37,6 +38,10 @@ func handleWS(c *gin.Context) {
|
||||
conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
return nil
|
||||
})
|
||||
|
||||
ctx, cancel := context.WithCancel(context.TODO())
|
||||
defer cancel()
|
||||
|
||||
for {
|
||||
req := new(model.Request)
|
||||
if err := conn.ReadJSON(req); err != nil {
|
||||
@ -49,7 +54,7 @@ func handleWS(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
ret := <-work.Submit(r)
|
||||
ret := <-work.Submit(ctx, r)
|
||||
execObserve(ret)
|
||||
resultCh <- model.ConvertResponse(ret)
|
||||
}()
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
package envexec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"golang.org/x/sync/errgroup"
|
||||
@ -32,7 +33,7 @@ type Pipe struct {
|
||||
}
|
||||
|
||||
// Run starts the cmd and returns exec results
|
||||
func (r *Group) Run() ([]Result, error) {
|
||||
func (r *Group) Run(ctx context.Context) ([]Result, error) {
|
||||
// prepare files
|
||||
fds, pipeToCollect, fileToClose, err := prepareFds(r)
|
||||
defer func() { closeFiles(fileToClose) }()
|
||||
@ -58,7 +59,7 @@ func (r *Group) Run() ([]Result, error) {
|
||||
for i, c := range r.Cmd {
|
||||
i, c := i, c
|
||||
g.Go(func() error {
|
||||
r, err := runSingle(ms[i], c, fds[i], pipeToCollect[i])
|
||||
r, err := runSingle(ctx, ms[i], c, fds[i], pipeToCollect[i])
|
||||
result[i] = r
|
||||
if err != nil {
|
||||
result[i].Status = StatusInternalError
|
||||
|
||||
@ -8,7 +8,7 @@ import (
|
||||
)
|
||||
|
||||
// runSingle runs Cmd inside the given environment and cgroup
|
||||
func runSingle(m Environment, c *Cmd, fds []*os.File, ptc []pipeCollector) (result Result, err error) {
|
||||
func runSingle(pc context.Context, m Environment, c *Cmd, fds []*os.File, ptc []pipeCollector) (result Result, err error) {
|
||||
fdToClose := fds
|
||||
defer func() { closeFiles(fdToClose) }()
|
||||
|
||||
@ -32,7 +32,7 @@ func runSingle(m Environment, c *Cmd, fds []*os.File, ptc []pipeCollector) (resu
|
||||
}
|
||||
|
||||
// start the cmd (they will be canceled in other goroutines)
|
||||
ctx, cancel := context.WithCancel(context.TODO())
|
||||
ctx, cancel := context.WithCancel(pc)
|
||||
waiterCtx, waiterCancel := context.WithCancel(ctx)
|
||||
|
||||
process, err := m.Execve(ctx, execParam)
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
package envexec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
@ -15,7 +16,7 @@ type Single struct {
|
||||
}
|
||||
|
||||
// Run starts the cmd and returns exec results
|
||||
func (s *Single) Run() (result Result, err error) {
|
||||
func (s *Single) Run(ctx context.Context) (result Result, err error) {
|
||||
// prepare files
|
||||
fd, fileToClose, pipeToCollect, err := prepareCmdFd(s.Cmd, len(s.Cmd.Files))
|
||||
defer func() { closeFiles(fileToClose) }()
|
||||
@ -31,7 +32,7 @@ func (s *Single) Run() (result Result, err error) {
|
||||
}
|
||||
defer s.EnvironmentPool.Put(m)
|
||||
|
||||
result, err = runSingle(m, s.Cmd, fd, pipeToCollect)
|
||||
result, err = runSingle(ctx, m, s.Cmd, fd, pipeToCollect)
|
||||
fileToClose = nil // already closed by runOne
|
||||
if err != nil {
|
||||
result.Status = StatusInternalError
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
package worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"path"
|
||||
"sync"
|
||||
@ -30,6 +31,7 @@ type Worker struct {
|
||||
|
||||
type workRequest struct {
|
||||
*Request
|
||||
context.Context
|
||||
resultCh chan<- Response
|
||||
}
|
||||
|
||||
@ -56,15 +58,32 @@ func (w *Worker) Start() {
|
||||
}
|
||||
|
||||
// Submit submits a single request
|
||||
func (w *Worker) Submit(req *Request) <-chan Response {
|
||||
func (w *Worker) Submit(ctx context.Context, req *Request) <-chan Response {
|
||||
ch := make(chan Response, 1)
|
||||
w.workCh <- workRequest{
|
||||
Request: req,
|
||||
Context: ctx,
|
||||
resultCh: ch,
|
||||
}
|
||||
return ch
|
||||
}
|
||||
|
||||
// Execute will execute the request in new goroutine (bypass the parallism limit)
|
||||
func (w *Worker) Execute(ctx context.Context, req *Request) <-chan Response {
|
||||
ch := make(chan Response, 1)
|
||||
w.wg.Add(1)
|
||||
go func() {
|
||||
defer w.wg.Done()
|
||||
wq := workRequest{
|
||||
Request: req,
|
||||
Context: ctx,
|
||||
resultCh: ch,
|
||||
}
|
||||
w.workDoCmd(wq)
|
||||
}()
|
||||
return ch
|
||||
}
|
||||
|
||||
// Shutdown waits all worker to finish
|
||||
func (w *Worker) Shutdown() {
|
||||
w.stopOnce.Do(func() {
|
||||
@ -91,15 +110,15 @@ func (w *Worker) loop() {
|
||||
func (w *Worker) workDoCmd(req workRequest) {
|
||||
var rt Response
|
||||
if len(req.Cmd) == 1 {
|
||||
rt = w.workDoSingle(req.Cmd[0])
|
||||
rt = w.workDoSingle(req.Context, req.Cmd[0])
|
||||
} else {
|
||||
rt = w.workDoGroup(req.Cmd, req.PipeMapping)
|
||||
rt = w.workDoGroup(req.Context, req.Cmd, req.PipeMapping)
|
||||
}
|
||||
rt.RequestID = req.RequestID
|
||||
req.resultCh <- rt
|
||||
}
|
||||
|
||||
func (w *Worker) workDoSingle(rc Cmd) (rt Response) {
|
||||
func (w *Worker) workDoSingle(ctx context.Context, rc Cmd) (rt Response) {
|
||||
c, copyOutSet, err := w.prepareCmd(rc)
|
||||
if err != nil {
|
||||
rt.Error = err
|
||||
@ -109,7 +128,7 @@ func (w *Worker) workDoSingle(rc Cmd) (rt Response) {
|
||||
EnvironmentPool: w.envPool,
|
||||
Cmd: c,
|
||||
}
|
||||
result, err := s.Run()
|
||||
result, err := s.Run(ctx)
|
||||
if err != nil {
|
||||
rt.Error = err
|
||||
return
|
||||
@ -119,7 +138,7 @@ func (w *Worker) workDoSingle(rc Cmd) (rt Response) {
|
||||
return
|
||||
}
|
||||
|
||||
func (w *Worker) workDoGroup(rc []Cmd, pm []PipeMap) (rt Response) {
|
||||
func (w *Worker) workDoGroup(ctx context.Context, rc []Cmd, pm []PipeMap) (rt Response) {
|
||||
var rts []Result
|
||||
p := preparePipeMapping(pm)
|
||||
cs := make([]*envexec.Cmd, 0, len(rc))
|
||||
@ -139,7 +158,7 @@ func (w *Worker) workDoGroup(rc []Cmd, pm []PipeMap) (rt Response) {
|
||||
Cmd: cs,
|
||||
Pipes: p,
|
||||
}
|
||||
results, err := g.Run()
|
||||
results, err := g.Run(ctx)
|
||||
if err != nil {
|
||||
rt.Error = err
|
||||
return
|
||||
|
||||
Loading…
Reference in New Issue
Block a user