diff --git a/internal/executor/sandbox/executor.go b/internal/executor/sandbox/executor.go index 65c3369..987382f 100644 --- a/internal/executor/sandbox/executor.go +++ b/internal/executor/sandbox/executor.go @@ -1,27 +1,16 @@ package sandbox import ( - "archive/tar" - "bytes" "context" "fmt" - "io" "log/slog" "maps" - "os" - "path/filepath" - "strings" "github.com/criyle/go-judge/pb" "github.com/joint-online-judge/JOJ3/internal/stage" "google.golang.org/protobuf/proto" ) -const ( - tarSizeThreshold = 128 * 1024 // 128 KB - tarCountThreshold = 100 // 100 files -) - func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { var err error if e.execClient == nil { @@ -31,51 +20,18 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { return nil, err } } + // cannot use range loop since we need to change the value for i := 0; i < len(cmds); i += 1 { - if cmd := &cmds[i]; cmd.CopyIn == nil { + cmd := &cmds[i] + if cmd.CopyIn == nil { cmd.CopyIn = make(map[string]stage.CmdFile) } - for k, v := range cmds[i].CopyInCached { + for k, v := range cmd.CopyInCached { if fileID, ok := e.cachedMap[v]; ok { - cmds[i].CopyIn[k] = stage.CmdFile{FileID: &fileID} + cmd.CopyIn[k] = stage.CmdFile{FileID: &fileID} } } } - if needTar, tarData := prepareTar(cmds); needTar { - return e.runWithTar(cmds, tarData) - } - return e.runUnary(cmds) -} - -func prepareTar(cmds []stage.Cmd) (bool, []byte) { - if len(cmds) == 0 { - return false, nil - } - for i := range cmds { - if shouldTar(&cmds[i]) { - tarData, keysInTar := createCopyInTar(&cmds[i]) - if tarData == nil { - return false, nil - } - slog.Debug("prepareTar", "tarSize", len(tarData), "keysInTar", len(keysInTar)) - for j := range cmds { - cmds[j].CopyInDir = "" - for _, k := range keysInTar { - delete(cmds[j].CopyIn, k) - } - } - return true, tarData - } - } - return false, nil -} - -func shouldTar(cmd *stage.Cmd) bool { - size, count := estimateCopyIn(cmd) - return size >= tarSizeThreshold || count >= tarCountThreshold -} - -func (e *Sandbox) runUnary(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { pbCmds := convertPBCmd(cmds) for i, pbCmd := range pbCmds { slog.Debug("sandbox execute", "i", i, "pbCmd size", proto.Size(pbCmd)) @@ -97,213 +53,6 @@ func (e *Sandbox) runUnary(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { return results, nil } -func (e *Sandbox) runWithTar(cmds []stage.Cmd, tarData []byte) ([]stage.ExecutorResult, error) { - fc := &pb.FileContent{} - fc.SetContent(tarData) - fileIDResp, err := e.execClient.FileAdd(context.TODO(), fc) - if err != nil { - return nil, fmt.Errorf("file add tar: %w", err) - } - fid := fileIDResp.GetFileID() - slog.Debug("sandbox tar uploaded", "fileID", fid, "tarSize", len(tarData)) - - defer func() { - deleteReq := &pb.FileID{} - deleteReq.SetFileID(fid) - if _, err := e.execClient.FileDelete(context.TODO(), deleteReq); err != nil { - slog.Warn("sandbox tar file delete", "fileID", fid, "error", err) - } - }() - - tarFileName := "/w/__joj3_copyin.tar" - script := fmt.Sprintf( - "/bin/tar xf %s -C / --no-same-owner && rm -f %s && exec \"$@\"", - tarFileName, tarFileName, - ) - - for i := range cmds { - if cmds[i].CopyIn == nil { - cmds[i].CopyIn = make(map[string]stage.CmdFile) - } - cmds[i].CopyIn[tarFileName] = stage.CmdFile{FileID: &fid} - cmds[i].Args = append([]string{ - "/bin/sh", "-c", script, "_", - }, cmds[i].Args...) - } - - slog.Debug("sandbox tar exec", "cmd", cmds[0].Args[:3]) - - return e.runUnary(cmds) -} - -func estimateCopyIn(cmd *stage.Cmd) (int, int) { - totalSize := 0 - totalCount := 0 - if cmd.CopyInDir != "" { - _ = filepath.Walk(cmd.CopyInDir, - func(path string, info os.FileInfo, err error) error { - if err != nil || info.IsDir() { - return nil - } - relPath, err := filepath.Rel(cmd.CopyInDir, path) - if err != nil { - return nil - } - if _, exists := cmd.CopyIn[relPath]; !exists { - totalSize += int(info.Size()) - totalCount++ - } - return nil - }) - } - for _, f := range cmd.CopyIn { - if f.Symlink != nil { - continue - } - if f.Src != nil { - if fi, err := os.Stat(*f.Src); err == nil { - totalSize += int(fi.Size()) - totalCount++ - } - } else if f.Content != nil { - totalSize += len(*f.Content) - totalCount++ - } - } - return totalSize, totalCount -} - -func formatTarPath(p string) string { - if filepath.IsAbs(p) { - return strings.TrimPrefix(p, "/") - } - return "w/" + p -} - -func createCopyInTar(cmd *stage.Cmd) ([]byte, []string) { - var buf bytes.Buffer - tw := tar.NewWriter(&buf) - tarKeys := make([]string, 0, len(cmd.CopyIn)) - - if cmd.CopyInDir != "" { - err := filepath.Walk(cmd.CopyInDir, - func(path string, info os.FileInfo, err error) error { - if err != nil { - return err - } - relPath, err := filepath.Rel(cmd.CopyInDir, path) - if err != nil { - return err - } - if relPath == "." { - return nil - } - if _, exists := cmd.CopyIn[relPath]; exists { - if info.IsDir() { - return filepath.SkipDir - } - return nil - } - if info.Mode()&os.ModeSymlink != 0 { - link, err := os.Readlink(path) - if err != nil { - return err - } - hdr, err := tar.FileInfoHeader(info, link) - if err != nil { - return err - } - hdr.Name = formatTarPath(relPath) - return tw.WriteHeader(hdr) - } - hdr, err := tar.FileInfoHeader(info, "") - if err != nil { - return err - } - hdr.Name = formatTarPath(relPath) - if info.IsDir() { - hdr.Name += "/" - } - if err := tw.WriteHeader(hdr); err != nil { - return err - } - if info.IsDir() { - return nil - } - f, err := os.Open(path) - if err != nil { - return err - } - _, err = io.Copy(tw, f) - f.Close() - return err - }) - if err != nil { - slog.Error("create copyIn tar walk", "error", err) - return nil, nil - } - } - - for k, f := range cmd.CopyIn { - if f.FileID != nil || f.Symlink != nil || f.StreamIn || f.StreamOut || f.Pipe { - continue - } - if f.Content != nil { - hdr := &tar.Header{ - Name: formatTarPath(k), - Mode: 0o644, - Size: int64(len(*f.Content)), - } - if err := tw.WriteHeader(hdr); err != nil { - slog.Error("create copyIn tar write header", "key", k, "error", err) - continue - } - if _, err := tw.Write([]byte(*f.Content)); err != nil { - slog.Error("create copyIn tar write content", "key", k, "error", err) - continue - } - tarKeys = append(tarKeys, k) - } else if f.Src != nil { - fi, err := os.Stat(*f.Src) - if err != nil { - slog.Error("create copyIn tar stat", "key", k, "src", *f.Src, "error", err) - continue - } - if fi.IsDir() { - continue - } - hdr, err := tar.FileInfoHeader(fi, "") - if err != nil { - slog.Error("create copyIn tar file info header", "key", k, "error", err) - continue - } - hdr.Name = formatTarPath(k) - if err := tw.WriteHeader(hdr); err != nil { - slog.Error("create copyIn tar write header", "key", k, "error", err) - continue - } - srcFile, err := os.Open(*f.Src) - if err != nil { - slog.Error("create copyIn tar open src", "key", k, "src", *f.Src, "error", err) - continue - } - _, err = io.Copy(tw, srcFile) - srcFile.Close() - if err != nil { - slog.Error("create copyIn tar copy", "key", k, "error", err) - continue - } - tarKeys = append(tarKeys, k) - } - } - - if err := tw.Close(); err != nil { - slog.Error("create copyIn tar close", "error", err) - return nil, nil - } - return buf.Bytes(), tarKeys -} - func (e *Sandbox) Cleanup() error { for k, fileID := range e.cachedMap { req := &pb.FileID{}