update package-lock.json

.github/workflows/build.yml

@@ -1,5 +1,9 @@
 name: Build
-on: [push, workflow_dispatch]
+on: 
+  push:
+  workflow_dispatch:
+  pull_request_target:
+
 jobs:
   Build-and-Test-CDK:
     runs-on: ubuntu-latest
@@ -12,3 +16,5 @@ jobs:
       - run: npm ci
       - run: npm run build
       - run: npm run test
+      - run: npx cdk synth
+

.gitignore

@@ -6,3 +6,7 @@ node_modules
 # CDK asset staging directory
 .cdk.staging
 cdk.out
+
+# idea / intelliJ
+.idea
+*.iml

README.md

@@ -1,21 +1,27 @@
 # Distributed Load Testing with Locust on Amazon ECS
 [![Build](https://github.com/aws-samples/distributed-load-testing-with-locust-on-ecs/actions/workflows/build.yml/badge.svg)](https://github.com/aws-samples/distributed-load-testing-with-locust-on-ecs/actions/workflows/build.yml)
 
-This sample shows you how to deploy [Locust](https://locust.io/), a modern load testing framework, to Amazon Elastic Container Service (ECS). It leverages a serverless compute engine [Fargate](https://aws.amazon.com/fargate/) with spot capacity, which allows you to run massive-scale load test without managing infrastructure and with relatively low cost (70% cheaper than using on-demand capacity).
+This sample shows you how to deploy [Locust](https://locust.io/), a modern load testing framework, to Amazon Elastic Container Service (ECS).
+
+* ✅ From small to massive-scale load test with AWS serverless technologies
+* ✅ Highly cost-efficient with Fargate spot capcity
+* ✅ Instant deployment using AWS CDK
+
+It leverages a serverless compute engine [Fargate](https://aws.amazon.com/fargate/) with spot capacity, which allows you to run massive-scale load test without managing infrastructure and with relatively low cost (70% cheaper than using on-demand capacity).
 
 ## How it works
 Below is the architecture diagram of this sample.
 
 ![architecture](imgs/architecture.png)
 
-We deploy Locust with distributed mode, so there are two ECS services, master service and worker service.
+We deploy Locust with distributed mode, hence two ECS services - Locust master and worker service.
 
-The number of  Locust master instance is always one, and it can be accessed via Application Load Balancer.
+The Locust master consists of a single Fargate task, and its Web GUI can be accessed via Application Load Balancer.
 
-On the other hand, there can be *N* Locust worker instances, which is usually the dominant factor of load test infrastructure cost.
-We use Fargate spot capacity for worker instances, which allows you to run load test at most 70% cheaper than on-demand capacity.
+Unlike master node, there can be *N* Locust worker nodes, which is usually the dominant factor of load test infrastructure cost.
+We use Fargate spot capacity for Locust workers, allowing you to run load tests at most 70% cheaper than on-demand capacity.
 
-Note that all the access from Locust workers go through NAT Gateway, which makes it easy to restrict access by IP addresses on load test target servers, because all the Locust workers shares the same outbound IP address among them.
+Note that all the access from Locust workers go through NAT Gateway, which makes it easy to restrict access by IP addresses on load test target servers, because all the Locust workers shares the same outbound IP address.
 
 ## Deploy
 To deploy this sample to your own AWS account, please follow the steps below.
@@ -30,7 +36,7 @@ Before you deploy, make sure you install the following tools in your local envir
 Also you need Administorator IAM policy to deploy this sample.
 
 ### 2. Set parameters
-Before deploy, you need to set some parameters.
+You need to set several parameters to configure the system.
 
 Please open [bin/load_test.ts](./bin/load_test.ts) and find property named `allowedCidrs`.
 This property specifies the CIDRs which can access the Locust web UI ALB.
@@ -90,9 +96,9 @@ Now the deployment is completed! You can start to use Locust load tester.
 There are a few things you need to know to use this sample effectively.
 
 ### Adjust the number of Locust worker tasks
-According to the amount of load you want to generate, you may need to increase Locust workers.
+Depending on the amount of load you want to generate, you may need to increase Locust worker capacity.
 
-It can be done with the following command:
+It can be adjusted with the following command:
 
 ```sh
 aws ecs update-service --cluster <EcsClusterArn> --service <WorkerServiceName> --desired-count <the number of workers>
@@ -104,8 +110,9 @@ Please also be aware that your default quota for the number of Fargate tasks is
 If you need more tasks, you can request a limit increase from [Service Quotas console](https://console.aws.amazon.com/servicequotas/home). You can read further detail [here](https://docs.aws.amazon.com/general/latest/gr/aws_service_limits.html).
 
 ### When Fargate spot is out of capacity
-It is expected that sometimes Fargate spot fails to run your Locust workers because of insufficient capacity.
-If such situation continues for unacceptable time, you can add on-demand instances to fill your desired task count.
+It is expected that sometimes Fargate spot fails to allocate the required capacity for your Locust workers because of insufficient capacity.
+That issue should resolve if you wait for certain time.
+However, if it continues for unacceptable time, you can always add on-demand capacity to fill your desired task count.
 
 Please open [`lib/constructs/locust_worker_service.ts`](lib/constructs/locust_worker_service.ts) and find the lines below:
 
@@ -122,7 +129,7 @@ Please open [`lib/constructs/locust_worker_service.ts`](lib/constructs/locust_wo
       ],
 ```
 
-You can specify the ratio of spot vs on-demand by `weight` property. The default is to use spot 100%.
+You can specify the ratio of spot (`FARGATE_SPOT`) vs on-demand (`FARGATE`) by the `weight` properties. The default is to use spot 100% (1:0).
 
 ### Modify Locust scenario
 Default locustfile is placed on [`./app/locustfile.py`](app/locustfile.py).

app/Dockerfile

@@ -1,8 +1,8 @@
-FROM --platform=linux/amd64 locustio/locust:latest
+FROM --platform=linux/amd64 locustio/locust:2.26.0
 COPY . ./
 
 # for standalone
-ENTRYPOINT ["locust", "-f", "./locustfile.py"]
+ENTRYPOINT ["locust", "-f", "./locustfile.py", "--modern-ui"]
 # for worker
 # CMD [ "--worker", "--master-host", "MASTER_HOST_NAME"]
 # for master

app/auth.py

@@ -0,0 +1,65 @@
+import os
+
+from locust import events
+from flask import Blueprint, redirect, request, session, url_for
+from flask_login import UserMixin, login_user
+
+
+"""
+Example of implementing authentication for Locust when the --web-login flag is given
+
+This is only to serve as a starting point, proper authentication should be implemented
+according to your projects specifications.
+
+For more information, see https://docs.locust.io/en/stable/extending-locust.html#authentication
+"""
+
+
+class AuthUser(UserMixin):
+    def __init__(self, username):
+        self.username = username
+
+    def get_id(self):
+        return self.username
+
+
+auth_blueprint = Blueprint("auth", "web_ui_auth")
+
+
+def load_user(user_id):
+    return AuthUser(session.get("username"))
+
+
+@events.init.add_listener
+def locust_init(environment, **kwargs):
+    if environment.web_ui:
+        environment.web_ui.login_manager.user_loader(load_user)
+
+        environment.web_ui.app.config["SECRET_KEY"] = os.getenv("FLASK_SECRET_KEY")
+        cfg_password = os.getenv("LOCUST_PASSWORD")
+        cfg_username = os.getenv("LOCUST_USERNAME")
+
+        environment.web_ui.auth_args = {
+            "username_password_callback": "/login_submit",
+        }
+
+        @auth_blueprint.route("/login_submit")
+        def login_submit():
+            username = request.args.get("username")
+            password = request.args.get("password")
+
+            # Implement real password verification here
+            if username == cfg_username and cfg_password == password:
+                session["username"] = username
+                login_user(AuthUser(username))
+
+                return redirect(url_for("index"))
+
+            environment.web_ui.auth_args = {
+                **environment.web_ui.auth_args,
+                "error": "Invalid username or password",
+            }
+
+            return redirect(url_for("login"))
+
+        environment.web_ui.app.register_blueprint(auth_blueprint)

app/locustfile.py

@@ -1,9 +1,11 @@
 from locust import HttpUser, task
 
+# enable password authentication
+import auth
+
+
 # Defining a sample task. You can write your own locust script here.
 class SampleUser(HttpUser):
     @task
     def get_index(self):
         self.client.get("/")
-        self.client.get("/actuator/health/liveness")
-        self.client.get("/actuator/health/readiness")

bin/load_test.ts

@@ -17,11 +17,16 @@ new LoadTestStack(app, 'LoadTestStack', {
 
   // CIDRs that can access Locust Web UI ALB.
   // It is highly recommended to set this CIDR as narrowly as possible
-  // since Locust Web UI does NOT have any authentication mechanism
-  // allow traffic from the world.
-  allowedCidrs: ['0.0.0.0/0'],
-
-  // You can enable basic auth for Locust web UI uncommenting lines below:
+  // when you do not enable the authentication option below.
+  allowedCidrs: ['127.0.0.1/32'],
+
+  // You can enable password auth for Locust web UI uncommenting lines below:
   // webUsername: 'admin',
   // webPassword: 'passw0rd',
+
+  // Any arbitrary command line options to pass to Locust.
+  // An example would be:
+  // Exclude Tags - List of tags to exclude from the test, so only tasks
+  // with no matching tags will be executed.
+  // additionalArguments: ['--exclude-tags', 'tag1', 'tag2'],
 });

cdk.json

@@ -17,14 +17,35 @@
     ]
   },
   "context": {
-    "@aws-cdk/aws-apigateway:usagePlanKeyOrderInsensitiveId": true,
-    "@aws-cdk/core:stackRelativeExports": true,
-    "@aws-cdk/aws-rds:lowercaseDbIdentifier": true,
-    "@aws-cdk/aws-lambda:recognizeVersionProps": true,
-    "@aws-cdk/aws-cloudfront:defaultSecurityPolicyTLSv1.2_2021": true,
+    "@aws-cdk/aws-lambda:recognizeLayerVersion": true,
+    "@aws-cdk/core:checkSecretUsage": true,
     "@aws-cdk/core:target-partitions": [
       "aws",
       "aws-cn"
-    ]
+    ],
+    "@aws-cdk-containers/ecs-service-extensions:enableDefaultLogDriver": true,
+    "@aws-cdk/aws-ec2:uniqueImdsv2TemplateName": true,
+    "@aws-cdk/aws-ecs:arnFormatIncludesClusterName": true,
+    "@aws-cdk/aws-iam:minimizePolicies": true,
+    "@aws-cdk/core:validateSnapshotRemovalPolicy": true,
+    "@aws-cdk/aws-codepipeline:crossAccountKeyAliasStackSafeResourceName": true,
+    "@aws-cdk/aws-s3:createDefaultLoggingPolicy": true,
+    "@aws-cdk/aws-sns-subscriptions:restrictSqsDescryption": true,
+    "@aws-cdk/aws-apigateway:disableCloudWatchRole": true,
+    "@aws-cdk/core:enablePartitionLiterals": true,
+    "@aws-cdk/aws-events:eventsTargetQueueSameAccount": true,
+    "@aws-cdk/aws-iam:standardizedServicePrincipals": true,
+    "@aws-cdk/aws-ecs:disableExplicitDeploymentControllerForCircuitBreaker": true,
+    "@aws-cdk/aws-iam:importedRoleStackSafeDefaultPolicyName": true,
+    "@aws-cdk/aws-s3:serverAccessLogsUseBucketPolicy": true,
+    "@aws-cdk/aws-route53-patters:useCertificate": true,
+    "@aws-cdk/customresources:installLatestAwsSdkDefault": false,
+    "@aws-cdk/aws-rds:databaseProxyUniqueResourceName": true,
+    "@aws-cdk/aws-codedeploy:removeAlarmsFromDeploymentGroup": true,
+    "@aws-cdk/aws-apigateway:authorizerChangeDeploymentLogicalId": true,
+    "@aws-cdk/aws-ec2:launchTemplateDefaultUserData": true,
+    "@aws-cdk/aws-secretsmanager:useAttachedSecretResourcePolicyForSecretTargetAttachments": true,
+    "@aws-cdk/aws-redshift:columnId": true,
+    "@aws-cdk/aws-stepfunctions-tasks:enableEmrServicePolicyV2": true
   }
 }

lib/constructs/locust_master_service.ts

@@ -8,13 +8,14 @@ import { IBucket } from 'aws-cdk-lib/aws-s3';
 import { Certificate } from 'aws-cdk-lib/aws-certificatemanager';
 import { RetentionDays } from 'aws-cdk-lib/aws-logs';
 import { LocustWorkerService } from './locust_worker_service';
+import { Platform } from 'aws-cdk-lib/aws-ecr-assets';
 
 export interface LocustMasterServiceProps {
-  readonly image: ecs.ContainerImage;
   readonly cluster: ecs.ICluster;
   readonly certificateArn?: string;
   readonly allowedCidrs: string[];
   readonly logBucket: IBucket;
+  readonly additionalArguments?: string[];
   readonly webUsername?: string;
   readonly webPassword?: string;
 }
@@ -26,9 +27,10 @@ export class LocustMasterService extends Construct {
   constructor(scope: Construct, id: string, props: LocustMasterServiceProps) {
     super(scope, id);
 
-    const { cluster, image, webUsername, webPassword } = props;
+    const { cluster, additionalArguments, webUsername, webPassword } = props;
 
     const configMapName = 'master';
+    const image = new ecs.AssetImage('app', { platform: Platform.LINUX_AMD64 });
 
     const protocol = props.certificateArn != null ? ApplicationProtocol.HTTPS : ApplicationProtocol.HTTP;
 
@@ -42,12 +44,17 @@ export class LocustMasterService extends Construct {
       memoryLimitMiB: 2048,
     });
 
+    let environment: { [key: string]: string } = {};
     const command = ['--master'];
     if (webUsername != null && webPassword != null) {
-      command.push('--web-auth');
-      command.push(`${webUsername}:${webPassword}`);
+      command.push('--web-login');
+      environment['LOCUST_USERNAME'] = webUsername;
+      environment['LOCUST_PASSWORD'] = webPassword;
+      environment['FLASK_SECRET_KEY'] = 'dummy'; // this is somehow required for Locust
+    }
+    if (additionalArguments != null) {
+      command.push(...additionalArguments);
     }
-
     masterTaskDefinition.addContainer('locust', {
       image,
       command,
@@ -60,6 +67,7 @@ export class LocustMasterService extends Construct {
           containerPort: 8089,
         },
       ],
+      environment,
     });
 
     const master = new ApplicationLoadBalancedFargateService(this, 'Service', {
@@ -87,8 +95,8 @@ export class LocustMasterService extends Construct {
     master.targetGroup.configureHealthCheck({
       interval: Duration.seconds(15),
       healthyThresholdCount: 2,
-      // regard 401 as healthy because we cannot use basic auth for health check
-      healthyHttpCodes: '200,401',
+      // regard 302 as healthy because Locust redirects unauthenticated requests
+      healthyHttpCodes: '200,302',
     });
 
     const port = protocol == ApplicationProtocol.HTTPS ? 443 : 80;

lib/constructs/locust_worker_service.ts

@@ -1,9 +1,9 @@
 import { Construct } from 'constructs';
 import * as ecs from 'aws-cdk-lib/aws-ecs';
 import { RetentionDays } from 'aws-cdk-lib/aws-logs';
+import { Platform } from 'aws-cdk-lib/aws-ecr-assets';
 
 export interface LocustWorkerServiceProps {
-  readonly image: ecs.ContainerImage;
   readonly cluster: ecs.ICluster;
   readonly locustMasterHostName: string;
 }
@@ -14,11 +14,14 @@ export class LocustWorkerService extends Construct {
   constructor(scope: Construct, id: string, props: LocustWorkerServiceProps) {
     super(scope, id);
 
-    const { cluster, image } = props;
+    const { cluster } = props;
+
+    const image = new ecs.AssetImage('app', { platform: Platform.LINUX_AMD64 });
 
     const workerTaskDefinition = new ecs.FargateTaskDefinition(this, 'TaskDefinition', {
-      cpu: 4096,
-      memoryLimitMiB: 8192,
+      // a locust worker can use only 1 core: https://github.com/locustio/locust/issues/1493
+      cpu: 1024,
+      memoryLimitMiB: 2048,
     });
 
     workerTaskDefinition
@@ -39,7 +42,7 @@ export class LocustWorkerService extends Construct {
       });
 
     const service = new ecs.FargateService(this, 'Service', {
-      desiredCount: 20, // set number of locust worker nodes
+      desiredCount: 2, // set number of locust worker nodes
       cluster,
       taskDefinition: workerTaskDefinition,
       // You can adjust spot:on-demand ratio here
@@ -53,6 +56,7 @@ export class LocustWorkerService extends Construct {
           weight: 0,
         },
       ],
+      minHealthyPercent: 0,
     });
 
     this.service = service;