Add environment variables to control container resource constraints (COMPLEMENT_CONTAINER_CPU_CORES, COMPLEMENT_CONTAINER_MEMORY)

ENVIRONMENT.md

@@ -16,6 +16,16 @@ If 1, always prints the Homeserver container logs even on success. When used wit
 This allows you to override the base image used for a particular named homeserver. For example, `COMPLEMENT_BASE_IMAGE_HS1=complement-dendrite:latest` would use `complement-dendrite:latest` for the `hs1` homeserver in blueprints, but not any other homeserver (e.g `hs2`). This matching is case-insensitive. This allows Complement to test how different homeserver implementations work with each other.  
 - Type: `map[string]string`
 
+#### `COMPLEMENT_CONTAINER_CPU_CORES`
+The number of CPU cores available for the container to use (can be fractional like 0.5). This is passed to Docker as the `--cpus`/`NanoCPUs` argument. If 0, no limit is set and the container can use all available host CPUs. This is useful to mimic a resource-constrained environment, like a CI environment.  
+- Type: `float64`
+- Default: 0
+
+#### `COMPLEMENT_CONTAINER_MEMORY`
+The maximum amount of memory the container can use (ex. "1GB"). Valid units are "B", (decimal: "KB", "MB", "GB, "TB, "PB"), (binary: "KiB", "MiB", "GiB", "TiB", "PiB") or no units (bytes) (case-insensitive). We also support "K", "M", "G" as per Docker's CLI. The number of bytes is passed to Docker as the `--memory`/`Memory` argument. If 0, no limit is set and the container can use all available host memory. This is useful to mimic a resource-constrained environment, like a CI environment.  
+- Type: `int64`
+- Default: 0
+
 #### `COMPLEMENT_DEBUG`
 If 1, prints out more verbose logging such as HTTP request/response bodies.  
 - Type: `bool`

cmd/gendoc/main.go

