package main import ( "context" "fmt" "path" "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<- result } 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) { var rt result if len(req.Cmd) == 1 { rt = workDoSingle(req.Cmd[0]) } else { rt = workDoGroup(req.Cmd, req.PipeMapping) } rt.RequestID = req.RequestID req.resultCh <- rt } func workDoSingle(rc cmd) (rt result) { c, copyOutSet, err := prepareCmd(rc) if err != nil { rt.Error = err return } s := &envexec.Single{ EnvironmentPool: envPool, Cmd: c, } result, err := s.Run() if err != nil { rt.Error = err return } res := convertResult(result, copyOutSet) rt.Response = []response{res} return } func workDoGroup(rc []cmd, pm []pipeMap) (rt result) { var rts []response p := preparePipeMapping(pm) cs := make([]*envexec.Cmd, 0, len(rc)) copyOutSets := make([]map[string]bool, 0, len(rc)) for _, cc := range rc { c, os, err := prepareCmd(cc) if err != nil { rt.Error = err return } cs = append(cs, c) copyOutSets = append(copyOutSets, os) } g := envexec.Group{ EnvironmentPool: envPool, Cmd: cs, Pipes: p, } results, err := g.Run() if err != nil { rt.Error = err return } rts = make([]response, 0, len(results)) for i, result := range results { res := convertResult(result, copyOutSets[i]) rts = append(rts, res) } rt.Response = rts return } func convertResult(result envexec.Result, copyOutSet map[string]bool) (res response) { res.Status = status(result.Status) res.ExitStatus = result.ExitStatus 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 := fs.Add(name, b) if err != nil { res.Status = status(envexec.StatusFileError) res.Error = err.Error() return } res.FileIDs[name] = id } } return res } func prepareCmd(rc cmd) (*envexec.Cmd, map[string]bool, error) { files, pipeFileName, err := prepareCmdFiles(rc.Files) if err != nil { return nil, nil, err } copyIn, err := prepareCopyIn(rc.CopyIn) if err != nil { return nil, nil, err } copyOutSet := make(map[string]bool) // pipe default copyout for k := range pipeFileName { copyOutSet[k] = true } 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) } else { delete(copyOutSet, fn) } } w := &waiter{ timeLimit: time.Duration(rc.CPULimit), realTimeLimit: time.Duration(rc.RealCPULimit), } var copyOutDir string if rc.CopyOutDir != "" { if path.IsAbs(rc.CopyOutDir) { copyOutDir = rc.CopyOutDir } else { copyOutDir = path.Join(*dir, rc.CopyOutDir) } } timeLimit := time.Duration(rc.CPULimit) if rc.RealCPULimit > rc.CPULimit { timeLimit = time.Duration(rc.RealCPULimit) } return &envexec.Cmd{ Args: rc.Args, Env: rc.Env, Files: files, TimeLimit: timeLimit, MemoryLimit: runner.Size(rc.MemoryLimit), ProcLimit: rc.ProcLimit, CopyIn: copyIn, CopyOut: copyOut, CopyOutDir: copyOutDir, Waiter: w.Wait, }, copyOutSet, nil } func preparePipeMapping(pm []pipeMap) []*envexec.Pipe { rt := make([]*envexec.Pipe, 0, len(pm)) for _, p := range pm { rt = append(rt, &envexec.Pipe{ In: envexec.PipeIndex{Index: p.In.Index, Fd: p.In.Fd}, Out: envexec.PipeIndex{Index: p.Out.Index, Fd: p.Out.Fd}, }) } return rt } 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 := fs.Get(*f.FileID) if fd == nil { return nil, fmt.Errorf("file not exists for %v", *f.FileID) } return fd, 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 result { ch := make(chan result, 1) workCh <- workRequest{ request: req, resultCh: ch, } return ch } func workerShutdown() { workStopOnce.Do(func() { workCtxCancel() workWg.Wait() }) }