go-judge/cmd/executorserver/worker.go
criyle 7951f37bd4 Implements the draft executor server
- Add new command executorserver as the draft executor server
- Move shared pool logic into /pkg/pool
2020-03-04 02:12:26 -05:00

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()
})
}