From e7b19ce8aeadd81b061b77f5e469f43abb6ec1ab Mon Sep 17 00:00:00 2001 From: BoYanZh Date: Mon, 9 Jun 2025 16:57:01 -0400 Subject: [PATCH] feat: support -no-fallback cli flag (#151) --- cmd/go-judge/config/config.go | 1 + cmd/go-judge/main.go | 9 ++++++--- env/config.go | 1 + env/env_cgroup_linux.go | 16 ++++++++++++++-- 4 files changed, 22 insertions(+), 5 deletions(-) diff --git a/cmd/go-judge/config/config.go b/cmd/go-judge/config/config.go index ace6079..0a42fad 100644 --- a/cmd/go-judge/config/config.go +++ b/cmd/go-judge/config/config.go @@ -21,6 +21,7 @@ type Config struct { Parallelism int `flagUsage:"control the # of concurrency execution (default equal to number of cpu)"` CgroupPrefix string `flagUsage:"control cgroup prefix" default:"gojudge"` ContainerCredStart int `flagUsage:"control the start uid&gid for container (0 uses unprivileged root)" default:"0"` + NoFallback bool `flagUsage:"exit if fallback to rlimit / rusage mode"` // file store SrcPrefix []string `flagUsage:"specifies directory prefix for source type copyin (example: -src-prefix=/home,/usr)"` diff --git a/cmd/go-judge/main.go b/cmd/go-judge/main.go index 58277c2..2be5fc7 100644 --- a/cmd/go-judge/main.go +++ b/cmd/go-judge/main.go @@ -149,8 +149,10 @@ func loadConf() *config.Config { return &conf } -type stopFunc func(ctx context.Context) error -type initFunc func() (start func(), cleanUp stopFunc) +type ( + stopFunc func(ctx context.Context) error + initFunc func() (start func(), cleanUp stopFunc) +) func cleanUpWorker(work worker.Worker) initFunc { return func() (start func(), cleanUp stopFunc) { @@ -474,7 +476,7 @@ func newFilsStore(conf *config.Config) (filestore.FileStore, func() error) { return os.RemoveAll(conf.Dir) } } - os.MkdirAll(conf.Dir, 0755) + os.MkdirAll(conf.Dir, 0o755) fs = filestore.NewFileLocalStore(conf.Dir) if conf.EnableMetrics { fs = newMetricsFileStore(fs) @@ -497,6 +499,7 @@ func newEnvBuilder(conf *config.Config) (pool.EnvBuilder, map[string]any) { EnableCPURate: conf.EnableCPURate, CPUCfsPeriod: conf.CPUCfsPeriod, SeccompConf: conf.SeccompConf, + NoFallback: conf.NoFallback, }, logger) if err != nil { logger.Fatal("create environment builder failed ", zap.Error(err)) diff --git a/env/config.go b/env/config.go index 7eebf71..07ec29c 100644 --- a/env/config.go +++ b/env/config.go @@ -14,4 +14,5 @@ type Config struct { ContainerCredStart int EnableCPURate bool CPUCfsPeriod time.Duration + NoFallback bool } diff --git a/env/env_cgroup_linux.go b/env/env_cgroup_linux.go index e7830d5..813c2f7 100644 --- a/env/env_cgroup_linux.go +++ b/env/env_cgroup_linux.go @@ -29,7 +29,7 @@ func setupCgroup(c Config, logger *zap.Logger) (cgroup.Cgroup, *cgroup.Controlle } } - return createAndNestCgroup(prefix, ct, logger) + return createAndNestCgroup(prefix, ct, c.NoFallback, logger) } func setupCgroupV2(prefix string, logger *zap.Logger) (string, *cgroup.Controllers, error) { @@ -99,7 +99,7 @@ func getControllersWithPrefix(prefix string, logger *zap.Logger) *cgroup.Control return ct } -func createAndNestCgroup(prefix string, ct *cgroup.Controllers, logger *zap.Logger) (cgroup.Cgroup, *cgroup.Controllers, error) { +func createAndNestCgroup(prefix string, ct *cgroup.Controllers, noFallback bool, logger *zap.Logger) (cgroup.Cgroup, *cgroup.Controllers, error) { cgb, err := cgroup.New(prefix, ct) if err != nil { if os.Getuid() == 0 { @@ -107,6 +107,9 @@ func createAndNestCgroup(prefix string, ct *cgroup.Controllers, logger *zap.Logg return nil, nil, err } logger.Warn("not running in root and have no permission on cgroup, falling back to rlimit / rusage mode", zap.Error(err)) + if noFallback { + return nil, nil, fmt.Errorf("failed to create cgroup with no fallback: %w", err) + } return nil, nil, nil } logger.Info("creating nesting api cgroup", zap.Any("cgroup", cgb)) @@ -114,6 +117,9 @@ func createAndNestCgroup(prefix string, ct *cgroup.Controllers, logger *zap.Logg if os.Getuid() != 0 { logger.Warn("creating api cgroup with error, falling back to rlimit / rusage mode", zap.Error(err)) cgb.Destroy() + if noFallback { + return nil, nil, fmt.Errorf("failed to create nesting api cgroup with no fallback: %w", err) + } return nil, nil, nil } } @@ -122,10 +128,16 @@ func createAndNestCgroup(prefix string, ct *cgroup.Controllers, logger *zap.Logg cg, err := cgb.New("containers") if err != nil { logger.Warn("creating containers cgroup with error, falling back to rlimit / rusage mode", zap.Error(err)) + if noFallback { + return nil, nil, fmt.Errorf("failed to create containers cgroup with no fallback: %w", err) + } return nil, nil, nil } if ct != nil && !ct.Memory { logger.Warn("memory cgroup is not enabled, falling back to rlimit / rusage mode") + if noFallback { + return nil, nil, fmt.Errorf("memory cgroup is not enabled with no fallback: %w", err) + } } if ct != nil && !ct.Pids { logger.Warn("pid cgroup is not enabled, proc limit does not have effect")