From 1dfe296820969f11e4fd61bd0eef23f65432dd03 Mon Sep 17 00:00:00 2001
From: Paul Pan <i@0x7f.app>
Date: Sun, 28 Apr 2024 01:11:57 +0800
Subject: [PATCH] feat: allow to configure runner concurrency

---
 config.docker.yaml                      | 1 +
 docker-entrypoint.sh                    | 1 +
 internal/model/config.go                | 3 ++-
 internal/service/runner/service.go      | 6 +++++-
 internal/service/runner/service_test.go | 5 ++++-
 5 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/config.docker.yaml b/config.docker.yaml
index 0797be1..6d7d379 100644
--- a/config.docker.yaml
+++ b/config.docker.yaml
@@ -39,6 +39,7 @@ Storage:
 
 Runner:
   CGroup: ${CGROUP_PATH}
+  Concurrency: ${RUNNER_CONCURRENCY}
 
 Metrics:
   Namespace: ${METRICS_NAMESPACE}
diff --git a/docker-entrypoint.sh b/docker-entrypoint.sh
index c4c73fa..b8af094 100755
--- a/docker-entrypoint.sh
+++ b/docker-entrypoint.sh
@@ -75,6 +75,7 @@ function extract_storage() {
 
 function extract_runner() {
     check_env "CGROUP_PATH" "/sys/fs/cgroup/nsjail" true
+    check_env "RUNNER_CONCURRENCY" 0 false
 }
 
 function extract_metrics() {
diff --git a/internal/model/config.go b/internal/model/config.go
index e6b5453..98c5032 100644
--- a/internal/model/config.go
+++ b/internal/model/config.go
@@ -50,7 +50,8 @@ type ConfigStorage struct {
 }
 
 type ConfigRunner struct {
-	CGroup string `yaml:"CGroup"`
+	CGroup      string `yaml:"CGroup"`
+	Concurrency int    `yaml:"Concurrency"`
 }
 
 type ConfigMetrics struct {
diff --git a/internal/service/runner/service.go b/internal/service/runner/service.go
index 394ce67..3abaa71 100644
--- a/internal/service/runner/service.go
+++ b/internal/service/runner/service.go
@@ -38,9 +38,13 @@ type Service interface {
 }
 
 func NewService(i *do.Injector) (Service, error) {
-	concurrency := utils.If(runtime.NumCPU() > 1, runtime.NumCPU()-1, 1)
 	cfg := do.MustInvoke[config.Service](i).GetConfig()
 
+	concurrency := cfg.Runner.Concurrency
+	if concurrency <= 0 {
+		concurrency = utils.If(runtime.NumCPU() > 1, runtime.NumCPU()-1, 1)
+	}
+
 	srv := &service{
 		log:     do.MustInvoke[log.Service](i).GetLogger("runner"),
 		pool:    pool.NewTaskPool(concurrency, concurrency),
diff --git a/internal/service/runner/service_test.go b/internal/service/runner/service_test.go
index ebfa729..70a3cd9 100644
--- a/internal/service/runner/service_test.go
+++ b/internal/service/runner/service_test.go
@@ -24,7 +24,10 @@ func GetService(dev bool) Service {
 
 	cfg := &FakeConfigService{conf: model.Config{
 		Development: dev,
-		Runner:      model.ConfigRunner{CGroup: "/sys/fs/cgroup/nsjail"},
+		Runner: model.ConfigRunner{
+			CGroup:      "/sys/fs/cgroup/nsjail",
+			Concurrency: 0,
+		},
 	}}
 
 	do.ProvideValue[config.Service](injector, cfg)