make consumable() return a list of non-consumable resource keys

singholt · Sep 11, 2024 · f357ff6 · f357ff6
1 parent edd5f2e
commit f357ff6
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 17 deletions.
diff --git a/agent/engine/host_resource_manager.go b/agent/engine/host_resource_manager.go
@@ -115,7 +115,7 @@ func (h *HostResourceManager) consume(taskArn string, resources map[string]*ecs.
 		return true, nil
 	}
 
-	ok, failedResourceKey, err := h.consumable(resources)
+	ok, failedResourceKeys, err := h.consumable(resources)
 	if err != nil {
 		logger.Error("Resources failing to consume, error in task resources", logger.Fields{
 			"taskArn":   taskArn,
@@ -140,8 +140,8 @@ func (h *HostResourceManager) consume(taskArn string, resources map[string]*ecs.
 		return true, nil
 	}
 	logger.Info("Resources not consumed, enough resources not available", logger.Fields{
-		"taskArn":  taskArn,
-		"resource": failedResourceKey,
+		"taskArn":   taskArn,
+		"resources": failedResourceKeys,
 	})
 	return false, nil
 }
@@ -254,31 +254,37 @@ func (h *HostResourceManager) checkResourcesHealth(resources map[string]*ecs.Res
 }
 
 // Helper function for consume to check if resources are consumable with the current account
-// we have for the host resources. Should not call host resource manager lock in this func
-// return values
-func (h *HostResourceManager) consumable(resources map[string]*ecs.Resource) (bool, string, error) {
+// we have for the host resources. Should not call host resource manager lock in this func return values
+// This function returns a bool (indicating whether ALL requested resources are consumable), a list of non-consumable
+// resource keys, and error, if any.
+func (h *HostResourceManager) consumable(resources map[string]*ecs.Resource) (bool, []string, error) {
 	err := h.checkResourcesHealth(resources)
 	if err != nil {
-		return false, "", err
+		return false, nil, err
 	}
 
+	var resourcesNotConsumable []string
 	for resourceKey := range resources {
 		if *resources[resourceKey].Type == "INTEGER" {
 			consumable := h.checkConsumableIntType(resourceKey, resources)
 			if !consumable {
-				return false, resourceKey, nil
+				resourcesNotConsumable = append(resourcesNotConsumable, resourceKey)
 			}
 		}
 
 		if *resources[resourceKey].Type == "STRINGSET" {
 			consumable := h.checkConsumableStringSetType(resourceKey, resources)
 			if !consumable {
-				return false, resourceKey, nil
+				resourcesNotConsumable = append(resourcesNotConsumable, resourceKey)
 			}
 		}
 	}
 
-	return true, "", nil
+	if resourcesNotConsumable != nil {
+		return false, resourcesNotConsumable, nil
+	} else {
+		return true, nil, nil
+	}
 }
 
 // Utility function to manage release of ports

diff --git a/agent/engine/host_resource_manager_test.go b/agent/engine/host_resource_manager_test.go
@@ -17,6 +17,7 @@
 package engine
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/aws/amazon-ecs-agent/agent/utils"
@@ -191,7 +192,7 @@ func TestConsumable(t *testing.T) {
 		portsUdp                  []uint16
 		gpus                      []string
 		canBeConsumed             bool
-		expectedFailedResourceKey string
+		expectedFailedResourceKey []string
 	}{
 		{
 			name:                      "consumable",
@@ -201,7 +202,7 @@ func TestConsumable(t *testing.T) {
 			portsUdp:                  []uint16{1003},
 			gpus:                      []string{"gpu1", "gpu2"},
 			canBeConsumed:             true,
-			expectedFailedResourceKey: "",
+			expectedFailedResourceKey: nil,
 		},
 		{
 			name:                      "cpu not consumable",
@@ -211,7 +212,7 @@ func TestConsumable(t *testing.T) {
 			portsUdp:                  []uint16{},
 			gpus:                      []string{},
 			canBeConsumed:             false,
-			expectedFailedResourceKey: "CPU",
+			expectedFailedResourceKey: []string{"CPU"},
 		},
 		{
 			name:                      "memory not consumable",
@@ -221,7 +222,7 @@ func TestConsumable(t *testing.T) {
 			portsUdp:                  []uint16{},
 			gpus:                      []string{},
 			canBeConsumed:             false,
-			expectedFailedResourceKey: "MEMORY",
+			expectedFailedResourceKey: []string{"MEMORY"},
 		},
 		{
 			name:                      "tcp ports not consumable",
@@ -231,7 +232,7 @@ func TestConsumable(t *testing.T) {
 			portsUdp:                  []uint16{},
 			gpus:                      []string{},
 			canBeConsumed:             false,
-			expectedFailedResourceKey: "PORTS_TCP",
+			expectedFailedResourceKey: []string{"PORTS_TCP"},
 		},
 		{
 			name:                      "udp ports not consumable",
@@ -241,7 +242,17 @@ func TestConsumable(t *testing.T) {
 			portsUdp:                  []uint16{1000},
 			gpus:                      []string{},
 			canBeConsumed:             false,
-			expectedFailedResourceKey: "PORTS_UDP",
+			expectedFailedResourceKey: []string{"PORTS_UDP"},
+		},
+		{
+			name:                      "multiple resources not consumable - cpu and udp ports",
+			cpu:                       int64(2500),
+			mem:                       int64(1024),
+			ports:                     []uint16{},
+			portsUdp:                  []uint16{1000},
+			gpus:                      []string{},
+			canBeConsumed:             false,
+			expectedFailedResourceKey: []string{"CPU", "PORTS_UDP"},
 		},
 	}
 
@@ -253,7 +264,7 @@ func TestConsumable(t *testing.T) {
 			"Error in checking if resources can be successfully consumed")
 		assert.Equal(t, nil, err,
 			"Error in checking if resources can be successfully consumed, error returned from consumable")
-		assert.Equal(t, tc.expectedFailedResourceKey, failedResourceKey)
+		assert.True(t, reflect.DeepEqual(tc.expectedFailedResourceKey, failedResourceKey))
 	}
 }