From e751c9a237a10a82f9804f16c5aec395b5d35d22 Mon Sep 17 00:00:00 2001 From: Boming Zhang Date: Thu, 23 Jul 2026 05:14:20 -0700 Subject: [PATCH] fix(sandbox): optimize tar preparation, file limit thresholds, and multi-cmd support --- internal/executor/sandbox/executor.go | 101 ++++++++++++++------------ 1 file changed, 54 insertions(+), 47 deletions(-) diff --git a/internal/executor/sandbox/executor.go b/internal/executor/sandbox/executor.go index 786b03b..65c3369 100644 --- a/internal/executor/sandbox/executor.go +++ b/internal/executor/sandbox/executor.go @@ -10,13 +10,17 @@ import ( "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 tarThreshold = 128 * 1024 +const ( + tarSizeThreshold = 128 * 1024 // 128 KB + tarCountThreshold = 100 // 100 files +) func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { var err error @@ -28,13 +32,12 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { } } for i := 0; i < len(cmds); i += 1 { - cmd := &cmds[i] - if cmd.CopyIn == nil { + if cmd := &cmds[i]; cmd.CopyIn == nil { cmd.CopyIn = make(map[string]stage.CmdFile) } - for k, v := range cmd.CopyInCached { + for k, v := range cmds[i].CopyInCached { if fileID, ok := e.cachedMap[v]; ok { - cmd.CopyIn[k] = stage.CmdFile{FileID: &fileID} + cmds[i].CopyIn[k] = stage.CmdFile{FileID: &fileID} } } } @@ -49,8 +52,7 @@ func prepareTar(cmds []stage.Cmd) (bool, []byte) { return false, nil } for i := range cmds { - if cmds[i].CopyInDir != "" && - estimateCopyInSize(&cmds[i]) >= tarThreshold { + if shouldTar(&cmds[i]) { tarData, keysInTar := createCopyInTar(&cmds[i]) if tarData == nil { return false, nil @@ -68,6 +70,11 @@ func prepareTar(cmds []stage.Cmd) (bool, []byte) { 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 { @@ -100,35 +107,38 @@ func (e *Sandbox) runWithTar(cmds []stage.Cmd, tarData []byte) ([]stage.Executor 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 %s && exec \"$@\"", + "/bin/tar xf %s -C / --no-same-owner && rm -f %s && exec \"$@\"", tarFileName, tarFileName, ) - cmds[0].CopyIn[tarFileName] = stage.CmdFile{FileID: &fid} - cmds[0].Args = append([]string{ - "/bin/sh", "-c", script, "--", - }, cmds[0].Args...) + 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]) - results, err := e.runUnary(cmds) - if err != nil { - return nil, err - } - - 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) - } - - return results, nil + return e.runUnary(cmds) } -func estimateCopyInSize(cmd *stage.Cmd) int { - total := 0 +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 { @@ -140,7 +150,8 @@ func estimateCopyInSize(cmd *stage.Cmd) int { return nil } if _, exists := cmd.CopyIn[relPath]; !exists { - total += int(info.Size()) + totalSize += int(info.Size()) + totalCount++ } return nil }) @@ -151,13 +162,22 @@ func estimateCopyInSize(cmd *stage.Cmd) int { } if f.Src != nil { if fi, err := os.Stat(*f.Src); err == nil { - total += int(fi.Size()) + totalSize += int(fi.Size()) + totalCount++ } } else if f.Content != nil { - total += len(*f.Content) + totalSize += len(*f.Content) + totalCount++ } } - return total + 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) { @@ -193,20 +213,14 @@ func createCopyInTar(cmd *stage.Cmd) ([]byte, []string) { if err != nil { return err } - hdr.Name = relPath - if !filepath.IsAbs(hdr.Name) { - hdr.Name = "w/" + hdr.Name - } + hdr.Name = formatTarPath(relPath) return tw.WriteHeader(hdr) } hdr, err := tar.FileInfoHeader(info, "") if err != nil { return err } - hdr.Name = relPath - if !filepath.IsAbs(hdr.Name) { - hdr.Name = "w/" + hdr.Name - } + hdr.Name = formatTarPath(relPath) if info.IsDir() { hdr.Name += "/" } @@ -235,12 +249,8 @@ func createCopyInTar(cmd *stage.Cmd) ([]byte, []string) { continue } if f.Content != nil { - name := k - if !filepath.IsAbs(name) { - name = "w/" + name - } hdr := &tar.Header{ - Name: name, + Name: formatTarPath(k), Mode: 0o644, Size: int64(len(*f.Content)), } @@ -267,10 +277,7 @@ func createCopyInTar(cmd *stage.Cmd) ([]byte, []string) { slog.Error("create copyIn tar file info header", "key", k, "error", err) continue } - hdr.Name = k - if !filepath.IsAbs(hdr.Name) { - hdr.Name = "w/" + hdr.Name - } + hdr.Name = formatTarPath(k) if err := tw.WriteHeader(hdr); err != nil { slog.Error("create copyIn tar write header", "key", k, "error", err) continue