build: Add per-step CPU and memory resource limits

Port of moby/buildkit#6569 to buildx. Adds --memory, --memory-swap, --cpu-shares, --cpu-period, --cpu-quota, --cpuset-cpus, and --cpuset-mems flags to build, plus the equivalent bake target attributes and compose x-bake fields. Signed-off-by: Jiří Moravčík <jiri.moravcik@gmail.com>
2026-06-08 15:39:14 +02:00
parent 4fbc1e5f2c
commit 80b2293881
14 changed files with 761 additions and 48 deletions
@@ -87,6 +87,7 @@ type Options struct {
 	Tags                       []string
 	Target                     string
 	Ulimits                    *opts.UlimitOpt
+	ResourceLimits             ResourceLimits

 	Session                []session.Attachable
 	Linked                 bool // Linked marks this target as exclusively linked (not requested by the user).
@@ -98,6 +99,18 @@ type Options struct {
 	Policy                 []buildflags.PolicyConfig
 }

+// ResourceLimits holds the cgroup resource constraints applied to individual
+// build steps (RUN instructions). They don't affect the build cache key.
+type ResourceLimits struct {
+	Memory     opts.MemBytes
+	MemorySwap opts.MemSwapBytes
+	CPUShares  int64
+	CPUPeriod  int64
+	CPUQuota   int64
+	CPUSetCPUs string
+	CPUSetMems string
+}
+
 type CallFunc struct {
 	Name         string
 	Format       string
@@ -610,6 +610,9 @@ func toSolveOpt(ctx context.Context, np *noderesolver.ResolvedNode, multiDriver
 		so.FrontendAttrs["ulimit"] = ulimits
 	}

+	// setup per-step resource limits
+	addResourceLimits(opt.ResourceLimits, so.FrontendAttrs)
+
 	// mark call request as internal
 	if opt.CallFunc != nil {
 		so.Internal = true
@@ -99,6 +99,97 @@ func toBuildkitUlimits(inp *opts.UlimitOpt) (string, error) {
 	return strings.Join(ulimits, ","), nil
 }

+// User-facing resource keys accepted in `--resource key=value` entries, mirroring docker run flag names.
+const (
+	resourceKeyMemory     = "memory"
+	resourceKeyMemorySwap = "memory-swap"
+	resourceKeyCPUShares  = "cpu-shares"
+	resourceKeyCPUPeriod  = "cpu-period"
+	resourceKeyCPUQuota   = "cpu-quota"
+	resourceKeyCPUSetCPUs = "cpuset-cpus"
+	resourceKeyCPUSetMems = "cpuset-mems"
+)
+
+// Frontend attribute keys, must match those parsed by BuildKit's dockerui frontend.
+const (
+	attrMemory     = "memory"
+	attrMemorySwap = "memswap"
+	attrCPUShares  = "cpushares"
+	attrCPUPeriod  = "cpuperiod"
+	attrCPUQuota   = "cpuquota"
+	attrCPUSetCPUs = "cpusetcpus"
+	attrCPUSetMems = "cpusetmems"
+)
+
+// ParseResourceLimits parses `key=value` entries from the `--resource` flag into ResourceLimits.
+func ParseResourceLimits(entries []string) (ResourceLimits, error) {
+	var rl ResourceLimits
+	for _, entry := range entries {
+		k, v, ok := strings.Cut(entry, "=")
+		if !ok {
+			return rl, errors.Errorf("invalid resource %q, expected key=value", entry)
+		}
+		k = strings.TrimSpace(k)
+		v = strings.TrimSpace(v)
+		switch k {
+		case resourceKeyMemory:
+			if err := rl.Memory.Set(v); err != nil {
+				return rl, errors.Wrapf(err, "invalid value %q for resource %s", v, k)
+			}
+		case resourceKeyMemorySwap:
+			if err := rl.MemorySwap.Set(v); err != nil {
+				return rl, errors.Wrapf(err, "invalid value %q for resource %s", v, k)
+			}
+		case resourceKeyCPUShares, resourceKeyCPUPeriod, resourceKeyCPUQuota:
+			n, err := strconv.ParseInt(v, 10, 64)
+			if err != nil {
+				return rl, errors.Wrapf(err, "invalid value %q for resource %s", v, k)
+			}
+			switch k {
+			case resourceKeyCPUShares:
+				rl.CPUShares = n
+			case resourceKeyCPUPeriod:
+				rl.CPUPeriod = n
+			case resourceKeyCPUQuota:
+				rl.CPUQuota = n
+			}
+		case resourceKeyCPUSetCPUs:
+			rl.CPUSetCPUs = v
+		case resourceKeyCPUSetMems:
+			rl.CPUSetMems = v
+		default:
+			return rl, errors.Errorf("unknown resource %q", k)
+		}
+	}
+	return rl, nil
+}
+
+// addResourceLimits sets the frontend attributes for the resource limits.
+// Only non-zero values are sent, so builds against a daemon without the feature keep working.
+func addResourceLimits(rl ResourceLimits, attrs map[string]string) {
+	if v := rl.Memory.Value(); v > 0 {
+		attrs[attrMemory] = strconv.FormatInt(v, 10)
+	}
+	if v := rl.MemorySwap.Value(); v != 0 {
+		attrs[attrMemorySwap] = strconv.FormatInt(v, 10)
+	}
+	if rl.CPUShares > 0 {
+		attrs[attrCPUShares] = strconv.FormatInt(rl.CPUShares, 10)
+	}
+	if rl.CPUPeriod > 0 {
+		attrs[attrCPUPeriod] = strconv.FormatInt(rl.CPUPeriod, 10)
+	}
+	if rl.CPUQuota > 0 {
+		attrs[attrCPUQuota] = strconv.FormatInt(rl.CPUQuota, 10)
+	}
+	if rl.CPUSetCPUs != "" {
+		attrs[attrCPUSetCPUs] = rl.CPUSetCPUs
+	}
+	if rl.CPUSetMems != "" {
+		attrs[attrCPUSetMems] = rl.CPUSetMems
+	}
+}
+
 func notSupported(f driver.Feature, d *driver.DriverHandle, docs string) error {
 	return errors.Errorf(`%s is not supported for the %s driver.
 Switch to a different driver, or turn on the containerd image store, and try again.
@@ -5,6 +5,7 @@ import (
 	"strings"
 	"testing"

+	dockeropts "github.com/docker/cli/opts"
 	"github.com/stretchr/testify/require"
 )

@@ -150,3 +151,109 @@ func TestToBuildkitExtraHosts(t *testing.T) {
 		})
 	}
 }
+
+func TestAddResourceLimits(t *testing.T) {
+	mustMemSwap := func(v string) dockeropts.MemSwapBytes {
+		var m dockeropts.MemSwapBytes
+		require.NoError(t, m.Set(v))
+		return m
+	}
+
+	tests := []struct {
+		name     string
+		limits   ResourceLimits
+		expected map[string]string
+	}{
+		{
+			name:     "empty",
+			limits:   ResourceLimits{},
+			expected: map[string]string{},
+		},
+		{
+			name: "all",
+			limits: ResourceLimits{
+				Memory:     dockeropts.MemBytes(2 * 1024 * 1024 * 1024),
+				MemorySwap: mustMemSwap("4g"),
+				CPUShares:  1024,
+				CPUPeriod:  100000,
+				CPUQuota:   50000,
+				CPUSetCPUs: "0-3",
+				CPUSetMems: "0,1",
+			},
+			expected: map[string]string{
+				"memory":     "2147483648",
+				"memswap":    "4294967296",
+				"cpushares":  "1024",
+				"cpuperiod":  "100000",
+				"cpuquota":   "50000",
+				"cpusetcpus": "0-3",
+				"cpusetmems": "0,1",
+			},
+		},
+		{
+			name:   "unlimited swap",
+			limits: ResourceLimits{MemorySwap: mustMemSwap("-1")},
+			expected: map[string]string{
+				"memswap": "-1",
+			},
+		},
+		{
+			name:   "partial",
+			limits: ResourceLimits{Memory: dockeropts.MemBytes(512 * 1024 * 1024)},
+			expected: map[string]string{
+				"memory": "536870912",
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			attrs := map[string]string{}
+			addResourceLimits(tc.limits, attrs)
+			require.Equal(t, tc.expected, attrs)
+		})
+	}
+}
+
+func TestParseResourceLimits(t *testing.T) {
+	t.Run("all", func(t *testing.T) {
+		rl, err := ParseResourceLimits([]string{
+			"memory=2g",
+			"memory-swap=4g",
+			"cpu-shares=1024",
+			"cpu-period=100000",
+			"cpu-quota=50000",
+			"cpuset-cpus=0-3",
+			"cpuset-mems=0,1",
+		})
+		require.NoError(t, err)
+		require.Equal(t, int64(2*1024*1024*1024), rl.Memory.Value())
+		require.Equal(t, int64(4*1024*1024*1024), rl.MemorySwap.Value())
+		require.Equal(t, int64(1024), rl.CPUShares)
+		require.Equal(t, int64(100000), rl.CPUPeriod)
+		require.Equal(t, int64(50000), rl.CPUQuota)
+		require.Equal(t, "0-3", rl.CPUSetCPUs)
+		require.Equal(t, "0,1", rl.CPUSetMems)
+	})
+
+	t.Run("unlimited swap", func(t *testing.T) {
+		rl, err := ParseResourceLimits([]string{"memory-swap=-1"})
+		require.NoError(t, err)
+		require.Equal(t, int64(-1), rl.MemorySwap.Value())
+	})
+
+	t.Run("missing value", func(t *testing.T) {
+		_, err := ParseResourceLimits([]string{"memory"})
+		require.Error(t, err)
+	})
+
+	t.Run("unknown key", func(t *testing.T) {
+		_, err := ParseResourceLimits([]string{"bogus=1"})
+		require.Error(t, err)
+	})
+
+	t.Run("invalid int", func(t *testing.T) {
+		_, err := ParseResourceLimits([]string{"cpu-shares=notanumber"})
+		require.Error(t, err)
+	})
+}