Bug: Windows numCPU returning incorrectly for values above 64

agent/api/task/task_windows.go

@@ -17,7 +17,6 @@
 package task
 
 import (
-	"runtime"
 	"time"
 
 	"github.com/aws/amazon-ecs-agent/agent/ecscni"
@@ -33,6 +32,7 @@ import (
 	taskresourcevolume "github.com/aws/amazon-ecs-agent/agent/taskresource/volume"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	ni "github.com/aws/amazon-ecs-agent/ecs-agent/netlib/model/networkinterface"
+	eautils "github.com/aws/amazon-ecs-agent/ecs-agent/utils"
 	"github.com/cihub/seelog"
 	dockercontainer "github.com/docker/docker/api/types/container"
 	"github.com/pkg/errors"
@@ -59,7 +59,7 @@ type PlatformFields struct {
 	MemoryUnbounded config.BooleanDefaultFalse `json:"memoryUnbounded"`
 }
 
-var cpuShareScaleFactor = runtime.NumCPU() * cpuSharesPerCore
+var cpuShareScaleFactor = eautils.GetNumCPU() * cpuSharesPerCore
 
 // adjustForPlatform makes Windows-specific changes to the task after unmarshal
 func (task *Task) adjustForPlatform(cfg *config.Config) {

agent/api/task/task_windows_test.go

@@ -19,7 +19,6 @@ package task
 import (
 	"encoding/json"
 	"fmt"
-	"runtime"
 	"testing"
 
 	apicontainer "github.com/aws/amazon-ecs-agent/agent/api/container"
@@ -33,6 +32,7 @@ import (
 	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
 	apitaskstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/task/status"
 	ni "github.com/aws/amazon-ecs-agent/ecs-agent/netlib/model/networkinterface"
+	eautils "github.com/aws/amazon-ecs-agent/ecs-agent/utils"
 	"github.com/golang/mock/gomock"
 
 	"github.com/aws/amazon-ecs-agent/agent/dockerclient"
@@ -222,7 +222,7 @@ func TestDockerHostConfigRawConfigMerging(t *testing.T) {
 }
 
 func TestCPUPercentBasedOnUnboundedEnabled(t *testing.T) {
-	cpuShareScaleFactor := runtime.NumCPU() * cpuSharesPerCore
+	cpuShareScaleFactor := eautils.GetNumCPU() * cpuSharesPerCore
 	testcases := []struct {
 		cpu          int64
 		cpuUnbounded config.BooleanDefaultFalse

agent/stats/utils.go

@@ -15,12 +15,12 @@ package stats
 
 import (
 	"math"
-	"runtime"
 
+	eautils "github.com/aws/amazon-ecs-agent/ecs-agent/utils"
 	"github.com/docker/docker/api/types"
 )
 
-var numCores = uint64(runtime.NumCPU())
+var numCores = uint64(eautils.GetNumCPU())
 
 // nan32 returns a 32bit NaN.
 func nan32() float32 {

ecs-agent/api/ecs/client/ecs_client.go

@@ -18,7 +18,6 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
-	"runtime"
 	"strings"
 	"time"
 
@@ -417,7 +416,7 @@ func getCpuAndMemory() (int64, int64) {
 		})
 	}
 
-	cpu := runtime.NumCPU() * 1024
+	cpu := utils.GetNumCPU() * 1024
 
 	return int64(cpu), mem
 }

ecs-agent/utils/utils_unix.go

@@ -0,0 +1,23 @@
+//go:build !windows
+// +build !windows
+
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file 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 utils
+
+import "runtime"
+
+func GetNumCPU() int {
+	return runtime.NumCPU()
+}

ecs-agent/utils/utils_windows.go

@@ -0,0 +1,42 @@
+//go:build windows
+// +build windows
+
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file 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 utils
+
+import (
+	"runtime"
+
+	"golang.org/x/sys/windows"
+)
+
+var (
+	win32APIGetAllActiveProcessorCount = windows.GetActiveProcessorCount
+	golangRuntimeNumCPU                = runtime.NumCPU
+)
+
+// GetNumCPU On Windows, runtime.NumCPU() does not return the correct value for vCPU > 64
+// To resolve this, we are using MSFT implementation of the function: https://github.com/microsoft/hcsshim/blob/dd45838a9bf9ff8f431847aaf3e4421763c15c49/internal/processorinfo/processor_count.go#L14
+// Background around issue: https://github.com/moby/moby/issues/38935, https://learn.microsoft.com/en-us/windows/win32/procthread/processor-groups
+// Explanation: As was mentioned above, many of the Windows processor affinity functions will only return the information for a single Processor Group.
+// Since a single group can only hold 64 logical processors, this means when there are more they will be divided into multiple groups.
+// Golang runtime.NumCPU will only return the value of one Processor Group (not the sum of all).
+func GetNumCPU() int {
+	if amount := win32APIGetAllActiveProcessorCount(windows.ALL_PROCESSOR_GROUPS); amount != 0 {
+		return int(amount)
+	}
+	// If the Win32 API does not return correctly, fall back to runtime.NumCPU()
+	return golangRuntimeNumCPU()
+}

ecs-agent/utils/utils_windows_test.go

@@ -0,0 +1,41 @@
+package utils
+
+import (
+	"runtime"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"golang.org/x/sys/windows"
+)
+
+func TestGetNumCPU(t *testing.T) {
+	testCases := []struct {
+		win32APIReturn      uint32
+		runtimeNumCPUReturn int
+		expectedAnswer      int
+		name                string
+	}{
+		{128, 64, 128, "Both are valid"},
+		{0, 64, 64, "win32 API invalid"},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			defer func() {
+				win32APIGetAllActiveProcessorCount = windows.GetActiveProcessorCount
+				golangRuntimeNumCPU = runtime.NumCPU
+			}()
+
+			win32APIGetAllActiveProcessorCount = func(groupNumber uint16) (ret uint32) {
+				return tc.win32APIReturn
+			}
+
+			golangRuntimeNumCPU = func() int {
+				return tc.runtimeNumCPUReturn
+			}
+
+			assert.Equal(t, tc.expectedAnswer, GetNumCPU(), tc.name)
+		})
+	}
+
+}