From 1c9c5a70e9cc8ac64bb14df13c36e8fd0df92697 Mon Sep 17 00:00:00 2001
From: Tatsuya Mori <tamori@amazon.com>
Date: Sat, 17 Jun 2023 10:27:08 +0000
Subject: [PATCH] fix(ecs): validation for task definition fails when
 task-level memory is defined but container-level memory and memoryReservation
 are not defined with EC2 compatibility

---
 .../aws-ecs/lib/base/task-definition.ts       | 11 +++--
 .../aws-ecs/test/task-definition.test.ts      | 42 ++++++++++++++++++-
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-ecs/lib/base/task-definition.ts b/packages/aws-cdk-lib/aws-ecs/lib/base/task-definition.ts
index 78fd14f3ce95e..1280af3e5cd09 100644
--- a/packages/aws-cdk-lib/aws-ecs/lib/base/task-definition.ts
+++ b/packages/aws-cdk-lib/aws-ecs/lib/base/task-definition.ts
@@ -400,6 +400,8 @@ export class TaskDefinition extends TaskDefinitionBase {
 
   private readonly _cpu?: string;
 
+  private readonly _memory?: string;
+
   /**
    * Constructs a new instance of the TaskDefinition class.
    */
@@ -461,6 +463,7 @@ export class TaskDefinition extends TaskDefinitionBase {
 
     this.runtimePlatform = props.runtimePlatform;
     this._cpu = props.cpu;
+    this._memory = props.memoryMiB;
 
     const taskDef = new CfnTaskDefinition(this, 'Resource', {
       containerDefinitions: Lazy.any({ produce: () => this.renderContainers() }, { omitEmptyArray: true }),
@@ -736,9 +739,11 @@ export class TaskDefinition extends TaskDefinitionBase {
       // EC2 mode validations
 
       // Container sizes
-      for (const container of this.containers) {
-        if (!container.memoryLimitSpecified) {
-          ret.push(`ECS Container ${container.containerName} must have at least one of 'memoryLimitMiB' or 'memoryReservationMiB' specified`);
+      if (!this._memory) {
+        for (const container of this.containers) {
+          if (!container.memoryLimitSpecified) {
+            ret.push(`ECS Container ${container.containerName} must have at least one of 'memoryLimitMiB' or 'memoryReservationMiB' specified`);
+          }
         }
       }
     }
diff --git a/packages/aws-cdk-lib/aws-ecs/test/task-definition.test.ts b/packages/aws-cdk-lib/aws-ecs/test/task-definition.test.ts
index 5f50ffe44a59b..571bf6ba61f6d 100644
--- a/packages/aws-cdk-lib/aws-ecs/test/task-definition.test.ts
+++ b/packages/aws-cdk-lib/aws-ecs/test/task-definition.test.ts
@@ -275,6 +275,46 @@ describe('task definition', () => {
         }],
       });
     });
+
+    test('You can omit container-level memory and memoryReservation parameters with EC2 compatibilities if task-level memory parameter is defined', () => {
+      // GIVEN
+      const stack = new cdk.Stack();
+
+      // WHEN
+      const taskDef = new ecs.TaskDefinition(stack, 'TD', {
+        cpu: '512',
+        memoryMiB: '512',
+        compatibility: ecs.Compatibility.EC2,
+      });
+      taskDef.addContainer('Container', {
+        image: ecs.ContainerImage.fromRegistry('/aws/aws-example-app'),
+      });
+
+      // THEN
+      Template.fromStack(stack).hasResourceProperties('AWS::ECS::TaskDefinition', {
+        Memory: '512',
+      });
+
+    });
+
+    test('A task definition where task-level memory, container-level memory and memoryReservation are not defined throws an error', () => {
+      // GIVEN
+      const stack = new cdk.Stack();
+
+      // WHEN
+      const taskDef = new ecs.TaskDefinition(stack, 'TD', {
+        cpu: '512',
+        compatibility: ecs.Compatibility.EC2,
+      });
+      taskDef.addContainer('Container', {
+        image: ecs.ContainerImage.fromRegistry('/aws/aws-example-app'),
+      });
+
+      // THEN
+      expect(() => {
+        Template.fromStack(stack);
+      }).toThrow("ECS Container Container must have at least one of 'memoryLimitMiB' or 'memoryReservationMiB' specified");
+    });
   });
 
   describe('When importing from an existing Task definition', () => {
@@ -476,4 +516,4 @@ describe('task definition', () => {
       });
     });
   });
-});
+});
\ No newline at end of file