From 4d00f5a6b9d572883bf66de02e405109359deeec Mon Sep 17 00:00:00 2001 From: Boming Zhang Date: Wed, 22 Jul 2026 04:04:00 -0700 Subject: [PATCH] fix(executor/sandbox): stream in tar file to avoid SOCK_SEQPACKET size limit --- internal/executor/sandbox/executor.go | 237 ++++++++++++++++++++++++++ 1 file changed, 237 insertions(+) diff --git a/internal/executor/sandbox/executor.go b/internal/executor/sandbox/executor.go index 987382f..127aeff 100644 --- a/internal/executor/sandbox/executor.go +++ b/internal/executor/sandbox/executor.go @@ -1,16 +1,26 @@ package sandbox import ( + "archive/tar" + "bytes" "context" "fmt" + "io" "log/slog" "maps" + "os" + "path/filepath" "github.com/criyle/go-judge/pb" "github.com/joint-online-judge/JOJ3/internal/stage" "google.golang.org/protobuf/proto" ) +const ( + tarStreamThreshold = 128 * 1024 + streamChunkSize = 32 * 1024 +) + func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { var err error if e.execClient == nil { @@ -32,6 +42,34 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { } } } + if needStream, tarData := prepareTarStream(cmds); needStream { + return e.runWithTarStream(cmds, tarData) + } + return e.runUnary(cmds) +} + +func prepareTarStream(cmds []stage.Cmd) (bool, []byte) { + if len(cmds) == 0 { + return false, nil + } + for i := range cmds { + if cmds[i].CopyInDir != "" && + estimateCopyInSize(&cmds[i]) >= tarStreamThreshold { + tarData, err := createCopyInTar(&cmds[i]) + if err != nil { + slog.Error("create copyIn tar", "error", err) + return false, nil + } + for j := range cmds { + cmds[j].CopyInDir = "" + } + return true, tarData + } + } + return false, nil +} + +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)) @@ -53,6 +91,205 @@ func (e *Sandbox) Run(cmds []stage.Cmd) ([]stage.ExecutorResult, error) { return results, nil } +func (e *Sandbox) runWithTarStream(cmds []stage.Cmd, tarData []byte) ([]stage.ExecutorResult, error) { + tarCmd := stage.Cmd{ + // FIXME: `/w` won't work if go-judge change default working directory + Args: []string{"/bin/tar", "xf", "-", "-C", "/w"}, + Env: cmds[0].Env, + Stdin: &stage.CmdFile{StreamIn: true}, + CPULimit: cmds[0].CPULimit, + ClockLimit: cmds[0].ClockLimit, + MemoryLimit: cmds[0].MemoryLimit, + StackLimit: cmds[0].StackLimit, + ProcLimit: cmds[0].ProcLimit, + CPURateLimit: cmds[0].CPURateLimit, + CPUSetLimit: cmds[0].CPUSetLimit, + } + + allCmds := append([]stage.Cmd{tarCmd}, cmds...) + pbCmds := convertPBCmd(allCmds) + + pbReq := &pb.Request{} + pbReq.SetCmd(pbCmds) + slog.Debug( + "sandbox stream execute", + "pbReq size", proto.Size(pbReq), + "tarSize", len(tarData), + ) + + stream, err := e.execClient.ExecStream(context.TODO()) + if err != nil { + return nil, fmt.Errorf("exec stream: %w", err) + } + + sr := &pb.StreamRequest{} + sr.SetExecRequest(pbReq) + if err := stream.Send(sr); err != nil { + return nil, fmt.Errorf("stream send request: %w", err) + } + + for offset := 0; offset < len(tarData); offset += streamChunkSize { + end := offset + streamChunkSize + if end > len(tarData) { + end = len(tarData) + } + input := pb.StreamRequest_builder{ + ExecInput: (&pb.StreamRequest_Input_builder{ + Index: 0, + Fd: 0, + Content: tarData[offset:end], + }).Build(), + }.Build() + if err := stream.Send(input); err != nil { + return nil, fmt.Errorf("stream send input: %w", err) + } + } + + if err := stream.CloseSend(); err != nil { + return nil, fmt.Errorf("stream close send: %w", err) + } + + var finalResponse *pb.Response + for { + resp, err := stream.Recv() + if err == io.EOF { + break + } + if err != nil { + return nil, fmt.Errorf("stream recv: %w", err) + } + if resp.HasExecResponse() { + finalResponse = resp.GetExecResponse() + } + } + + if finalResponse == nil { + return nil, fmt.Errorf("stream: no exec response received") + } + if finalResponse.GetError() != "" { + return nil, fmt.Errorf("stream execute error: %s", finalResponse.GetError()) + } + + results := convertPBResult(finalResponse.GetResults()) + if len(results) == 0 { + return nil, fmt.Errorf("stream: empty results") + } + tarResult := results[0] + if tarResult.Status != stage.StatusAccepted || tarResult.ExitStatus != 0 { + return nil, fmt.Errorf( + "tar extraction failed: status=%v, exit=%d, error=%s", + tarResult.Status, tarResult.ExitStatus, tarResult.Error, + ) + } + + results = results[1:] + for _, result := range results { + maps.Copy(e.cachedMap, result.FileIDs) + } + return results, nil +} + +func estimateCopyInSize(cmd *stage.Cmd) int { + total := 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 { + total += int(info.Size()) + } + 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 { + total += int(fi.Size()) + } + } else if f.Content != nil { + total += len(*f.Content) + } + } + return total +} + +func createCopyInTar(cmd *stage.Cmd) ([]byte, error) { + var buf bytes.Buffer + tw := tar.NewWriter(&buf) + + 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 = relPath + return tw.WriteHeader(hdr) + } + hdr, err := tar.FileInfoHeader(info, "") + if err != nil { + return err + } + hdr.Name = 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 { + return nil, fmt.Errorf("walk copyInDir: %w", err) + } + } + + if err := tw.Close(); err != nil { + return nil, fmt.Errorf("close tar: %w", err) + } + return buf.Bytes(), nil +} + func (e *Sandbox) Cleanup() error { for k, fileID := range e.cachedMap { req := &pb.FileID{}