@@ -1,3 +1,5 @@
+// Usage: `go run ./cmd/gendoc --config config/config.go > ENVIRONMENT.md`
+
 package main
 
 import (

config/config.go

@@ -11,6 +11,7 @@ import (
 	"math/big"
 	"os"
 	"regexp"
+	"sort"
 	"strconv"
 	"strings"
 	"time"
@@ -52,6 +53,23 @@ type Complement struct {
 	// starting the container. Responsiveness is detected by `HEALTHCHECK` being healthy *and*
 	// the `/versions` endpoint returning 200 OK.
 	SpawnHSTimeout time.Duration
+	// Name: COMPLEMENT_CONTAINER_CPU_CORES
+	// Default: 0
+	// Description: The number of CPU cores available for the container to use (can be
+	// fractional like 0.5). This is passed to Docker as the `--cpus`/`NanoCPUs` argument.
+	// If 0, no limit is set and the container can use all available host CPUs. This is
+	// useful to mimic a resource-constrained environment, like a CI environment.
+	ContainerCPUCores float64
+	// Name: COMPLEMENT_CONTAINER_MEMORY
+	// Default: 0
+	// Description: The maximum amount of memory the container can use (ex. "1GB"). Valid
+	// units are "B", (decimal: "KB", "MB", "GB, "TB, "PB"), (binary: "KiB", "MiB", "GiB",
+	// "TiB", "PiB") or no units (bytes) (case-insensitive). We also support "K", "M", "G"
+	// as per Docker's CLI. The number of bytes is passed to Docker as the
+	// `--memory`/`Memory` argument. If 0, no limit is set and the container can use all
+	// available host memory. This is useful to mimic a resource-constrained environment,
+	// like a CI environment.
+	ContainerMemoryBytes int64
 	// Name: COMPLEMENT_KEEP_BLUEPRINTS
 	// Description: A list of space separated blueprint names to not clean up after running. For example,
 	// `one_to_one_room alice` would not delete the homeserver images for the blueprints `alice` and
@@ -145,8 +163,13 @@ func NewConfigFromEnvVars(pkgNamespace, baseImageURI string) *Complement {
 		// each iteration had a 50ms sleep between tries so the timeout is 50 * iteration ms
 		cfg.SpawnHSTimeout = time.Duration(50*parseEnvWithDefault("COMPLEMENT_VERSION_CHECK_ITERATIONS", 100)) * time.Millisecond
 	}
+	cfg.ContainerCPUCores = parseEnvAsFloatWithDefault("COMPLEMENT_CONTAINER_CPU_CORES", 0)
+	parsedMemoryBytes, err := parseByteSizeString(os.Getenv("COMPLEMENT_CONTAINER_MEMORY"))
+	if err != nil {
+		panic("COMPLEMENT_CONTAINER_MEMORY parse error: " + err.Error())
+	}
+	cfg.ContainerMemoryBytes = parsedMemoryBytes
 	cfg.KeepBlueprints = strings.Split(os.Getenv("COMPLEMENT_KEEP_BLUEPRINTS"), " ")
-	var err error
 	hostMounts := os.Getenv("COMPLEMENT_HOST_MOUNTS")
 	if hostMounts != "" {
 		cfg.HostMounts, err = newHostMounts(strings.Split(hostMounts, ";"))
@@ -214,17 +237,132 @@ func (c *Complement) CAPrivateKeyBytes() ([]byte, error) {
 	return caKey.Bytes(), err
 }
 
-func parseEnvWithDefault(key string, def int) int {
-	s := os.Getenv(key)
-	if s != "" {
-		i, err := strconv.Atoi(s)
-		if err != nil {
-			// Don't bother trying to report it
-			return def
+func parseEnvWithDefault(key string, defaultValue int) int {
+	inputString := os.Getenv(key)
+	if inputString == "" {
+		return defaultValue
+	}
+
+	parsedNumber, err := strconv.Atoi(inputString)
+	if err != nil {
+		panic(key + " parse error: " + err.Error())
+	}
+	return parsedNumber
+}
+
+func parseEnvAsFloatWithDefault(key string, defaultValue float64) float64 {
+	inputString := os.Getenv(key)
+	if inputString == "" {
+		return defaultValue
+	}
+
+	parsedNumber, err := strconv.ParseFloat(inputString, 64)
+	if err != nil {
+		panic(key + " parse error: " + err.Error())
+	}
+	return parsedNumber
+}
+
+// parseByteSizeString parses a byte size string (case insensitive) like "512MB"
+// or "2GB" into bytes. If the string is empty, 0 is returned. Returns an error if the
+// string does not match one of the valid units or is an invalid integer.
+//
+// Valid units are "B", (decimal: "KB", "MB", "GB, "TB, "PB"), (binary: "KiB", "MiB",
+// "GiB", "TiB", "PiB") or no units (bytes). We also support "K", "M", "G" as per
+// Docker's CLI.
+func parseByteSizeString(inputString string) (int64, error) {
+	// Strip spaces and normalize to lowercase
+	normalizedString := strings.TrimSpace(strings.ToLower(inputString))
+	if normalizedString == "" {
+		return 0, nil
+	}
+	unitToByteMultiplierMap := map[string]int64{
+		// No unit (bytes)
+		"":    1,
+		"b":   1,
+		"kb":  intPow(10, 3),
+		"mb":  intPow(10, 6),
+		"gb":  intPow(10, 9),
+		"tb":  intPow(10, 12),
+		"kib": 1024,
+		"mib": intPow(1024, 2),
+		"gib": intPow(1024, 3),
+		"tib": intPow(1024, 4),
+		// These are also supported to match Docker's CLI
+		"k": 1024,
+		"m": intPow(1024, 2),
+		"g": intPow(1024, 3),
+	}
+	availableUnitsSorted := make([]string, 0, len(unitToByteMultiplierMap))
+	for unit := range unitToByteMultiplierMap {
+		availableUnitsSorted = append(availableUnitsSorted, unit)
+	}
+	// Sort units by length descending so that longer units are matched first
+	// (e.g "mib" before "b")
+	sort.Slice(availableUnitsSorted, func(i, j int) bool {
+		return len(availableUnitsSorted[i]) > len(availableUnitsSorted[j])
+	})
+
+	// Find the number part of the string and the unit used
+	numberPart := ""
+	byteUnit := ""
+	byteMultiplier := int64(0)
+	for _, unit := range availableUnitsSorted {
+		if strings.HasSuffix(normalizedString, unit) {
+			byteUnit = unit
+			// Handle the case where there is a space between the number and the unit (e.g "512 MB")
+			numberPart = strings.TrimSpace(normalizedString[:len(normalizedString)-len(unit)])
+			byteMultiplier = unitToByteMultiplierMap[unit]
+			break
 		}
-		return i
 	}
-	return def
+
+	// Failed to find a valid unit
+	if byteUnit == "" {
+		return 0, fmt.Errorf("parseByteSizeString: invalid byte unit used in string: %s (supported units: %s)",
+			inputString,
+			strings.Join(availableUnitsSorted, ", "),
+		)
+	}
+	// Assert to sanity check our logic above is sound
+	if byteMultiplier == 0 {
+		panic(fmt.Sprintf(
+			"parseByteSizeString: byteMultiplier is unexpectedly 0 for unit: %s. "+
+				"This is probably a problem with the function itself.", byteUnit,
+		))
+	}
+
+	// Parse the number part as an int64
+	parsedNumber, err := strconv.ParseInt(strings.TrimSpace(numberPart), 10, 64)
+	if err != nil {
+		return 0, fmt.Errorf("parseByteSizeString: failed to parse number part of string: %s (%w)",
+			numberPart,
+			err,
+		)
+	}
+
+	// Calculate the total bytes
+	totalBytes := parsedNumber * byteMultiplier
+	return totalBytes, nil
+}
+
+// intPow calculates n to the mth power. Since the result is an int, it is assumed that m is a positive power
+//
+// via https://stackoverflow.com/questions/64108933/how-to-use-math-pow-with-integers-in-go/66429580#66429580
+func intPow(n, m int64) int64 {
+	if m == 0 {
+		return 1
+	}
+
+	if m == 1 {
+		return n
+	}
+
+	result := n
+	for i := int64(2); i <= m; i++ {
+		result *= n
+	}
+	return result
 }
 
 func newHostMounts(mounts []string) ([]HostMount, error) {

internal/docker/deployer.go

@@ -399,6 +399,18 @@ func deployImage(
 		PublishAllPorts: true,
 		ExtraHosts:      extraHosts,
 		Mounts:          mounts,
+		// https://docs.docker.com/engine/containers/resource_constraints/
+		Resources: container.Resources{
+			// Constrain the the number of CPU cores this container can use
+			//
+			// The number of CPU cores in 1e9 increments
+			//
+			// `NanoCPUs` is the option that is "Applicable to all platforms" instead of
+			// `CPUPeriod`/`CPUQuota` (Unix only) or `CPUCount`/`CPUPercent` (Windows only).
+			NanoCPUs: int64(cfg.ContainerCPUCores * 1e9),
+			// Constrain the maximum memory the container can use
+			Memory: cfg.ContainerMemoryBytes,
+		},
 	}, &network.NetworkingConfig{
 		EndpointsConfig: map[string]*network.EndpointSettings{
 			networkName: {
@@ -415,7 +427,20 @@ func deployImage(
 
 	containerID := body.ID
 	if cfg.DebugLoggingEnabled {
-		log.Printf("%s: Created container '%s' using image '%s' on network '%s'", contextStr, containerID, imageID, networkName)
+		constraintStrings := []string{}
+		if cfg.ContainerCPUCores > 0 {
+			constraintStrings = append(constraintStrings, fmt.Sprintf("%.1f CPU cores", cfg.ContainerCPUCores))
+		}
+		if cfg.ContainerMemoryBytes > 0 {
+			// TODO: It would be nice to pretty print this in MB/GB etc.
+			constraintStrings = append(constraintStrings, fmt.Sprintf("%d bytes of memory", cfg.ContainerMemoryBytes))
+		}
+		constrainedResourcesDisplayString := ""
+		if len(constraintStrings) > 0 {
+			constrainedResourcesDisplayString = fmt.Sprintf("(%s)", strings.Join(constraintStrings, ", "))
+		}
+
+		log.Printf("%s: Created container '%s' using image '%s' on network '%s' %s", contextStr, containerID, imageID, networkName, constrainedResourcesDisplayString)
 	}
 	stubDeployment := &HomeserverDeployment{
 		ContainerID: containerID,