From a3ca8a0dfc4f2e79487d2479fbb3ee41bc75227b Mon Sep 17 00:00:00 2001
From: Michael Crosby <crosbymichael@gmail.com>
Date: Thu, 5 Mar 2020 13:02:10 -0500
Subject: [PATCH] Add linux resource oci.SpecOpts

Signed-off-by: Michael Crosby <crosbymichael@gmail.com>
---
 oci/spec_opts.go       | 15 +++++++++++
 oci/spec_opts_linux.go | 61 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/oci/spec_opts.go b/oci/spec_opts.go
index 4c14c5570..59dbfdb19 100644
--- a/oci/spec_opts.go
+++ b/oci/spec_opts.go
@@ -91,6 +91,21 @@ func setResources(s *Spec) {
 	}
 }
 
+// nolint
+func setCPU(s *Spec) {
+	setResources(s)
+	if s.Linux != nil {
+		if s.Linux.Resources.CPU == nil {
+			s.Linux.Resources.CPU = &specs.LinuxCPU{}
+		}
+	}
+	if s.Windows != nil {
+		if s.Windows.Resources.CPU == nil {
+			s.Windows.Resources.CPU = &specs.WindowsCPUResources{}
+		}
+	}
+}
+
 // setCapabilities sets Linux Capabilities to empty if unset
 func setCapabilities(s *Spec) {
 	setProcess(s)
diff --git a/oci/spec_opts_linux.go b/oci/spec_opts_linux.go
index 1448ee78f..b1eab6f9f 100644
--- a/oci/spec_opts_linux.go
+++ b/oci/spec_opts_linux.go
@@ -119,3 +119,64 @@ func deviceFromPath(path, permissions string) (*specs.LinuxDevice, error) {
 		GID:      &stat.Gid,
 	}, nil
 }
+
+// WithMemorySwap sets the container's swap in bytes
+func WithMemorySwap(swap int64) SpecOpts {
+	return func(ctx context.Context, _ Client, c *containers.Container, s *Spec) error {
+		setResources(s)
+		if s.Linux.Resources.Memory == nil {
+			s.Linux.Resources.Memory = &specs.LinuxMemory{}
+		}
+		s.Linux.Resources.Memory.Swap = &swap
+		return nil
+	}
+}
+
+// WithPidsLimit sets the container's pid limit or maximum
+func WithPidsLimit(limit int64) SpecOpts {
+	return func(ctx context.Context, _ Client, c *containers.Container, s *Spec) error {
+		setResources(s)
+		if s.Linux.Resources.Pids == nil {
+			s.Linux.Resources.Pids = &specs.LinuxPids{}
+		}
+		s.Linux.Resources.Pids.Limit = limit
+		return nil
+	}
+}
+
+// WithCPUShares sets the container's cpu shares
+func WithCPUShares(shares uint64) SpecOpts {
+	return func(ctx context.Context, _ Client, c *containers.Container, s *Spec) error {
+		setCPU(s)
+		s.Linux.Resources.CPU.Shares = &shares
+		return nil
+	}
+}
+
+// WithCPUs sets the container's cpus/cores for use by the container
+func WithCPUs(cpus string) SpecOpts {
+	return func(ctx context.Context, _ Client, c *containers.Container, s *Spec) error {
+		setCPU(s)
+		s.Linux.Resources.CPU.Cpus = cpus
+		return nil
+	}
+}
+
+// WithCPUsMems sets the container's cpu mems for use by the container
+func WithCPUsMems(mems string) SpecOpts {
+	return func(ctx context.Context, _ Client, c *containers.Container, s *Spec) error {
+		setCPU(s)
+		s.Linux.Resources.CPU.Mems = mems
+		return nil
+	}
+}
+
+// WithCPUCFS sets the container's Completely fair scheduling (CFS) quota and period
+func WithCPUCFS(quota int64, period uint64) SpecOpts {
+	return func(ctx context.Context, _ Client, c *containers.Container, s *Spec) error {
+		setCPU(s)
+		s.Linux.Resources.CPU.Quota = &quota
+		s.Linux.Resources.CPU.Period = &period
+		return nil
+	}
+}