diff --git a/internal/executor/sandbox/executor.go b/internal/executor/sandbox/executor.go index 3c94326..c037bec 100644 --- a/internal/executor/sandbox/executor.go +++ b/internal/executor/sandbox/executor.go @@ -6,6 +6,8 @@ import ( "fmt" "log/slog" "maps" + "math" + "time" "github.com/criyle/go-judge/pb" "github.com/joint-online-judge/JOJ3/internal/stage" @@ -43,7 +45,7 @@ func (e *Sandbox) Run(ctx context.Context, cmds []stage.Cmd) ([]stage.ExecutorRe pbReq := &pb.Request{} pbReq.SetCmd(pbCmds) slog.Debug("sandbox execute", "pbReq size", proto.Size(pbReq)) - callCtx, cancel := context.WithTimeout(ctx, e.timeout) + callCtx, cancel := context.WithTimeout(ctx, execRPCTimeout(cmds)) defer cancel() pbRet, err := e.execClient.Exec(callCtx, pbReq) if err != nil { @@ -64,7 +66,7 @@ func (e *Sandbox) Cleanup(ctx context.Context) error { for k, fileID := range e.cachedMap { req := &pb.FileID{} req.SetFileID(fileID) - callCtx, cancel := context.WithTimeout(ctx, e.timeout) + callCtx, cancel := context.WithTimeout(ctx, rpcTimeoutMargin) _, err := e.execClient.FileDelete(callCtx, req) cancel() if err != nil { @@ -80,3 +82,29 @@ func (e *Sandbox) Cleanup(ctx context.Context) error { } return cleanupErr } + +func execRPCTimeout(cmds []stage.Cmd) time.Duration { + var maxLimit uint64 + for _, cmd := range cmds { + limit := cmd.ClockLimit + if limit == 0 { + // Match the local executor's default wall-clock allowance when only a + // CPU limit is specified. + if cmd.CPULimit > math.MaxUint64/2 { + limit = math.MaxUint64 + } else { + limit = cmd.CPULimit * 2 + } + } + if limit > maxLimit { + maxLimit = limit + } + } + + margin := uint64(rpcTimeoutMargin) + if maxLimit > uint64(math.MaxInt64)-margin { + return time.Duration(math.MaxInt64) + } + // The bound above guarantees the conversion fits in time.Duration. + return time.Duration(maxLimit + margin) // #nosec G115 +} diff --git a/internal/executor/sandbox/executor_test.go b/internal/executor/sandbox/executor_test.go new file mode 100644 index 0000000..0a164aa --- /dev/null +++ b/internal/executor/sandbox/executor_test.go @@ -0,0 +1,42 @@ +package sandbox + +import ( + "math" + "testing" + "time" + + "github.com/joint-online-judge/JOJ3/internal/stage" +) + +func TestExecRPCTimeout(t *testing.T) { + tests := []struct { + name string + cmds []stage.Cmd + want time.Duration + }{ + {name: "margin only", want: rpcTimeoutMargin}, + { + name: "largest clock limit", + cmds: []stage.Cmd{{ClockLimit: uint64(time.Minute)}, {ClockLimit: uint64(2 * time.Minute)}}, + want: 2*time.Minute + rpcTimeoutMargin, + }, + { + name: "cpu limit fallback", + cmds: []stage.Cmd{{CPULimit: uint64(time.Minute)}}, + want: 2*time.Minute + rpcTimeoutMargin, + }, + { + name: "duration overflow", + cmds: []stage.Cmd{{ClockLimit: math.MaxUint64}}, + want: time.Duration(math.MaxInt64), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := execRPCTimeout(tt.cmds); got != tt.want { + t.Fatalf("execRPCTimeout() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/internal/executor/sandbox/meta.go b/internal/executor/sandbox/meta.go index f327eec..0fe91c8 100644 --- a/internal/executor/sandbox/meta.go +++ b/internal/executor/sandbox/meta.go @@ -19,15 +19,15 @@ type Sandbox struct { cachedMap map[string]string execClient pb.ExecutorClient conn *grpc.ClientConn - timeout time.Duration } +const rpcTimeoutMargin = 30 * time.Second + func init() { stage.RegisterExecutor(name, &Sandbox{ execServer: "localhost:5051", token: "", cachedMap: make(map[string]string), - timeout: 30 * time.Second, }) } @@ -37,6 +37,5 @@ func InitWithConf(execServer, token string) { execServer: execServer, token: token, cachedMap: make(map[string]string), - timeout: 30 * time.Second, }) }