feat: cleanup idle stages by project tag (#35)

Author: drewhan90

cdk/Makefile

@@ -10,7 +10,7 @@ PUBLISH					 ?=	false
 STACK 					 ?=	UGC-$(STAGE)
 COGNITO_CLEANUP_SCHEDULE ?= rate(48 hours)
 STAGE_CLEANUP_SCHEDULE   ?= rate(24 hours)
-CDK_OPTIONS 			 =	$(if $(AWS_PROFILE),$(AWS_PROFILE_FLAG)) -c stage=$(STAGE) -c publish=$(PUBLISH) -c stackName=$(STACK) -c cognitoCleanupScheduleExp="$(strip $(COGNITO_CLEANUP_SCHEDULE))"
+CDK_OPTIONS 			 =	$(if $(AWS_PROFILE),$(AWS_PROFILE_FLAG)) -c stage=$(STAGE) -c publish=$(PUBLISH) -c stackName=$(STACK) -c cognitoCleanupScheduleExp="$(strip $(COGNITO_CLEANUP_SCHEDULE))" -c stageCleanupScheduleExp="$(strip $(STAGE_CLEANUP_SCHEDULE))"
 FE_DEPLOYMENT_STACK		 =	UGC-Frontend-Deployment-$(STAGE)
 SEED_COUNT 				 ?=	50
 OFFLINE_SESSION_COUNT	 ?=	1

cdk/api/stages/authRouter/createStage.ts

@@ -64,7 +64,8 @@ const handler = async (request: FastifyRequest, reply: FastifyReply) => {
       ],
       tags: {
         creationDate: stageCreationDate,
-        stageOwnerChannelId: channelId
+        stageOwnerChannelId: channelId,
+        stack: process.env.STACK as string
       }
     };
 

cdk/bin/cdk.ts

@@ -12,9 +12,8 @@ const app = new App();
 const stage = app.node.tryGetContext('stage');
 const stackName = app.node.tryGetContext('stackName');
 const shouldPublish = app.node.tryGetContext('publish') === 'true';
-const cognitoCleanupScheduleExp = app.node.tryGetContext(
-  'cognitoCleanupScheduleExp'
-);
+const cognitoCleanupScheduleExp = app.node.tryGetContext('cognitoCleanupScheduleExp');
+const stageCleanupScheduleExp = app.node.tryGetContext('stageCleanupScheduleExp')
 // Get the config for the current stage
 const { resourceConfig }: { resourceConfig: UGCResourceWithChannelsConfig } =
   app.node.tryGetContext(stage);
