chore: test memory limits are being set on inner container (#44)

- Fixed an issue where large integer CLI arguments were not parsed correctly
- Added an integration test to validate CPU/memory are set correctly in
  the inner container.

Author: sreya

cli/cliflag/cliflag.go

@@ -110,13 +110,13 @@ func IntVarP(flagset *pflag.FlagSet, ptr *int, name string, shorthand string, en
 		return
 	}
 
-	vi64, err := strconv.ParseUint(val, 10, 8)
+	vi64, err := strconv.Atoi(val)
 	if err != nil {
 		flagset.IntVarP(ptr, name, shorthand, def, fmtUsage(usage, env))
 		return
 	}
 
-	flagset.IntVarP(ptr, name, shorthand, int(vi64), fmtUsage(usage, env))
+	flagset.IntVarP(ptr, name, shorthand, vi64, fmtUsage(usage, env))
 }
 
 func Bool(flagset *pflag.FlagSet, name, shorthand, env string, def bool, usage string) {

integration/docker_test.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -80,8 +81,10 @@ func TestDocker(t *testing.T) {
 		require.NoError(t, err)
 
 		var (
-			tmpdir = integrationtest.TmpDir(t)
-			binds  = integrationtest.DefaultBinds(t, tmpdir)
+			tmpdir              = integrationtest.TmpDir(t)
+			binds               = integrationtest.DefaultBinds(t, tmpdir)
+			expectedMemoryLimit = "1073741824"
+			expectedCPULimit    = 1
 		)
 
 		homeDir := filepath.Join(tmpdir, "home")
@@ -116,6 +119,8 @@ func TestDocker(t *testing.T) {
 				"TEST_VAR=hello=world",
 				// Add a mount mapping to the inner container.
 				fmt.Sprintf("%s=%s:%s,%s:%s:ro", cli.EnvMounts, "/home/coder", "/home/coder", "/var/secrets", "/var/secrets"),
+				fmt.Sprintf("%s=%s", cli.EnvMemory, expectedMemoryLimit),
+				fmt.Sprintf("%s=%d", cli.EnvCPUs, expectedCPULimit),
 			}
 		)
 
@@ -141,6 +146,7 @@ func TestDocker(t *testing.T) {
 			AddFUSE:         true,
 			AddTUN:          true,
 			BootstrapScript: bootstrapScript,
+			CPUs:            expectedCPULimit,
 		})
 
 		// Validate that the envs are set correctly.
@@ -227,6 +233,35 @@ func TestDocker(t *testing.T) {
 		})
 		require.NoError(t, err)
 		require.Equal(t, "1000", strings.TrimSpace(string(out)))
+
+		// Validate that the bootstrap script ran.
+		out, err = integrationtest.ExecInnerContainer(t, pool, integrationtest.ExecConfig{
+			ContainerID: resource.Container.ID,
+			Cmd:         []string{"cat", "/sys/fs/cgroup/memory/memory.limit_in_bytes"},
+		})
+		require.NoError(t, err)
+		require.Equal(t, expectedMemoryLimit, strings.TrimSpace(string(out)))
+
+		// Validate that the bootstrap script ran.
+		periodStr, err := integrationtest.ExecInnerContainer(t, pool, integrationtest.ExecConfig{
+			ContainerID: resource.Container.ID,
+			Cmd:         []string{"cat", "/sys/fs/cgroup/cpu,cpuacct/cpu.cfs_period_us"},
+		})
+		require.NoError(t, err)
+		period, err := strconv.ParseInt(strings.TrimSpace(string(periodStr)), 10, 64)
+		require.NoError(t, err)
+
+		// Validate that the bootstrap script ran.
+		quotaStr, err := integrationtest.ExecInnerContainer(t, pool, integrationtest.ExecConfig{
+			ContainerID: resource.Container.ID,
+			Cmd:         []string{"cat", "/sys/fs/cgroup/cpu/cpu.cfs_quota_us"},
+		})
+		require.NoError(t, err)
+		quota, err := strconv.ParseInt(strings.TrimSpace(string(quotaStr)), 10, 64)
+		require.NoError(t, err)
+
+		actualLimit := float64(quota) / float64(period)
+		require.Equal(t, expectedCPULimit, int(actualLimit))
 	})
 }
 

integration/integrationtest/docker.go

@@ -20,6 +20,7 @@ import (
 
 	"github.com/coder/envbox/buildlog"
 	"github.com/coder/envbox/cli"
+	"github.com/coder/envbox/dockerutil"
 	"github.com/coder/retry"
 )
 
@@ -47,6 +48,7 @@ type CreateDockerCVMConfig struct {
 	Mounts          []string
 	AddFUSE         bool
 	AddTUN          bool
+	CPUs            int
 }
 
 func (c CreateDockerCVMConfig) validate(t *testing.T) {
@@ -85,6 +87,8 @@ func RunEnvbox(t *testing.T, pool *dockertest.Pool, conf *CreateDockerCVMConfig)
 	}, func(host *docker.HostConfig) {
 		host.Binds = conf.Binds
 		host.Privileged = true
+		host.CPUPeriod = int64(dockerutil.DefaultCPUPeriod)
+		host.CPUQuota = int64(conf.CPUs) * int64(dockerutil.DefaultCPUPeriod)
 	})
 	require.NoError(t, err)
 	// t.Cleanup(func() { _ = pool.Purge(resource) })