From b4350e349e2e0d78ada72c7fa06810351cd8580e Mon Sep 17 00:00:00 2001
From: Kazuho CryerShinozuka <malaysia.cryer@gmail.com>
Date: Sat, 9 Dec 2023 22:12:10 +0900
Subject: [PATCH] feat(ecs-patterns): add taskDefinition props to
 QueueProcessingEc2Service

---
 .../lib/ecs/queue-processing-ecs-service.ts   | 47 +++++++++++-------
 .../ec2/queue-processing-ecs-service.test.ts  | 48 +++++++++++++++++++
 2 files changed, 79 insertions(+), 16 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-ecs-patterns/lib/ecs/queue-processing-ecs-service.ts b/packages/aws-cdk-lib/aws-ecs-patterns/lib/ecs/queue-processing-ecs-service.ts
index fa53d91f0898d..f1afd90f2641f 100644
--- a/packages/aws-cdk-lib/aws-ecs-patterns/lib/ecs/queue-processing-ecs-service.ts
+++ b/packages/aws-cdk-lib/aws-ecs-patterns/lib/ecs/queue-processing-ecs-service.ts
@@ -8,6 +8,13 @@ import { QueueProcessingServiceBase, QueueProcessingServiceBaseProps } from '../
  * The properties for the QueueProcessingEc2Service service.
  */
 export interface QueueProcessingEc2ServiceProps extends QueueProcessingServiceBaseProps {
+  /**
+  * The task definition to use for tasks in the service. TaskDefinition or Image must be specified, but not both..
+  *
+  * @default - none
+  */
+  readonly taskDefinition?: Ec2TaskDefinition;
+
   /**
    * The number of cpu units used by the task.
    *
@@ -106,23 +113,31 @@ export class QueueProcessingEc2Service extends QueueProcessingServiceBase {
   constructor(scope: Construct, id: string, props: QueueProcessingEc2ServiceProps) {
     super(scope, id, props);
 
-    const containerName = props.containerName ?? 'QueueProcessingContainer';
+    if (props.taskDefinition != null && props.image != null) {
+      throw new Error('You must specify only one of TaskDefinition or Image');
+    } else if (props.taskDefinition != null) {
+      this.taskDefinition = props.taskDefinition;
+    } else if (props.image != null) {
+      // Create a Task Definition for the container to start
+      this.taskDefinition = new Ec2TaskDefinition(this, 'QueueProcessingTaskDef', {
+        family: props.family,
+      });
 
-    // Create a Task Definition for the container to start
-    this.taskDefinition = new Ec2TaskDefinition(this, 'QueueProcessingTaskDef', {
-      family: props.family,
-    });
-    this.taskDefinition.addContainer(containerName, {
-      image: props.image,
-      memoryLimitMiB: props.memoryLimitMiB,
-      memoryReservationMiB: props.memoryReservationMiB,
-      cpu: props.cpu,
-      gpuCount: props.gpuCount,
-      command: props.command,
-      environment: this.environment,
-      secrets: this.secrets,
-      logging: this.logDriver,
-    });
+      const containerName = props.containerName ?? 'QueueProcessingContainer';
+      this.taskDefinition.addContainer(containerName, {
+        image: props.image,
+        memoryLimitMiB: props.memoryLimitMiB,
+        memoryReservationMiB: props.memoryReservationMiB,
+        cpu: props.cpu,
+        gpuCount: props.gpuCount,
+        command: props.command,
+        environment: this.environment,
+        secrets: this.secrets,
+        logging: this.logDriver,
+      });
+    } else {
+      throw new Error('You must specify one of: taskDefinition or image');
+    }
 
     // The desiredCount should be removed from the fargate service when the feature flag is removed.
     const desiredCount = FeatureFlags.of(this).isEnabled(cxapi.ECS_REMOVE_DEFAULT_DESIRED_COUNT) ? undefined : this.desiredCount;
diff --git a/packages/aws-cdk-lib/aws-ecs-patterns/test/ec2/queue-processing-ecs-service.test.ts b/packages/aws-cdk-lib/aws-ecs-patterns/test/ec2/queue-processing-ecs-service.test.ts
index 7be7e67aa95c5..459c9bf301b3c 100644
--- a/packages/aws-cdk-lib/aws-ecs-patterns/test/ec2/queue-processing-ecs-service.test.ts
+++ b/packages/aws-cdk-lib/aws-ecs-patterns/test/ec2/queue-processing-ecs-service.test.ts
@@ -568,3 +568,51 @@ it('throws validation errors of the specific queue prop, when setting queue and
     });
   }).toThrow(new Error('visibilityTimeout can be set only when queue is not set. Specify them in the QueueProps of the queue'));
 });
+
+test('test ECS queue worker service construct - with task definition', () => {
+  // GIVEN
+  const stack = new cdk.Stack();
+  const vpc = new ec2.Vpc(stack, 'VPC');
+  const cluster = new ecs.Cluster(stack, 'Cluster', { vpc });
+  cluster.addAsgCapacityProvider(new AsgCapacityProvider(stack, 'DefaultAutoScalingGroupProvider', {
+    autoScalingGroup: new AutoScalingGroup(stack, 'DefaultAutoScalingGroup', {
+      vpc,
+      instanceType: new ec2.InstanceType('t2.micro'),
+      machineImage: MachineImage.latestAmazonLinux(),
+    }),
+  }));
+
+  // WHEN
+  const taskDefinition = new ecs.Ec2TaskDefinition(stack, 'TaskDef', {
+    family: 'MyTaskDefinitionFamily',
+  });
+
+  taskDefinition.addContainer('TheContainer', {
+    image: ecs.ContainerImage.fromRegistry('test'),
+    memoryLimitMiB: 1024,
+    logging: new ecs.AwsLogDriver({ streamPrefix: 'QueueProcessingFargateService' }),
+  });
+
+  new ecsPatterns.QueueProcessingEc2Service(stack, 'Service', {
+    cluster,
+    taskDefinition,
+  });
+
+  // THEN
+  Template.fromStack(stack).hasResourceProperties('AWS::ECS::TaskDefinition', {
+    ContainerDefinitions: [
+      Match.objectLike({
+        Essential: true,
+        Image: 'test',
+        LogConfiguration: {},
+        Memory: 1024,
+        Name: 'TheContainer',
+      }),
+    ],
+    ExecutionRoleArn: { 'Fn::GetAtt': ['TaskDefExecutionRoleB4775C97', 'Arn'] },
+    Family: 'MyTaskDefinitionFamily',
+    NetworkMode: 'bridge',
+    RequiresCompatibilities: ['EC2'],
+    TaskRoleArn: { 'Fn::GetAtt': ['TaskDefTaskRole1EDB4A67', 'Arn'] },
+  });
+});