@@ -26,7 +25,8 @@ new UGCStack(app, stackName, {
   tags: { stage, project: 'ugc' },
   resourceConfig,
   shouldPublish,
-  cognitoCleanupScheduleExp
+  cognitoCleanupScheduleExp,
+  stageCleanupScheduleExp
 });
 
 new UGCFrontendDeploymentStack(app, `UGC-Frontend-Deployment-${stage}`, {

cdk/lambdas/cleanupIdleStages.ts

@@ -0,0 +1,50 @@
+import { ListStagesCommand } from '@aws-sdk/client-ivs-realtime';
+
+import {
+  ivsRealTimeClient,
+  getIdleStageArns,
+  deleteStagesWithRetry,
+  updateMultipleChannelDynamoItems,
+  getBatchChannelWriteUpdates,
+  getIdleStages
+} from './helpers';
+
+export const handler = async () => {
+  try {
+    const deleteIdleStages = async (nextToken = '') => {
+      const listStagesCommand = new ListStagesCommand({
+        maxResults: 100,
+        nextToken
+      });
+      const response = await ivsRealTimeClient.send(listStagesCommand);
+
+      const stages = response?.stages || [];
+      const _nextToken = response?.nextToken || '';
+
+      if (stages.length) {
+        // Filter list of stages by stack name
+        const projectStages = stages.filter(
+          ({ tags }) => !!tags?.stack && tags.stack === process.env.STACK_TAG
+        );
+        const idleStages = getIdleStages(projectStages);
+        const idleStageArns = getIdleStageArns(idleStages);
+        const batchChannelWriteUpdates =
+          getBatchChannelWriteUpdates(idleStages);
+        await Promise.all([
+          deleteStagesWithRetry(idleStageArns),
+          updateMultipleChannelDynamoItems(batchChannelWriteUpdates)
+        ]);
+      }
+
+      if (_nextToken) await deleteIdleStages(_nextToken);
+    };
+
+    await deleteIdleStages();
+  } catch (error) {
+    console.error(error);
+
+    throw new Error('Failed to remove idle stages due to unexpected error');
+  }
+};
+
+export default handler;

cdk/lambdas/helpers.ts

@@ -87,7 +87,7 @@ export const batchDeleteItemsWithRetry = async (
  * Cleanup Idle Stages
  */
 
-export const getIdleStages = (stages: StageSummary[]) => {
+export const getIdleStages = (stages: StageSummary[] = []) => {
   const currentTimestamp = Date.now();
   const millisecondsPerHour = 60 * 60 * 1000;
   const hoursThreshold = 1;
@@ -105,7 +105,9 @@ export const getIdleStages = (stages: StageSummary[]) => {
     });
 };
 
-export const getBatchChannelWriteUpdates = (idleStages: StageSummary[]) => {
+export const getBatchChannelWriteUpdates = (
+  idleStages: StageSummary[] = []
+) => {
   const channelWriteUpdates: WriteRequest[] = [];
 
   idleStages.forEach((idleAndOldStage) => {
@@ -132,7 +134,7 @@ export const getBatchChannelWriteUpdates = (idleStages: StageSummary[]) => {
   return channelWriteUpdates;
 };
 
-export const getIdleStageArns = (idleStages: StageSummary[]) =>
+export const getIdleStageArns = (idleStages: StageSummary[] = []) =>
   idleStages
     .map((idleAndOldStage) => idleAndOldStage.arn)
     .filter((arn) => typeof arn === 'string') as string[];
@@ -209,7 +211,7 @@ const analyzeDeleteStageResponse = (
   };
 };
 
-export const deleteStagesWithRetry = async (stageArns: string[]) => {
+export const deleteStagesWithRetry = async (stageArns: string[] = []) => {
   if (!stageArns.length) return;
 
   const stagesToDelete = chunkIntoArrayBatches(stageArns, 5);
@@ -282,8 +284,10 @@ export const updateDynamoItemAttributes = ({
 };
 
 export const updateMultipleChannelDynamoItems = (
-  idleStagesChannelArns: WriteRequest[]
+  idleStagesChannelArns: WriteRequest[] = []
 ) => {
+  if (!idleStagesChannelArns.length) return;
+
   const batchWriteInput: BatchWriteItemInput = {
     RequestItems: {
       [process.env.CHANNELS_TABLE_NAME as string]: idleStagesChannelArns

cdk/lib/ChannelsStack/cdk-channels-stack.ts

@@ -7,9 +7,11 @@ import {
   aws_events as events,
   aws_events_targets as targets,
   aws_iam as iam,
+  aws_lambda as lambda,
   aws_lambda_nodejs as nodejsLambda,
   aws_s3 as s3,
   aws_s3_notifications as s3n,
+  aws_sqs as sqs,
   Duration,
   NestedStack,
   NestedStackProps,
@@ -41,6 +43,7 @@ interface ChannelsStackProps extends NestedStackProps {
   resourceConfig: ChannelsResourceConfig;
   tags: { [key: string]: string };
   cognitoCleanupScheduleExp: string;
+  stageCleanupScheduleExp: string;
 }
 
 export class ChannelsStack extends NestedStack {
@@ -67,7 +70,12 @@ export class ChannelsStack extends NestedStack {
     const region = Stack.of(this.nestedStackParent!).region;
     const nestedStackName = 'Channels';
     const stackNamePrefix = `${parentStackName}-${nestedStackName}`;
-    const { resourceConfig, cognitoCleanupScheduleExp, tags } = props;
+    const {
+      resourceConfig,
+      cognitoCleanupScheduleExp,
+      stageCleanupScheduleExp,
+      tags
+    } = props;
 
     // Configuration variables based on the stage (dev or prod)
     const {
@@ -471,7 +479,7 @@ export class ChannelsStack extends NestedStack {
 
     // Cleanup idle stages users policies
     const deleteIdleStagesIvsPolicyStatement = new iam.PolicyStatement({
-      actions: ['ivs:ListStages', 'ivs:DeleteStage'],
+      actions: ['ivs:ListStages', 'ivs:DeleteStage', 'dynamodb:BatchWriteItem'],
       effect: iam.Effect.ALLOW,
       resources: ['*']
     });
@@ -488,6 +496,23 @@ export class ChannelsStack extends NestedStack {
       resources: [userPool.userPoolArn]
     });
 
+    // Cleanup idle stages lambda
+    const cleanupIdleStagesHandler = new nodejsLambda.NodejsFunction(
+      this,
+      `${stackNamePrefix}-CleanupIdleStages-Handler`,
+      {
+        ...defaultLambdaParams,
+        logRetention: RetentionDays.ONE_WEEK,
+        functionName: `${stackNamePrefix}-CleanupIdleStages`,
+        entry: getLambdaEntryPath('cleanupIdleStages'),
+        timeout: Duration.minutes(10),
+        initialPolicy: [deleteIdleStagesIvsPolicyStatement],
+        environment: {
+          STACK_TAG: parentStackName
+        }
+      }
+    );
+
     // Cleanup unverified users lambda
     const cleanupUnverifiedUsersHandler = new nodejsLambda.NodejsFunction(
       this,
@@ -505,6 +530,18 @@ export class ChannelsStack extends NestedStack {
       }
     );
 
+    // Scheduled cleanup idle stages lambda function
+    new events.Rule(this, 'Cleanup-Idle-Stages-Schedule-Rule', {
+      schedule: events.Schedule.expression(stageCleanupScheduleExp),
+      ruleName: `${stackNamePrefix}-CleanupIdleStages-Schedule`,
+      targets: [
+        new targets.LambdaFunction(cleanupIdleStagesHandler, {
+          maxEventAge: Duration.minutes(2),
+          retryAttempts: 2
+        })
+      ]
+    });
+
     // Scheduled cleanup unverified users lambda function
     new events.Rule(this, 'Cleanup-Unverified-Users-Schedule-Rule', {
       schedule: events.Schedule.expression(cognitoCleanupScheduleExp),

cdk/lib/cdk-ugc-stack.ts

@@ -28,6 +28,7 @@ const DEFAULT_CLIENT_BASE_URLS = ['', 'http://localhost:3000'];
 interface UGCDashboardStackProps extends StackProps {
   resourceConfig: UGCResourceWithChannelsConfig;
   cognitoCleanupScheduleExp: string;
+  stageCleanupScheduleExp: string;
   shouldPublish: boolean;
 }
 
@@ -38,6 +39,7 @@ export class UGCStack extends Stack {
     const {
       resourceConfig,
       cognitoCleanupScheduleExp,
+      stageCleanupScheduleExp,
       shouldPublish,
       tags = {}
     } = props;
@@ -141,6 +143,7 @@ export class UGCStack extends Stack {
     const channelsStack = new ChannelsStack(this, 'Channels', {
       resourceConfig,
       cognitoCleanupScheduleExp,
+      stageCleanupScheduleExp,
       tags
     });
     const {
@@ -203,7 +206,8 @@ export class UGCStack extends Stack {
       PRODUCT_LINK_REGION_CODE: productLinkRegionCode,
       ENABLE_AMAZON_PRODUCT_STREAM_ACTION: `${enableAmazonProductStreamAction}`,
       PRODUCT_API_SECRET_NAME: productApiSecretName,
-      APPSYNC_GRAPHQL_API_SECRET_NAME: appSyncGraphQlApi.secretName
+      APPSYNC_GRAPHQL_API_SECRET_NAME: appSyncGraphQlApi.secretName,
+      STACK: stackNamePrefix
     };
     const sharedContainerEnv = {
       ...baseContainerEnv,