Adding ctr memory and cpu flags

Adds ctr run --memory-limit for all platforms.
Adds ctr run --cpu-count for Windows platforms.

Signed-off-by: Justin Terry (VM) <juterry@microsoft.com>

cmd/ctr/commands/commands.go

@@ -124,6 +124,10 @@ var (
 			Name:  "allow-new-privs",
 			Usage: "turn off OCI spec's NoNewPrivileges feature flag",
 		},
+		cli.Uint64Flag{
+			Name:  "memory-limit",
+			Usage: "memory limit (in bytes) for the container",
+		},
 	}
 )
 

cmd/ctr/commands/commands_windows.go

@@ -0,0 +1,30 @@
+// +build windows
+
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package commands
+
+import (
+	"github.com/urfave/cli"
+)
+
+func init() {
+	ContainerFlags = append(ContainerFlags, cli.Uint64Flag{
+		Name:  "cpu-count",
+		Usage: "number of CPUs available to the container",
+	})
+}

cmd/ctr/commands/run/run_unix.go

@@ -139,6 +139,10 @@ func NewContainer(ctx gocontext.Context, client *containerd.Client, context *cli
 			// NOTE: can be set to "" explicitly for disabling cgroup.
 			opts = append(opts, oci.WithCgroup(context.String("cgroup")))
 		}
+		limit := context.Uint64("memory-limit")
+		if limit != 0 {
+			opts = append(opts, oci.WithMemoryLimit(limit))
+		}
 	}
 
 	cOpts = append(cOpts, containerd.WithRuntime(context.String("runtime"), nil))

cmd/ctr/commands/run/run_windows.go

@@ -105,6 +105,14 @@ func NewContainer(ctx gocontext.Context, client *containerd.Client, context *cli
 		if context.Bool("isolated") {
 			opts = append(opts, oci.WithWindowsHyperV)
 		}
+		limit := context.Uint64("memory-limit")
+		if limit != 0 {
+			opts = append(opts, oci.WithMemoryLimit(limit))
+		}
+		ccount := context.Uint64("cpu-count")
+		if ccount != 0 {
+			opts = append(opts, oci.WithWindowsCPUCount(ccount))
+		}
 	}
 
 	cOpts = append(cOpts, containerd.WithContainerLabels(commands.LabelArgs(context.StringSlice("label"))))

oci/spec_opts.go

@@ -1026,3 +1026,32 @@ func WithWindowsHyperV(_ context.Context, _ Client, _ *containers.Container, s *
 	}
 	return nil
 }
+
+// WithMemoryLimit sets the `Linux.LinuxResources.Memory.Limit` section to the
+// `limit` specified if the `Linux` section is not `nil`. Additionally sets the
+// `Windows.WindowsResources.Memory.Limit` section if the `Windows` section is
+// not `nil`.
+func WithMemoryLimit(limit uint64) SpecOpts {
+	return func(_ context.Context, _ Client, _ *containers.Container, s *Spec) error {
+		if s.Linux != nil {
+			if s.Linux.Resources == nil {
+				s.Linux.Resources = &specs.LinuxResources{}
+			}
+			if s.Linux.Resources.Memory == nil {
+				s.Linux.Resources.Memory = &specs.LinuxMemory{}
+			}
+			l := int64(limit)
+			s.Linux.Resources.Memory.Limit = &l
+		}
+		if s.Windows != nil {
+			if s.Windows.Resources == nil {
+				s.Windows.Resources = &specs.WindowsResources{}
+			}
+			if s.Windows.Resources.Memory == nil {
+				s.Windows.Resources.Memory = &specs.WindowsMemoryResources{}
+			}
+			s.Windows.Resources.Memory.Limit = &limit
+		}
+		return nil
+	}
+}

oci/spec_opts_test.go

@@ -174,3 +174,61 @@ func TestWithSpecFromFile(t *testing.T) {
 		t.Fatalf("spec from option differs from default: \n%#v != \n%#v", &s, expected)
 	}
 }
