feat: support -no-fallback cli flag (#151)

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)"`

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))

env/config.go

@@ -14,4 +14,5 @@ type Config struct {
 	ContainerCredStart int
 	EnableCPURate      bool
 	CPUCfsPeriod       time.Duration
+	NoFallback         bool
 }

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,21 +99,27 @@ 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 {
 			logger.Error("failed to create cgroup", zap.String("prefix", prefix), zap.Error(err))
 			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))
 	if _, err = cgb.Nest("api"); err != nil {
 		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")