From 82ec0ff68fe31ee4b4fdc04a9dbd35131d234677 Mon Sep 17 00:00:00 2001
From: cohalz <cohalzweb@gmail.com>
Date: Sun, 6 Jan 2019 21:14:10 +0900
Subject: [PATCH] feat(aws-ecs): add support Amazon Linux 2 (#1484)

Fixes #1483.
---
 packages/@aws-cdk/aws-ecs/README.md      |  2 ++
 packages/@aws-cdk/aws-ecs/lib/cluster.ts | 25 +++++++++++++++++++++---
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/packages/@aws-cdk/aws-ecs/README.md b/packages/@aws-cdk/aws-ecs/README.md
index 83479b96eb56c..6fa8e9f251116 100644
--- a/packages/@aws-cdk/aws-ecs/README.md
+++ b/packages/@aws-cdk/aws-ecs/README.md
@@ -88,6 +88,8 @@ const autoScalingGroup = new autoscaling.AutoScalingGroup(this, 'ASG', {
   vpc,
   instanceType: new ec2.InstanceType('t2.xlarge'),
   machineImage: new EcsOptimizedAmi(),
+  // Or use ECS-Optimized Amazon Linux 2 AMI
+  // machineImage: new EcsOptimizedAmi({ generation: ec2.AmazonLinuxGeneration.AmazonLinux2 }),
   desiredCapacity: 3,
   // ... other options here ...
 });
diff --git a/packages/@aws-cdk/aws-ecs/lib/cluster.ts b/packages/@aws-cdk/aws-ecs/lib/cluster.ts
index ac90dc3660d60..dfc6d5cc3fe0f 100644
--- a/packages/@aws-cdk/aws-ecs/lib/cluster.ts
+++ b/packages/@aws-cdk/aws-ecs/lib/cluster.ts
@@ -184,18 +184,37 @@ export class Cluster extends cdk.Construct implements ICluster {
   }
 }
 
+export interface EcsOptimizedAmiProps {
+  /**
+   * What generation of Amazon Linux to use
+   *
+   * @default AmazonLinux
+   */
+  generation?: ec2.AmazonLinuxGeneration;
+}
+
 /**
  * Construct a Linux machine image from the latest ECS Optimized AMI published in SSM
  */
-export class EcsOptimizedAmi implements ec2.IMachineImageSource  {
-  private static AmiParameterName = "/aws/service/ecs/optimized-ami/amazon-linux/recommended";
+export class EcsOptimizedAmi implements ec2.IMachineImageSource {
+  private readonly generation: ec2.AmazonLinuxGeneration;
+  private readonly amiParameterName: string;
+
+  constructor(props?: EcsOptimizedAmiProps) {
+    this.generation = (props && props.generation) || ec2.AmazonLinuxGeneration.AmazonLinux;
+    if (this.generation === ec2.AmazonLinuxGeneration.AmazonLinux2) {
+      this.amiParameterName = "/aws/service/ecs/optimized-ami/amazon-linux-2/recommended";
+    } else {
+      this.amiParameterName = "/aws/service/ecs/optimized-ami/amazon-linux/recommended";
+    }
+  }
 
   /**
    * Return the correct image
    */
   public getImage(scope: cdk.Construct): ec2.MachineImage {
     const ssmProvider = new cdk.SSMParameterProvider(scope, {
-        parameterName: EcsOptimizedAmi.AmiParameterName
+      parameterName: this.amiParameterName
     });
 
     const json = ssmProvider.parameterValue("{\"image_id\": \"\"}");