feat(inbox): Queue depth lambda fully working!

Author: m-szaf

package.json

@@ -1,5 +1,3 @@
-import * as awsLambda from 'aws-lambda';
-
 import { getSecrets, getRedisConnection, getQueueDepth, publishMetrics } from './methods';
 
 export interface Result {
@@ -24,18 +22,13 @@ export const handler = async () => {
       throw new Error('REDIS_PORT environment variable does not parseInt');
     }
 
-    // parse queue names
     const queuesBlob = process.env.QUEUE_NAMES;
     if (!queuesBlob) {
       throw new Error('QUEUE_NAMES environment variable not set');
     }
     const queueNames: string[] = JSON.parse(queuesBlob);
 
-    // get username/password
     const redisSecretArn = process.env.REDIS_SECRET_ARN;
-    if (!redisSecretArn) {
-      throw new Error('REDIS_SECRET_ARN environment variable not set');
-    }
 
     const redisSecretPasswordPath = process.env.REDIS_SECRET_PASSWORD_PATH;
     if (!redisSecretPasswordPath) {
@@ -47,40 +40,58 @@ export const handler = async () => {
       throw new Error('REDIS_SECRET_USERNAME_PATH environment variable not set');
     }
 
-    const { username, password } = await getSecrets({
-      redisSecretArn,
-      redisSecretPasswordPath,
-      redisSecretUsernamePath,
-    });
+    const cwService = process.env.CW_SERVICE;
+    if (!cwService) {
+      throw new Error('CW_SERVICE environment variable not set');
+    }
 
-    // connect to redis instance, note this waits until the connection is 'ready'.
-    const redisConnection = await getRedisConnection({
+    const cwNamespace = process.env.CW_NAMESPACE;
+    if (!cwNamespace) {
+      throw new Error('CW_NAMESPACE environment variable not set');
+    }
+
+    console.debug(`Connecting to redis at ${host}:${port}`);
+
+    const connectionOptions = {
       connectionName: 'queueDepthMetricPublisher',
       host,
       port,
-      username,
-      password,
-      tls: {}, // always use tls, but be easy about it.
-    });
+    } as any;
+
+    if (redisSecretArn) {
+      const { username, password } = await getSecrets({
+        redisSecretArn,
+        redisSecretPasswordPath,
+        redisSecretUsernamePath,
+      });
+      connectionOptions.username = username;
+      connectionOptions.password = password;
+    }
 
-    // capture the values for the queues
+    // connect to redis instance, note this waits until the connection is 'ready'.
+    const redisConnection = await getRedisConnection(connectionOptions);
+
+    // capture the depth of all queues
     const results = await Promise.all(
       queueNames.map(async (queueName) => {
-        let depth = await getQueueDepth({ redisConnection, queueName });
-        if (!depth) {
-          console.error(`Error fetching queue depth for ${queueName}`);
-          depth = -1;
+        try {
+          let depth = await getQueueDepth({ redisConnection, queueName });
+          return { queueName, depth };
+        } catch (error) {
+          console.error(`Error fetching queue depth for ${queueName}`, error);
         }
-        return { queueName, depth };
+        return { queueName, depth: -1 };
       }),
     );
 
     if (results === undefined) {
       throw new Error('Error fetching ALL queue depths');
     }
 
-    // publish the metrics
-    publishMetrics(results);
+    publishMetrics(results, {
+      namespace: cwNamespace,
+      serviceName: cwService,
+    });
   } catch (error) {
     console.error('Error publishing metric data', error);
     throw error;

src/construct.handler.ts

@@ -1,4 +1,4 @@
-import { Duration } from 'aws-cdk-lib';
+import { Duration, Stack, aws_ec2 } from 'aws-cdk-lib';
 import { PolicyStatement } from 'aws-cdk-lib/aws-iam';
 import { NodejsFunction } from 'aws-cdk-lib/aws-lambda-nodejs';
 import { Runtime } from 'aws-cdk-lib/aws-lambda';
@@ -19,6 +19,11 @@ export interface RedisQueueDepthMetricPublisherProps {
    * @default 'RedisQueueDepth'
    */
   readonly cwNamespace?: string;
+  /**
+   * The CloudWatch service to publish metrics to.
+   * @default 'RedisQueueDepthMetricPublisher'
+   */
+  readonly cwService?: string;
   /**
    * Time intervals that Lambda will be triggered to publish metric in CloudWatch.
    * @default Duration.minutes(1)
@@ -54,7 +59,7 @@ export interface RedisQueueDepthMetricPublisherProps {
    * You can override these paths using {@link redisSecretPasswordPath}, and
    * {@link redisSecretUsernamePath} respectively.
    */
-  readonly redisSecretArn: string;
+  readonly redisSecretArn?: string;
   /**
    * In the best possible world, we would be using ABAC to allow decryption of SecretsManager payload.
    * If that is an option, leave this undefined.
@@ -68,13 +73,26 @@ export interface RedisQueueDepthMetricPublisherProps {
    *
    * @default 'password'
    */
-  readonly redisSecretPasswordPath: string;
+  readonly redisSecretPasswordPath?: string;
   /**
    * Override the JSON path for the username in the {@link redisSecretArn}
    *
    * @default 'username'
    */
-  readonly redisSecretUsernamePath: string;
+  readonly redisSecretUsernamePath?: string;
+  /**
+   * The regions to publish metrics to/observe metrics from.
+   * @default ['us-west-2']
+   */
+  readonly regions?: string[];
+  /**
+   * The VPC to run the Lambda in.
+   */
+  readonly vpc?: any;
+  /**
+   * The SecurityGroupId to grant the lambda to access redis clusters
+   */
+  readonly securityGroupId?: string;
 }
 
 /**
@@ -86,6 +104,7 @@ export class RedisQueueDepthMetricPublisher extends Construct {
   readonly handler: NodejsFunction;
   readonly rule: Rule;
   readonly cwNamespace: string;
+  readonly cwService: string;
 
   /**
    * Creates a new instance of RedisQueueDepthMetricPublisher.
@@ -96,8 +115,15 @@ export class RedisQueueDepthMetricPublisher extends Construct {
 
   constructor(scope: Construct, id: Namer, props: RedisQueueDepthMetricPublisherProps) {
     super(scope, id.pascal);
+
+    if (!props.redisSecretArn && (!props.vpc || !props.securityGroupId)) {
+      throw new Error('Either a secretArn or a vpc/securityGroupId must be provided');
+    }
+
     this.publishFrequency = props.publishFrequency ?? Duration.minutes(1);
     this.cwNamespace = props.cwNamespace ?? 'RedisQueueDepth';
+    this.cwService = props.cwService ?? 'RedisQueueDepthMetricPublisher';
+    this.regions = props.regions ?? ['us-west-2'];
     const myConstruct = this;
 
     // Note the name implies where we fetch the code from
@@ -108,16 +134,12 @@ export class RedisQueueDepthMetricPublisher extends Construct {
       },
       logRetention: props.cloudwatchLogsRetention ?? RetentionDays.THREE_MONTHS,
       memorySize: 512,
-      runtime: Runtime.NODEJS_LATEST, // Should be at least node20, but let's be aggressive about this.
+      runtime: Runtime.NODEJS_18_X, // Should be at least node20, but let's be aggressive about this.
       timeout: Duration.seconds(45),
+      ...(props.vpc ? { vpc: props.vpc } : {}),
     });
 
-    [
-      new PolicyStatement({
-        sid: 'fetchRedisSecret',
-        actions: ['secretsmanager:GetSecretValue'],
-        resources: [props.redisSecretArn],
-      }),
+    const policies = [
       new PolicyStatement({
         sid: 'putRedisQueueDepth',
         actions: ['cloudwatch:PutMetricData'],
@@ -128,29 +150,53 @@ export class RedisQueueDepthMetricPublisher extends Construct {
           },
         },
       }),
-    ].forEach((policy) => this.handler.addToRolePolicy(policy));
+    ];
 
+    if (props.redisSecretArn) {
+      policies.push(
+        new PolicyStatement({
+          sid: 'fetchRedisSecret',
+          actions: ['secretsmanager:GetSecretValue'],
+          resources: [props.redisSecretArn],
+        }),
+      );
+    }
     if (props.redisSecretKeyArn) {
-      this.handler.addToRolePolicy(
+      policies.push(
         new PolicyStatement({
           sid: 'decryptRedisSecret',
           actions: ['kms:Decrypt'],
           resources: [props.redisSecretKeyArn],
         }),
       );
     }
+    policies.forEach((policy) => this.handler.addToRolePolicy(policy));
 
     this.handler
       .addEnvironment('QUEUE_NAMES', JSON.stringify(props.queueNames))
       .addEnvironment('REDIS_ADDR', props.redisAddr)
       .addEnvironment('REDIS_PORT', props.redisPort ?? '6379')
-      .addEnvironment('REDIS_SECRET_ARN', props.redisSecretArn)
+      .addEnvironment('REDIS_SECRET_ARN', props.redisSecretArn! || '')
       .addEnvironment('REDIS_SECRET_PASSWORD_PATH', props.redisSecretPasswordPath ?? 'password')
-      .addEnvironment('REDIS_SECRET_USERNAME_PATH', props.redisSecretUsernamePath ?? 'username');
+      .addEnvironment('REDIS_SECRET_USERNAME_PATH', props.redisSecretUsernamePath ?? 'username')
+      .addEnvironment('CW_SERVICE', props.cwService ?? this.cwService)
+      .addEnvironment('CW_NAMESPACE', props.cwNamespace ?? this.cwNamespace);
+
+    if (props.securityGroupId) {
+      const securityGroup = aws_ec2.SecurityGroup.fromSecurityGroupId(this, 'securityGroup', props.securityGroupId);
+      this.handler.connections.allowTo(securityGroup, aws_ec2.Port.tcp(6379));
+    }
 
     this.rule = new Rule(this, 'rule', {
       schedule: Schedule.rate(this.publishFrequency),
     });
     this.rule.addTarget(new LambdaFunction(this.handler));
   }
 }
+
+export class RedisQueueDepthMetricPublisherStack extends Stack {
+  constructor(scope: Construct, id: string, props: RedisQueueDepthMetricPublisherProps) {
+    super(scope, id);
+    new RedisQueueDepthMetricPublisher(this, new Namer([id, 'lambda']), props);
+  }
+}

src/construct.ts

@@ -73,16 +73,8 @@ export interface GetQueueDepthProps {
 }
 
 export async function getQueueDepth(props: GetQueueDepthProps): Promise<number> {
-  let count = 0;
-  let cursor = '0';
-
-  do {
-    const result = await props.redisConnection.scan(cursor, 'MATCH', `${props.queueName}*`, 'COUNT', '1000');
-    cursor = result[0];
-    count += result[1].length;
-  } while (cursor !== '0');
-
-  return count;
+  console.debug(`Getting queue depth for ${props.queueName}`);
+  return props.redisConnection.llen(props.queueName);
 }
 
 export interface PublishMetricsProps {}
@@ -93,24 +85,21 @@ export interface PublishMetricsPayload {
 export async function publishMetrics(payload: PublishMetricsPayload[], options?: MetricsOptions): Promise<void> {
   const metrics = new Metrics(options);
   payload.forEach((p) => {
-    if (p.depth >= 0) {
-      metrics.addMetric(p.queueName, MetricUnits.Count, p.depth);
-    }
+    metrics.addMetric(`queue_depth:${p.queueName}`, MetricUnits.Count, p.depth);
   });
   metrics.publishStoredMetrics();
-  return;
 }
 
-async function localTest() {
-  const conn = await getRedisConnection({
-    host: 'localhost',
-  });
-  const depth = await getQueueDepth({
-    redisConnection: conn,
-    queueName: 'NOTIFQUEUE:NOTIFQUEUE',
-  });
-
-  console.log({ depth });
-}
+// async function localTest() {
+//   const conn = await getRedisConnection({
+//     host: 'localhost',
+//   });
+//   const depth = await getQueueDepth({
+//     redisConnection: conn,
+//     queueName: 'NOTIFQUEUE:NOTIFQUEUE',
+//   });
+
+//   console.log({ depth });
+// }
 
 // localTest();