From 227156dac6d9204dbb718d6ff85527a06846aabb Mon Sep 17 00:00:00 2001
From: Justin Terry <jlterry@amazon.com>
Date: Thu, 14 Apr 2022 15:21:21 -0700
Subject: [PATCH] Add ctr support for CPUMax and CPUShares

Adds CPU.Maximum and CPU.Shares support to the ctr
cmdline for testing

Signed-off-by: Justin Terry <jlterry@amazon.com>
---
 cmd/ctr/commands/commands_windows.go | 23 ++++++++++++++++-------
 cmd/ctr/commands/run/run_windows.go  |  8 ++++++++
 oci/spec_opts.go                     |  9 +++++++++
 oci/spec_opts_windows.go             | 27 +++++++++++++++++++++------
 4 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/cmd/ctr/commands/commands_windows.go b/cmd/ctr/commands/commands_windows.go
index 10a5d6a5d..179cce501 100644
--- a/cmd/ctr/commands/commands_windows.go
+++ b/cmd/ctr/commands/commands_windows.go
@@ -21,11 +21,20 @@ import (
 )
 
 func init() {
-	ContainerFlags = append(ContainerFlags, cli.Uint64Flag{
-		Name:  "cpu-count",
-		Usage: "number of CPUs available to the container",
-	}, cli.StringSliceFlag{
-		Name:  "device",
-		Usage: "identifier of a device to add to the container  (e.g. class://5B45201D-F2F2-4F3B-85BB-30FF1F953599)",
-	})
+	ContainerFlags = append(ContainerFlags,
+		cli.Uint64Flag{
+			Name:  "cpu-count",
+			Usage: "number of CPUs available to the container",
+		},
+		cli.Uint64Flag{
+			Name:  "cpu-shares",
+			Usage: "The relative number of CPU shares given to the container relative to other workloads. Between 0 and 10,000.",
+		},
+		cli.Uint64Flag{
+			Name:  "cpu-max",
+			Usage: "The number of processor cycles threads in a container can use per 10,000 cycles. Set to a percentage times 100. Between 1 and 10,000",
+		}, cli.StringSliceFlag{
+			Name:  "device",
+			Usage: "identifier of a device to add to the container  (e.g. class://5B45201D-F2F2-4F3B-85BB-30FF1F953599)",
+		})
 }
diff --git a/cmd/ctr/commands/run/run_windows.go b/cmd/ctr/commands/run/run_windows.go
index 60efca42f..52af73a0f 100644
--- a/cmd/ctr/commands/run/run_windows.go
+++ b/cmd/ctr/commands/run/run_windows.go
@@ -143,6 +143,14 @@ func NewContainer(ctx gocontext.Context, client *containerd.Client, context *cli
 		if ccount != 0 {
 			opts = append(opts, oci.WithWindowsCPUCount(ccount))
 		}
+		cshares := context.Uint64("cpu-shares")
+		if cshares != 0 {
+			opts = append(opts, oci.WithWindowsCPUShares(uint16(cshares)))
+		}
+		cmax := context.Uint64("cpu-max")
+		if cmax != 0 {
+			opts = append(opts, oci.WithWindowsCPUMaximum(uint16(cmax)))
+		}
 		for _, dev := range context.StringSlice("device") {
 			parts := strings.Split(dev, "://")
 			if len(parts) != 2 {
diff --git a/oci/spec_opts.go b/oci/spec_opts.go
index 9911e6299..371d2ae44 100644
--- a/oci/spec_opts.go
+++ b/oci/spec_opts.go
@@ -83,6 +83,10 @@ func setResources(s *Spec) {
 			s.Linux.Resources = &specs.LinuxResources{}
 		}
 	}
+}
+
+// nolint
+func setResourcesWindows(s *Spec) {
 	if s.Windows != nil {
 		if s.Windows.Resources == nil {
 			s.Windows.Resources = &specs.WindowsResources{}
@@ -98,6 +102,11 @@ func setCPU(s *Spec) {
 			s.Linux.Resources.CPU = &specs.LinuxCPU{}
 		}
 	}
+}
+
+// nolint
+func setCPUWindows(s *Spec) {
+	setResourcesWindows(s)
 	if s.Windows != nil {
 		if s.Windows.Resources.CPU == nil {
 			s.Windows.Resources.CPU = &specs.WindowsCPUResources{}
diff --git a/oci/spec_opts_windows.go b/oci/spec_opts_windows.go
index 4ddb13d3f..94cb39312 100644
--- a/oci/spec_opts_windows.go
+++ b/oci/spec_opts_windows.go
@@ -31,17 +31,32 @@ import (
 // `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{}
-		}
+		setCPUWindows(s)
 		s.Windows.Resources.CPU.Count = &count
 		return nil
 	}
 }
 
+// WithWindowsCPUShares sets the `Windows.Resources.CPU.Shares` section to the
+// `shares` specified.
+func WithWindowsCPUShares(shares uint16) SpecOpts {
+	return func(_ context.Context, _ Client, _ *containers.Container, s *Spec) error {
+		setCPUWindows(s)
+		s.Windows.Resources.CPU.Shares = &shares
+		return nil
+	}
+}
+
+// WithWindowsCPUMaximum sets the `Windows.Resources.CPU.Maximum` section to the
+// `max` specified.
+func WithWindowsCPUMaximum(max uint16) SpecOpts {
+	return func(_ context.Context, _ Client, _ *containers.Container, s *Spec) error {
+		setCPUWindows(s)
+		s.Windows.Resources.CPU.Maximum = &max
+		return nil
+	}
+}
+
 // WithWindowsIgnoreFlushesDuringBoot sets `Windows.IgnoreFlushesDuringBoot`.
 func WithWindowsIgnoreFlushesDuringBoot() SpecOpts {
 	return func(_ context.Context, _ Client, _ *containers.Container, s *Spec) error {