From 40cfdd406d5834933ff28dcdd25c8f9f8ee7ee5d Mon Sep 17 00:00:00 2001
From: hveldstra <h@veldstra.org>
Date: Tue, 6 Aug 2024 13:32:21 +0100
Subject: [PATCH] feat(ecs): add ability to set DNS servers and ephemeral
 storage size

- --container-dns-servers can now be used to set the dnsServers setting
  for the Artillery container
- --ephemeral-storage can now be used to set the amount of ephemeral storage for the worker task
---
 packages/artillery/lib/cmds/run-fargate.js             | 10 ++++++++++
 .../lib/platform/aws-ecs/legacy/run-cluster.js         | 10 ++++++++++
 2 files changed, 20 insertions(+)

diff --git a/packages/artillery/lib/cmds/run-fargate.js b/packages/artillery/lib/cmds/run-fargate.js
index 000b1418be..29b8530749 100644
--- a/packages/artillery/lib/cmds/run-fargate.js
+++ b/packages/artillery/lib/cmds/run-fargate.js
@@ -129,6 +129,16 @@ RunCommand.flags = {
     description:
       'JSON to customize launch configuration of ECS/Fargate tasks (see https://www.artillery.io/docs/reference/cli/run-fargate#using---launch-config)'
   }),
+  'container-dns-servers': Flags.string({
+    description:
+      'Comma-separated list of DNS servers for Artillery container. Maps to dnsServers parameter in ECS container definition'
+  }),
+  'task-ephemeral-storage': Flags.string({
+    description:
+      'Ephemeral storage in GiB for the worker task. Maps to ephemeralStorage parameter in ECS container definition. Fargate-only.',
+    type: 'integer',
+  }),
+
   'subnet-ids': Flags.string({
     description:
       'Comma-separated list of AWS VPC subnet IDs to launch Fargate tasks in'
diff --git a/packages/artillery/lib/platform/aws-ecs/legacy/run-cluster.js b/packages/artillery/lib/platform/aws-ecs/legacy/run-cluster.js
index 2a8fc6cef9..b3e3219bad 100644
--- a/packages/artillery/lib/platform/aws-ecs/legacy/run-cluster.js
+++ b/packages/artillery/lib/platform/aws-ecs/legacy/run-cluster.js
@@ -1203,6 +1203,10 @@ async function ensureTaskExists(context) {
       }
     };
 
+    if (context.cliOptions.containerDnsServers) {
+      artilleryContainerDefinition.dnsServers = context.cliOptions.containerDnsServers.split(',');
+    }
+  
     let taskDefinition = {
       family: context.taskName,
       containerDefinitions: [artilleryContainerDefinition],
@@ -1495,6 +1499,12 @@ async function generateTaskOverrides(context) {
     overrides.taskRoleArn = context.customRoleArn;
   }
 
+  if (context.cliOptions.taskEphemeralStorage) {
+    overrides.ephemeralStorage = {
+      sizeInGiB: context.cliOptions.taskEphemeralStorage
+    };
+  }
+
   overrides.containerOverrides[0].environment.push({
     name: 'USE_V2',
     value: 'true'