From 156a6d0c95dadb3883cd9284061dd8b14930c3f2 Mon Sep 17 00:00:00 2001
From: Robert Jandow <robert.jandow@tum.de>
Date: Sat, 27 Jan 2024 19:53:24 +0100
Subject: [PATCH] Add CPU and RAM limits for Docker containers

---
 HadesScheduler/docker/docker.go | 15 +++++++++++++++
 shared/utils/config.go          | 22 ++++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/HadesScheduler/docker/docker.go b/HadesScheduler/docker/docker.go
index df29674..eee0017 100644
--- a/HadesScheduler/docker/docker.go
+++ b/HadesScheduler/docker/docker.go
@@ -27,6 +27,8 @@ type DockerConfig struct {
 	DockerHost           string `env:"DOCKER_HOST" envDefault:"unix:///var/run/docker.sock"`
 	ContainerAutoremove  bool   `env:"CONTAINER_AUTOREMOVE" envDefault:"true"`
 	DockerScriptExecutor string `env:"DOCKER_SCRIPT_EXECUTOR" envDefault:"/bin/bash -c"`
+	CPU_limit            uint   `env:"DOCKER_CPU_LIMIT"`
+	RAM_limit            string `env:"DOCKER_RAM_LIMIT"`
 }
 
 func init() {
@@ -104,6 +106,19 @@ func executeStep(ctx context.Context, client *client.Client, step payload.Step,
 		AutoRemove: container_autoremove, // Remove the container after it is done only if the config is set to true
 	}
 
+	// Limit the resource usage of the containers
+	if DockerCfg.CPU_limit != 0 {
+		host_config.Resources.NanoCPUs = int64(DockerCfg.CPU_limit * 1e9)
+	}
+	if DockerCfg.RAM_limit != "" {
+		bytes, err := utils.ParseMemoryLimit(DockerCfg.RAM_limit)
+		if err != nil {
+			log.WithError(err).Errorf("Failed to parse RAM limit %s", DockerCfg.RAM_limit)
+		} else {
+			host_config.Resources.Memory = bytes
+		}
+	}
+
 	// Create the bash script if there is one
 	if step.Script != "" {
 		// Overwrite the default entrypoint
diff --git a/shared/utils/config.go b/shared/utils/config.go
index bb21924..6ee4c46 100644
--- a/shared/utils/config.go
+++ b/shared/utils/config.go
@@ -1,6 +1,10 @@
 package utils
 
 import (
+	"fmt"
+	"strconv"
+	"strings"
+
 	"github.com/caarlos0/env/v9"
 	"github.com/joho/godotenv"
 	log "github.com/sirupsen/logrus"
@@ -32,3 +36,21 @@ func LoadConfig(cfg interface{}) {
 
 	log.Debug("Config loaded: ", cfg)
 }
+
+func ParseMemoryLimit(limit string) (int64, error) {
+	unit := limit[len(limit)-2:]
+	number := limit[:len(limit)-2]
+	value, err := strconv.ParseInt(number, 10, 64)
+	if err != nil {
+		return 0, err
+	}
+
+	switch strings.ToUpper(unit) {
+	case "GB", "G":
+		return value * 1024 * 1024 * 1024, nil
+	case "MB", "M":
+		return value * 1024 * 1024, nil
+	default:
+		return 0, fmt.Errorf("unknown unit: %s", unit)
+	}
+}