feat(cli): put s3 bucket lifecycle policy for aws tests (#2434)

* feat(fargate): set s3 lifecycle rules for object expiration

* refactor(lambda): use central util getAccountId and ensureS3BucketExists

* feat(lambda): set s3 lifecycle rules for object expiration

Author: bernardobridge

packages/artillery/lib/platform/aws-ecs/ecs.js

@@ -27,6 +27,21 @@ class PlatformECS {
     if (!this.testRunId) {
       throw new Error('testRunId is required');
     }
+
+    this.s3LifecycleConfigurationRules = [
+      {
+        Expiration: { Days: 2 },
+        Filter: { Prefix: 'tests/' },
+        ID: 'RemoveAdHocTestData',
+        Status: 'Enabled'
+      },
+      {
+        Expiration: { Days: 7 },
+        Filter: { Prefix: 'test-runs/' },
+        ID: 'RemoveTestRunMetadata',
+        Status: 'Enabled'
+      }
+    ];
   }
 
   async init() {
@@ -35,7 +50,7 @@ class PlatformECS {
     this.accountId = await getAccountId();
 
     await ensureSSMParametersExist(this.platformOpts.region);
-    await ensureS3BucketExists();
+    await ensureS3BucketExists('global', this.s3LifecycleConfigurationRules);
     await createIAMResources(this.accountId);
   }
 

packages/artillery/lib/platform/aws-lambda/index.js

@@ -34,6 +34,8 @@ const prices = require('./prices');
 const { STATES } = require('../local/artillery-worker-local');
 
 const { SQS_QUEUES_NAME_PREFIX } = require('../aws/constants');
+const ensureS3BucketExists = require('../aws/aws-ensure-s3-bucket-exists');
+const getAccountId = require('../aws/aws-get-account-id');
 
 const createSQSQueue = require('../aws/aws-create-sqs-queue');
 
@@ -91,6 +93,20 @@ class PlatformLambda {
       platformConfig['lambda-role-arn'] || platformConfig['lambdaRoleArn'];
 
     this.platformOpts = platformOpts;
+    this.s3LifecycleConfigurationRules = [
+      {
+        Expiration: { Days: 2 },
+        Filter: { Prefix: '/lambda' },
+        ID: 'RemoveAdHocTestData',
+        Status: 'Enabled'
+      },
+      {
+        Expiration: { Days: 7 },
+        Filter: { Prefix: '/' },
+        ID: 'RemoveTestRunMetadata',
+        Status: 'Enabled'
+      }
+    ];
 
     this.artilleryArgs = [];
   }
@@ -103,7 +119,7 @@ class PlatformLambda {
     artillery.log('λ Creating AWS Lambda function...');
 
     await setDefaultAWSCredentials(AWS);
-    this.accountId = await this.getAccountId();
+    this.accountId = await getAccountId();
 
     const dirname = temp.mkdirSync(); // TODO: May want a way to override this by the user
     const zipfile = temp.path({ suffix: '.zip' });
@@ -313,7 +329,10 @@ class PlatformLambda {
     await this.createZip(dirname, zipfile);
 
     artillery.log('Preparing AWS environment...');
-    const bucketName = await this.ensureS3BucketExists();
+    const bucketName = await ensureS3BucketExists(
+      this.region,
+      this.s3LifecycleConfigurationRules
+    );
     this.bucketName = bucketName;
 
     const s3path = await this.uploadLambdaZip(bucketName, zipfile);
@@ -639,42 +658,6 @@ class PlatformLambda {
     });
   }
 
-  // TODO: Move into reusable platform util
-  async ensureS3BucketExists() {
-    const accountId = await this.getAccountId();
-    // S3 and Lambda have to be in the same region, which means we can't reuse
-    // the bucket created by Pro to store Lambda deployment zips
-    const bucketName = `artilleryio-test-data-${this.region}-${accountId}`;
-    const s3 = new AWS.S3({ region: this.region });
-
-    try {
-      await s3.listObjectsV2({ Bucket: bucketName, MaxKeys: 1 }).promise();
-    } catch (s3Err) {
-      if (s3Err.code === 'NoSuchBucket') {
-        const res = await s3.createBucket({ Bucket: bucketName }).promise();
-      } else {
-        throw s3Err;
-      }
-    }
-
-    return bucketName;
-  }
-
-  // TODO: Move into reusable platform util
-  async getAccountId() {
-    let stsOpts = {};
-    if (process.env.ARTILLERY_STS_OPTS) {
-      stsOpts = Object.assign(
-        stsOpts,
-        JSON.parse(process.env.ARTILLERY_STS_OPTS)
-      );
-    }
-
-    const sts = new AWS.STS(stsOpts);
-    const awsAccountId = (await sts.getCallerIdentity({}).promise()).Account;
-    return awsAccountId;
-  }
-
   async createLambdaRole() {
     const ROLE_NAME = 'artilleryio-default-lambda-role-20230116';
     const POLICY_NAME = 'artilleryio-lambda-policy-20230116';

packages/artillery/lib/platform/aws/aws-ensure-s3-bucket-exists.js

@@ -10,12 +10,34 @@ const getAWSAccountId = require('./aws-get-account-id');
 
 const { S3_BUCKET_NAME_PREFIX } = require('./constants');
 
+const setBucketLifecyclePolicy = async (
+  bucketName,
+  lifecycleConfigurationRules
+) => {
+  const s3 = new AWS.S3();
+  const params = {
+    Bucket: bucketName,
+    LifecycleConfiguration: {
+      Rules: lifecycleConfigurationRules
+    }
+  };
+  try {
+    await s3.putBucketLifecycleConfiguration(params).promise();
+  } catch (err) {
+    debug('Error setting lifecycle policy');
+    debug(err);
+  }
+};
+
 // Create an S3 bucket in the given region if it doesn't already exist.
 // By default, the bucket will be created without specifying a specific region.
 // Sometimes we need to use region-specific buckets, e.g. when
 // creating Lambda functions from a zip file in S3 the region of the
 // Lambda and the region of the S3 bucket must match.
-module.exports = async function ensureS3BucketExists(region = 'global') {
+module.exports = async function ensureS3BucketExists(
+  region = 'global',
+  lifecycleConfigurationRules = []
+) {
   const accountId = await getAWSAccountId();
   let bucketName = `${S3_BUCKET_NAME_PREFIX}-${accountId}`;
   if (region !== 'global') {
@@ -34,6 +56,10 @@ module.exports = async function ensureS3BucketExists(region = 'global') {
     }
   }
 
+  if (lifecycleConfigurationRules.length > 0) {
+    await setBucketLifecyclePolicy(bucketName, lifecycleConfigurationRules);
+  }
+
   debug(bucketName);
   return bucketName;
 };