runc exec: implement CPU affinity

As per
- opencontainers/runtime-spec#1253
- opencontainers/runtime-spec#1261

Add some tests (alas it's impossible to test initial CPU affinity
without adding debug logging).

Signed-off-by: Kir Kolyshkin <kolyshkin@gmail.com>

Author: kolyshkin

libcontainer/configs/config.go

@@ -225,6 +225,9 @@ type Config struct {
 
 	// IOPriority is the container's I/O priority.
 	IOPriority *IOPriority `json:"io_priority,omitempty"`
+
+	// ExecCPUAffinity is CPU affinity for a non-init process to be run in the container.
+	ExecCPUAffinity *CPUAffinity `json:"exec_cpu_affinity,omitempty"`
 }
 
 // Scheduler is based on the Linux sched_setattr(2) syscall.
@@ -286,7 +289,10 @@ func ToSchedAttr(scheduler *Scheduler) (*unix.SchedAttr, error) {
 	}, nil
 }
 
-type IOPriority = specs.LinuxIOPriority
+type (
+	IOPriority  = specs.LinuxIOPriority
+	CPUAffinity = specs.CPUAffinity
+)
 
 type (
 	HookName string

libcontainer/configs/config_linux.go

@@ -4,6 +4,10 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"strconv"
+	"strings"
+
+	"golang.org/x/sys/unix"
 )
 
 var (
@@ -95,3 +99,40 @@ func (c Config) hostIDFromMapping(containerID int64, uMap []IDMap) (int64, bool)
 	}
 	return -1, false
 }
+
+// ToCPUSet converts [CPUAffinity] (initial or final) value to [unix.CPUSet].
+func ToCPUSet(str string) (*unix.CPUSet, error) {
+	s := new(unix.CPUSet)
+	for _, r := range strings.Split(str, ",") {
+		// Allow extra spaces around.
+		r = strings.TrimSpace(r)
+		// Allow empty elements (extra commas).
+		if r == "" {
+			continue
+		}
+		if r0, r1, found := strings.Cut(r, "-"); found {
+			start, err := strconv.ParseUint(r0, 10, 32)
+			if err != nil {
+				return nil, err
+			}
+			end, err := strconv.ParseUint(r1, 10, 32)
+			if err != nil {
+				return nil, err
+			}
+			if start > end {
+				return nil, errors.New("invalid range: " + r)
+			}
+			for i := int(start); i <= int(end); i++ {
+				s.Set(i)
+			}
+		} else {
+			val, err := strconv.ParseUint(r, 10, 32)
+			if err != nil {
+				return nil, err
+			}
+			s.Set(int(val))
+		}
+	}
+
+	return s, nil
+}

libcontainer/container_linux.go

