internal/executor/sandbox/executor.go

@@ -1,16 +1,27 @@
 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 {
@@ -20,18 +31,51 @@ 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 {
-		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}
 			}
 		}
 	}
+	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))
@@ -53,6 +97,213 @@ func (e *Sandbox) Run(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{}