+
+func TestWithMemoryLimit(t *testing.T) {
+	var (
+		ctx = namespaces.WithNamespace(context.Background(), "testing")
+		c   = containers.Container{ID: t.Name()}
+		m   = uint64(768 * 1024 * 1024)
+		o   = WithMemoryLimit(m)
+	)
+	// Test with all three supported scenarios
+	platforms := []string{"", "linux/amd64", "windows/amd64"}
+	for _, p := range platforms {
+		var spec *Spec
+		var err error
+		if p == "" {
+			t.Log("Testing GenerateSpec default platform")
+			spec, err = GenerateSpec(ctx, nil, &c, o)
+
+			// Convert the platform to the default based on GOOS like
+			// GenerateSpec does.
+			switch runtime.GOOS {
+			case "linux":
+				p = "linux/amd64"
+			case "windows":
+				p = "windows/amd64"
+			}
+		} else {
+			t.Logf("Testing GenerateSpecWithPlatform with platform: '%s'", p)
+			spec, err = GenerateSpecWithPlatform(ctx, nil, p, &c, o)
+		}
+		if err != nil {
+			t.Fatalf("failed to generate spec with: %v", err)
+		}
+		switch p {
+		case "linux/amd64":
+			if *spec.Linux.Resources.Memory.Limit != int64(m) {
+				t.Fatalf("spec.Linux.Resources.Memory.Limit expected: %v, got: %v", m, *spec.Linux.Resources.Memory.Limit)
+			}
+			// If we are linux/amd64 on Windows GOOS it is LCOW
+			if runtime.GOOS == "windows" {
+				// Verify that we also set the Windows section.
+				if *spec.Windows.Resources.Memory.Limit != m {
+					t.Fatalf("for LCOW spec.Windows.Resources.Memory.Limit is also expected: %v, got: %v", m, *spec.Windows.Resources.Memory.Limit)
+				}
+			} else {
+				if spec.Windows != nil {
+					t.Fatalf("spec.Windows section should not be set for linux/amd64 spec on non-windows platform")
+				}
+			}
+		case "windows/amd64":
+			if *spec.Windows.Resources.Memory.Limit != m {
+				t.Fatalf("spec.Windows.Resources.Memory.Limit expected: %v, got: %v", m, *spec.Windows.Resources.Memory.Limit)
+			}
+			if spec.Linux != nil {
+				t.Fatalf("spec.Linux section should not be set for windows/amd64 spec ever")
+			}
+		}
+	}
+}

oci/spec_opts_windows.go

@@ -0,0 +1,41 @@
+// +build windows
+
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package oci
+
+import (
+	"context"
+
+	"github.com/containerd/containerd/containers"
+	specs "github.com/opencontainers/runtime-spec/specs-go"
+)
+
+// WithWindowsCPUCount sets the `Windows.Resources.CPU.Count` section to the
+// `count` specified.
+func WithWindowsCPUCount(count uint64) SpecOpts {
+	return func(_ context.Context, _ Client, _ *containers.Container, s *Spec) error {
+		if s.Windows.Resources == nil {
+			s.Windows.Resources = &specs.WindowsResources{}
+		}
+		if s.Windows.Resources.CPU == nil {
+			s.Windows.Resources.CPU = &specs.WindowsCPUResources{}
+		}
+		s.Windows.Resources.CPU.Count = &count
+		return nil
+	}
+}

oci/spec_opts_windows_test.go

@@ -0,0 +1,58 @@
+// +build windows
+
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package oci
+
+import (
+	"context"
+	"testing"
+
+	"github.com/containerd/containerd/containers"
+	"github.com/containerd/containerd/namespaces"
+)
+
+func TestWithCPUCount(t *testing.T) {
+	var (
+		ctx = namespaces.WithNamespace(context.Background(), "testing")
+		c   = containers.Container{ID: t.Name()}
+		cpu = uint64(8)
+		o   = WithWindowsCPUCount(cpu)
+	)
+	// Test with all three supported scenarios
+	platforms := []string{"", "linux/amd64", "windows/amd64"}
+	for _, p := range platforms {
+		var spec *Spec
+		var err error
+		if p == "" {
+			t.Log("Testing GenerateSpec default platform")
+			spec, err = GenerateSpec(ctx, nil, &c, o)
+		} else {
+			t.Logf("Testing GenerateSpecWithPlatform with platform: '%s'", p)
+			spec, err = GenerateSpecWithPlatform(ctx, nil, p, &c, o)
+		}
+		if err != nil {
+			t.Fatalf("failed to generate spec with: %v", err)
+		}
+		if *spec.Windows.Resources.CPU.Count != cpu {
+			t.Fatalf("spec.Windows.Resources.CPU.Count expected: %v, got: %v", cpu, *spec.Windows.Resources.CPU.Count)
+		}
+		if spec.Linux != nil && spec.Linux.Resources != nil && spec.Linux.Resources.CPU != nil {
+			t.Fatalf("spec.Linux.Resources.CPU section should not be set on GOOS=windows")
+		}
+	}
+}