mirror of
https://github.com/criyle/go-judge.git
synced 2025-11-04 14:50:02 +08:00
- Add new command executorserver as the draft executor server - Move shared pool logic into /pkg/pool
212 lines
4.4 KiB
Go
212 lines
4.4 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/criyle/go-judge/file"
|
|
"github.com/criyle/go-judge/pkg/envexec"
|
|
"github.com/criyle/go-sandbox/runner"
|
|
)
|
|
|
|
const maxWaiting = 512
|
|
|
|
type workRequest struct {
|
|
*request
|
|
resultCh chan<- response
|
|
}
|
|
|
|
var (
|
|
workStartOnce sync.Once
|
|
workStopOnce sync.Once
|
|
|
|
workCh chan workRequest
|
|
workWg sync.WaitGroup
|
|
|
|
workCtx context.Context
|
|
workCtxCancel func()
|
|
)
|
|
|
|
func startWorkers() {
|
|
workStartOnce.Do(func() {
|
|
workCtx, workCtxCancel = context.WithCancel(context.Background())
|
|
workCh = make(chan workRequest, maxWaiting)
|
|
workWg.Add(*parallism)
|
|
for i := 0; i < *parallism; i++ {
|
|
go workerLoop()
|
|
}
|
|
})
|
|
}
|
|
|
|
func workerLoop() {
|
|
defer workWg.Done()
|
|
for {
|
|
var req workRequest
|
|
select {
|
|
case <-workCtx.Done():
|
|
return
|
|
case req = <-workCh:
|
|
}
|
|
workDoCmd(req)
|
|
}
|
|
}
|
|
|
|
func workDoCmd(req workRequest) {
|
|
if len(req.Cmd) == 1 {
|
|
req.resultCh <- workDoSingle(req.Cmd[0])
|
|
} else {
|
|
req.resultCh <- response{
|
|
Error: "not implemented yet TAT",
|
|
}
|
|
}
|
|
}
|
|
|
|
func workDoSingle(rc cmd) (res response) {
|
|
files, pipeFileName, err := prepareCmdFiles(rc.Files)
|
|
if err != nil {
|
|
res.Status = status(envexec.StatusInternalError)
|
|
res.Error = err.Error()
|
|
return
|
|
}
|
|
copyIn, err := prepareCopyIn(rc.CopyIn)
|
|
|
|
copyOutSet := make(map[string]bool)
|
|
copyOut := make([]string, 0, len(rc.CopyOut)+len(rc.CopyOutCached))
|
|
for _, fn := range rc.CopyOut {
|
|
if !pipeFileName[fn] {
|
|
copyOut = append(copyOut, fn)
|
|
}
|
|
copyOutSet[fn] = true
|
|
}
|
|
for _, fn := range rc.CopyOutCached {
|
|
if !pipeFileName[fn] {
|
|
copyOut = append(copyOut, fn)
|
|
}
|
|
}
|
|
|
|
w := &waiter{
|
|
timeLimit: time.Duration(rc.CPULimit * float64(time.Second)),
|
|
realTimeLimit: time.Duration(rc.RealCPULimit * float64(time.Second)),
|
|
}
|
|
|
|
c := &envexec.Cmd{
|
|
Args: rc.Args,
|
|
Env: rc.Env,
|
|
Files: files,
|
|
MemoryLimit: runner.Size(rc.MemoryLimit),
|
|
ProcLimit: rc.ProcLimit,
|
|
CopyIn: copyIn,
|
|
CopyOut: copyOut,
|
|
Waiter: w.Wait,
|
|
}
|
|
s := &envexec.Single{
|
|
CgroupPool: cgroupPool,
|
|
EnvironmentPool: envPool,
|
|
Cmd: c,
|
|
}
|
|
result, err := s.Run()
|
|
if err != nil {
|
|
res.Status = status(envexec.StatusInternalError)
|
|
res.Error = err.Error()
|
|
return
|
|
}
|
|
res.Status = status(result.Status)
|
|
res.Error = result.Error
|
|
res.Time = uint64(result.Time)
|
|
res.Memory = uint64(result.Memory)
|
|
res.Files = make(map[string]string)
|
|
res.FileIDs = make(map[string]string)
|
|
|
|
for name, fi := range result.Files {
|
|
b, err := fi.Content()
|
|
if err != nil {
|
|
res.Status = status(envexec.StatusFileError)
|
|
res.Error = err.Error()
|
|
return
|
|
}
|
|
if copyOutSet[name] {
|
|
res.Files[name] = string(b)
|
|
} else {
|
|
id, err := addFile(b, name)
|
|
if err != nil {
|
|
res.Status = status(envexec.StatusFileError)
|
|
res.Error = err.Error()
|
|
return
|
|
}
|
|
res.FileIDs[name] = id
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
func prepareCopyIn(cf map[string]cmdFile) (map[string]file.File, error) {
|
|
rt := make(map[string]file.File)
|
|
for name, f := range cf {
|
|
pcf, err := prepareCmdFile(&f)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
fi, ok := pcf.(file.File)
|
|
if !ok {
|
|
return nil, fmt.Errorf("pipe type cannot be used for copyIn %s", name)
|
|
}
|
|
rt[name] = fi
|
|
}
|
|
return rt, nil
|
|
}
|
|
|
|
func prepareCmdFiles(files []*cmdFile) ([]interface{}, map[string]bool, error) {
|
|
rt := make([]interface{}, 0, len(files))
|
|
pipeFileName := make(map[string]bool)
|
|
for _, f := range files {
|
|
cf, err := prepareCmdFile(f)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
rt = append(rt, cf)
|
|
if t, ok := cf.(envexec.PipeCollector); ok {
|
|
pipeFileName[t.Name] = true
|
|
}
|
|
}
|
|
return rt, pipeFileName, nil
|
|
}
|
|
|
|
func prepareCmdFile(f *cmdFile) (interface{}, error) {
|
|
switch {
|
|
case f == nil:
|
|
return nil, nil
|
|
case f.Src != nil:
|
|
return file.NewLocalFile(*f.Src, *f.Src), nil
|
|
case f.Content != nil:
|
|
return file.NewMemFile("file", []byte(*f.Content)), nil
|
|
case f.FileID != nil:
|
|
fd, ok := getFile(*f.FileID)
|
|
if !ok {
|
|
return nil, fmt.Errorf("file not exists for %v", *f.FileID)
|
|
}
|
|
return file.NewMemFile(fd.FileName, fd.Content), nil
|
|
case f.Max != nil && f.Name != nil:
|
|
return envexec.PipeCollector{Name: *f.Name, SizeLimit: *f.Max}, nil
|
|
default:
|
|
return nil, fmt.Errorf("file is not valid for cmd")
|
|
}
|
|
}
|
|
|
|
func submitRequest(req *request) <-chan response {
|
|
ch := make(chan response, 1)
|
|
workCh <- workRequest{
|
|
request: req,
|
|
resultCh: ch,
|
|
}
|
|
return ch
|
|
}
|
|
|
|
func workerShutdown() {
|
|
workStopOnce.Do(func() {
|
|
workCtxCancel()
|
|
workWg.Wait()
|
|
})
|
|
}
|