jobcontainers: wire Windows CPU affinity to JobObject limits (microsoft#2603)

Signed-off-by: Weyang1 <weyang1@microsoft.com>
Co-authored-by: Weyang1 <weyang1@microsoft.com>

Author: wenyang

internal/jobcontainers/oci.go

@@ -4,6 +4,7 @@ package jobcontainers
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/Microsoft/hcsshim/internal/hcsoci"
 	"github.com/Microsoft/hcsshim/internal/jobobject"
@@ -40,6 +41,21 @@ func specToLimits(ctx context.Context, cid string, s *specs.Spec) (*jobobject.Jo
 		return nil, err
 	}
 
+	var cpuAffinity uint64
+	if s.Windows != nil && s.Windows.Resources != nil && s.Windows.Resources.CPU != nil && len(s.Windows.Resources.CPU.Affinity) > 0 {
+		affinity := s.Windows.Resources.CPU.Affinity
+		if len(affinity) != 1 {
+			return nil, fmt.Errorf("cpu affinity with multiple processor groups is not supported")
+		}
+		if affinity[0].Group != 0 {
+			return nil, fmt.Errorf("cpu affinity processor group %d is not supported", affinity[0].Group)
+		}
+		if affinity[0].Mask == 0 {
+			return nil, fmt.Errorf("cpu affinity mask must be non-zero")
+		}
+		cpuAffinity = affinity[0].Mask
+	}
+
 	realCPULimit, realCPUWeight := uint32(cpuLimit), uint32(cpuWeight)
 	if cpuCount != 0 {
 		// Job object API does not support "CPU count". Instead, we translate the notion of "count" into
@@ -61,6 +77,7 @@ func specToLimits(ctx context.Context, cid string, s *specs.Spec) (*jobobject.Jo
 	return &jobobject.JobLimits{
 		CPULimit:           realCPULimit,
 		CPUWeight:          realCPUWeight,
+		CPUAffinity:        cpuAffinity,
 		MaxIOPS:            maxIops,
 		MaxBandwidth:       maxBandwidth,
 		MemoryLimitInBytes: memLimitMB * memory.MiB,

internal/jobcontainers/oci_test.go

@@ -0,0 +1,100 @@
+//go:build windows
+
+package jobcontainers
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	specs "github.com/opencontainers/runtime-spec/specs-go"
+)
+
+func TestSpecToLimits_CPUAffinity_Group0MaskSet(t *testing.T) {
+	s := &specs.Spec{
+		Windows: &specs.Windows{
+			Resources: &specs.WindowsResources{
+				CPU: &specs.WindowsCPUResources{
+					Affinity: []specs.WindowsCPUGroupAffinity{
+						{Mask: 0x3, Group: 0},
+					},
+				},
+			},
+		},
+	}
+
+	limits, err := specToLimits(context.Background(), "cid", s)
+	if err != nil {
+		t.Fatalf("specToLimits failed: %v", err)
+	}
+	if limits.CPUAffinity != 0x3 {
+		t.Fatalf("unexpected cpu affinity: got %d want %d", limits.CPUAffinity, uint64(0x3))
+	}
+}
+
+func TestSpecToLimits_CPUAffinity_MultiGroupRejected(t *testing.T) {
+	s := &specs.Spec{
+		Windows: &specs.Windows{
+			Resources: &specs.WindowsResources{
+				CPU: &specs.WindowsCPUResources{
+					Affinity: []specs.WindowsCPUGroupAffinity{
+						{Mask: 0x1, Group: 0},
+						{Mask: 0x1, Group: 1},
+					},
+				},
+			},
+		},
+	}
+
+	_, err := specToLimits(context.Background(), "cid", s)
+	if err == nil {
+		t.Fatal("expected error for multiple affinity entries")
+	}
+	if !strings.Contains(err.Error(), "multiple processor groups") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}
+
+func TestSpecToLimits_CPUAffinity_NonZeroGroupRejected(t *testing.T) {
+	s := &specs.Spec{
+		Windows: &specs.Windows{
+			Resources: &specs.WindowsResources{
+				CPU: &specs.WindowsCPUResources{
+					Affinity: []specs.WindowsCPUGroupAffinity{
+						{Mask: 0x1, Group: 1},
+					},
+				},
+			},
+		},
+	}
+
+	_, err := specToLimits(context.Background(), "cid", s)
+	if err == nil {
+		t.Fatal("expected error for non-zero affinity group")
+	}
+	if !strings.Contains(err.Error(), "processor group") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}
+
+func TestSpecToLimits_CPUAffinity_ZeroMaskRejected(t *testing.T) {
+	s := &specs.Spec{
+		Windows: &specs.Windows{
+			Resources: &specs.WindowsResources{
+				CPU: &specs.WindowsCPUResources{
+					Affinity: []specs.WindowsCPUGroupAffinity{
+						{Mask: 0, Group: 0},
+					},
+				},
+			},
+		},
+	}
+
+	_, err := specToLimits(context.Background(), "cid", s)
+	if err == nil {
+		t.Fatal("expected error for zero affinity mask")
+	}
+	if !strings.Contains(err.Error(), "mask must be non-zero") {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}

internal/jobobject/jobobject.go

@@ -32,6 +32,7 @@ type JobObject struct {
 type JobLimits struct {
 	CPULimit           uint32
 	CPUWeight          uint32
+	CPUAffinity        uint64
 	MemoryLimitInBytes uint64
 	MaxIOPS            int64
 	MaxBandwidth       int64

internal/jobobject/jobobject_test.go

@@ -252,6 +252,27 @@ func TestSetMultipleExtendedLimits(t *testing.T) {
 	}
 }
 
+func TestSetResourceLimitsCPUAffinity(t *testing.T) {
+	job, err := Create(context.Background(), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer job.Close()
+
+	limits := &JobLimits{CPUAffinity: 0x3}
+	if err := job.SetResourceLimits(limits); err != nil {
+		t.Fatalf("failed to set resource limits with cpu affinity: %v", err)
+	}
+
+	affinity, err := job.GetCPUAffinity()
+	if err != nil {
+		t.Fatalf("failed to query cpu affinity: %v", err)
+	}
+	if affinity != limits.CPUAffinity {
+		t.Fatalf("unexpected cpu affinity: got %d want %d", affinity, limits.CPUAffinity)
+	}
+}
+
 func TestNoMoreProcessesMessageKill(t *testing.T) {
 	// Test that we receive the no more processes in job message after killing all of
 	// the processes in the job.

internal/jobobject/limits.go

@@ -38,6 +38,12 @@ func (job *JobObject) SetResourceLimits(limits *JobLimits) error {
 		}
 	}
 
+	if limits.CPUAffinity != 0 {
+		if err := job.SetCPUAffinity(limits.CPUAffinity); err != nil {
+			return fmt.Errorf("failed to set job object cpu affinity: %w", err)
+		}
+	}
+
 	if limits.MaxBandwidth != 0 || limits.MaxIOPS != 0 {
 		if err := job.SetIOLimit(limits.MaxBandwidth, limits.MaxIOPS); err != nil {
 			return fmt.Errorf("failed to set io limit on job object: %w", err)