@@ -697,6 +697,7 @@ func (c *Container) newInitConfig(process *Process) *initConfig {
 		AppArmorProfile:  c.config.AppArmorProfile,
 		ProcessLabel:     c.config.ProcessLabel,
 		Rlimits:          c.config.Rlimits,
+		CPUAffinity:      c.config.ExecCPUAffinity,
 		CreateConsole:    process.ConsoleSocket != nil,
 		ConsoleWidth:     process.ConsoleWidth,
 		ConsoleHeight:    process.ConsoleHeight,
@@ -713,6 +714,9 @@ func (c *Container) newInitConfig(process *Process) *initConfig {
 	if len(process.Rlimits) > 0 {
 		cfg.Rlimits = process.Rlimits
 	}
+	if process.CPUAffinity != nil {
+		cfg.CPUAffinity = process.CPUAffinity
+	}
 	if cgroups.IsCgroup2UnifiedMode() {
 		cfg.Cgroup2Path = c.cgroupManager.Path("")
 	}

libcontainer/init_linux.go

@@ -72,6 +72,7 @@ type initConfig struct {
 	RootlessCgroups  bool                  `json:"rootless_cgroups,omitempty"`
 	SpecState        *specs.State          `json:"spec_state,omitempty"`
 	Cgroup2Path      string                `json:"cgroup2_path,omitempty"`
+	CPUAffinity      *specs.CPUAffinity    `json:"cpu_affinity,omitempty"`
 }
 
 // Init is part of "runc init" implementation.
@@ -199,6 +200,23 @@ func startInitialization() (retErr error) {
 		}
 	}()
 
+	// See tests/integration/cpu_affinity.bats.
+	if logrus.GetLevel() >= logrus.DebugLevel {
+		var cpus unix.CPUSet
+		err := unix.SchedGetaffinity(0, &cpus)
+		if err != nil {
+			logrus.Debugf("sched_getaffinity: error %v", err)
+		} else {
+			var list []int
+			for i := 0; i < 256; i++ {
+				if cpus.IsSet(i) {
+					list = append(list, i)
+				}
+			}
+			logrus.Debugf("Initial CPUs: %v", list)
+		}
+	}
+
 	var config initConfig
 	if err := json.NewDecoder(initPipe).Decode(&config); err != nil {
 		return err

libcontainer/process.go

@@ -102,6 +102,8 @@ type Process struct {
 	Scheduler *configs.Scheduler
 
 	IOPriority *configs.IOPriority
+
+	CPUAffinity *configs.CPUAffinity
 }
 
 // Wait waits for the process to exit.

libcontainer/process_linux.go

@@ -163,13 +163,62 @@ type setnsProcess struct {
 	initProcessPid  int
 }
 
+// Starts setns process with specified initial CPU affinity.
+func (p *setnsProcess) startWithCPUAffinity() error {
+	aff := p.config.CPUAffinity
+	if aff == nil || aff.Initial == "" {
+		return p.cmd.Start()
+	}
+	cpus, err := configs.ToCPUSet(aff.Initial)
+	if err != nil {
+		return fmt.Errorf("invalid CPUAffinity.initial: %w", err)
+	}
+
+	errCh := make(chan error)
+	defer close(errCh)
+
+	// Use a goroutine to dedicate an OS thread.
+	go func() {
+		runtime.LockOSThread()
+		// Command inherits the CPU affinity.
+		if err := unix.SchedSetaffinity(unix.Gettid(), cpus); err != nil {
+			runtime.UnlockOSThread()
+			errCh <- fmt.Errorf("setting initial CPU affinity: %w", err)
+			return
+		}
+
+		errCh <- p.cmd.Start()
+		// Deliberately omit runtime.UnlockOSThread here.
+		// https://pkg.go.dev/runtime#LockOSThread says:
+		// "If the calling goroutine exits without unlocking the
+		// thread, the thread will be terminated".
+	}()
+
+	return <-errCh
+}
+
+func (p *setnsProcess) setFinalCPUAffinity() error {
+	aff := p.config.CPUAffinity
+	if aff == nil || aff.Final == "" {
+		return nil
+	}
+	cpus, err := configs.ToCPUSet(aff.Final)
+	if err != nil {
+		return fmt.Errorf("invalid CPUAffinity.final: %w", err)
+	}
+	if err := unix.SchedSetaffinity(p.pid(), cpus); err != nil {
+		return fmt.Errorf("setting final CPU affinity: %w", err)
+	}
+	return nil
+}
+
 func (p *setnsProcess) start() (retErr error) {
 	defer p.comm.closeParent()
 
-	// get the "before" value of oom kill count
+	// Get the "before" value of oom kill count.
 	oom, _ := p.manager.OOMKillCount()
-	err := p.cmd.Start()
-	// close the child-side of the pipes (controlled by child)
+	err := p.startWithCPUAffinity()
+	// Close the child-side of the pipes (controlled by child).
 	p.comm.closeChild()
 	if err != nil {
 		return fmt.Errorf("error starting setns process: %w", err)
@@ -228,6 +277,9 @@ func (p *setnsProcess) start() (retErr error) {
 			}
 		}
 	}
+	if err := p.setFinalCPUAffinity(); err != nil {
+		return err
+	}
 
 	if err := utils.WriteJSON(p.comm.initSockParent, p.config); err != nil {
 		return fmt.Errorf("error writing config to pipe: %w", err)

libcontainer/specconv/spec_linux.go

@@ -556,6 +556,11 @@ func CreateLibcontainerConfig(opts *CreateOpts) (*configs.Config, error) {
 			ioPriority := *spec.Process.IOPriority
 			config.IOPriority = &ioPriority
 		}
+		if spec.Process.ExecCPUAffinity != nil {
+			a := *spec.Process.ExecCPUAffinity
+			config.ExecCPUAffinity = &a
+		}
+
 	}
 	createHooks(spec, config)
 	config.Version = specs.Version

tests/integration/cpu_affinity.bats

@@ -0,0 +1,85 @@
+#!/usr/bin/env bats
+# Exec CPU affinity tests. For more details, see:
+#  - https://github.com/opencontainers/runtime-spec/pull/1253
+
+load helpers
+
+function setup() {
+	requires smp cgroups_cpuset
+	setup_busybox
+}
+
+function teardown() {
+	teardown_bundle
+}
+
+function all_cpus() {
+	cat /sys/devices/system/cpu/online
+}
+
+function first_cpu() {
+	all_cpus | sed 's/[-,].*//g'
+}
+
+@test "runc exec [CPU affinity inherited from runc]" {
+	requires root # For taskset.
+
+	first="$(first_cpu)"
+
+	# Container's process CPU affinity is inherited from that of runc.
+	taskset -p -c "$first" $$
+
+	runc run -d --console-socket "$CONSOLE_SOCKET" ct1
+	[ "$status" -eq 0 ]
+
+	# Check init.
+	runc exec ct1 grep "Cpus_allowed_list:" /proc/1/status
+	[ "$status" -eq 0 ]
+	[[ "${lines[0]}" == "Cpus_allowed_list:	$first" ]]
+
+	# Check exec.
+	runc exec ct1 grep "Cpus_allowed_list:" /proc/self/status
+	[ "$status" -eq 0 ]
+	[[ "${lines[0]}" == "Cpus_allowed_list:	$first" ]]
+}
+
+@test "runc exec [CPU affinity, initial set via process.json]" {
+	first="$(first_cpu)"
+	second=$((first+1)) # Hacky; might not work in all environments.
+
+	runc run -d --console-socket "$CONSOLE_SOCKET" ct1
+	[ "$status" -eq 0 ]
+
+	for cpus in "$first" "$first-$second" "$first,$second" "$second"; do
+		proc='
+{
+    "terminal": false,
+    "execCPUAffinity": {
+	    "initial": "'$cpus'"
+    },
+    "args": [ "/bin/true" ],
+    "cwd": "/"
+}'
+		exp=${cpus//,/-} # 1. "," --> "-".
+		exp=${exp//-/ }  # 2. "-" --> " ".
+		echo "CPUS: $cpus, exp: $exp"
+		runc --debug exec --process <(echo "$proc") ct1
+		[[ "$output" == *"Initial CPUs: [$exp]"* ]]
+	done
+}
+
+@test "runc exec [CPU affinity, initial and final are set]" {
+	first="$(first_cpu)"
+	second=$((first+1)) # Hacky; might not work in all environments.
+
+	update_config "	  .process.execCPUAffinity.initial = \"$first\"
+			| .process.execCPUAffinity.final = \"$second\""
+
+	runc run -d --console-socket "$CONSOLE_SOCKET" ct1
+	[ "$status" -eq 0 ]
+
+	runc --debug exec ct1 grep "Cpus_allowed_list:" /proc/self/status
+	[ "$status" -eq 0 ]
+	[[ "$output" == *"Initial CPUs: [$first]"* ]]
+	[[ "$output}" == *"Cpus_allowed_list:	$second"* ]] # Mind the literal tab.
+}

utils_linux.go

@@ -57,6 +57,7 @@ func newProcess(p specs.Process) (*libcontainer.Process, error) {
 		AppArmorProfile: p.ApparmorProfile,
 		Scheduler:       p.Scheduler,
 		IOPriority:      p.IOPriority,
+		CPUAffinity:     p.ExecCPUAffinity,
 	}
 
 	if p.ConsoleSize != nil {