Lambda warmer code

Author: devksingh4

modules/lambda-warmer/lambda/lambda.mjs

@@ -0,0 +1,8 @@
+import path from 'path';
+      import { fileURLToPath } from 'url';
+      import { createRequire as topLevelCreateRequire } from 'module';
+      const require = topLevelCreateRequire(import.meta.url);
+      const __filename = fileURLToPath(import.meta.url);
+      const __dirname = path.dirname(__filename);
+import{LambdaClient as f,InvokeCommand as v}from"@aws-sdk/client-lambda";import{TextDecoder as u}from"util";var g=new f({}),p=new u;async function I(d,a){if(a<=0)return new Set;console.log(`Firing a batch of ${a} concurrent invocations...`);let c=Array.from({length:a},()=>{let e=new v({FunctionName:d,Payload:JSON.stringify({action:"warmer"})});return g.send(e)}),i=await Promise.allSettled(c),n=new Set;return i.forEach(e=>{if(e.status==="fulfilled"&&e.value.Payload)try{let o=p.decode(e.value.Payload),t=JSON.parse(o);t.instanceId&&n.add(t.instanceId)}catch(o){console.error("Error parsing payload from target function:",o)}else e.status==="rejected"&&console.error("Invocation failed:",e.reason.message)}),n}var S=async d=>{let{lambdaName:a,numInstancesStr:c,maxWavesStr:i}={lambdaName:process.env.LAMBDA_NAME,numInstancesStr:process.env.NUM_INSTANCES,maxWavesStr:process.env.MAX_WAVES};if(!a||!c)throw new Error("Env vars 'LAMBDA_NAME' and 'NUM_INSTANCES' are required.");let n=parseInt(c,10),e=parseInt(i||"5",10),o=0,t=0,s=new Set;for(let r=1;r<=e;r++){if(t=r,n-s.size<=0){console.log("Target met. No more waves needed.");break}console.log(`--- Wave ${r} of ${e} ---`);let l=await I(a,n);o+=n,l.forEach(m=>s.add(m)),console.log(`Wave ${r} complete. Found ${s.size} of ${n} unique instances.`)}return console.log(`Warming complete. Found ${s.size} unique instances from ${o} total invocations over ${t} waves.`),{statusCode:200,body:JSON.stringify({targetInstances:n,warmedInstances:s.size,totalInvocations:o,wavesCompleted:t,instanceIds:[...s]})}};import.meta.url===`file://${process.argv[1]}`&&(process.env.LAMBDA_NAME="infra-core-api-lambda",process.env.NUM_INSTANCES="3",process.env.MAX_WAVES="3",console.log(await S({})));export{S as handler};
+//# sourceMappingURL=lambda.mjs.map

modules/lambda-warmer/lambda/lambda.mjs.map

@@ -0,0 +1,7 @@
+{
+  "version": 3,
+  "sources": ["../../src/api/warmer/lambda.ts"],
+  "sourcesContent": ["import { LambdaClient, InvokeCommand } from \"@aws-sdk/client-lambda\";\nimport { TextDecoder } from \"util\";\n\nconst lambdaClient = new LambdaClient({});\nconst textDecoder = new TextDecoder();\n\n/**\n * Invokes a batch of lambdas concurrently and returns the unique instance IDs found.\n */\nasync function invokeBatch(\n  lambdaName: string,\n  count: number,\n): Promise<Set<string>> {\n  if (count <= 0) {\n    return new Set();\n  }\n\n  console.log(`Firing a batch of ${count} concurrent invocations...`);\n\n  const invocationPromises = Array.from({ length: count }, () => {\n    const command = new InvokeCommand({\n      FunctionName: lambdaName,\n      Payload: JSON.stringify({ action: \"warmer\" }),\n    });\n    return lambdaClient.send(command);\n  });\n\n  const results = await Promise.allSettled(invocationPromises);\n  const foundInstanceIds = new Set<string>();\n\n  results.forEach((result) => {\n    if (result.status === \"fulfilled\" && result.value.Payload) {\n      try {\n        const payloadString = textDecoder.decode(result.value.Payload);\n        const body = JSON.parse(payloadString);\n        if (body.instanceId) {\n          foundInstanceIds.add(body.instanceId);\n        }\n      } catch (e) {\n        console.error(\"Error parsing payload from target function:\", e);\n      }\n    } else if (result.status === \"rejected\") {\n      console.error(\"Invocation failed:\", result.reason.message);\n    }\n  });\n\n  return foundInstanceIds;\n}\n\nexport const handler = async (event: {}) => {\n  const { lambdaName, numInstancesStr, maxWavesStr } = {\n    lambdaName: process.env.LAMBDA_NAME,\n    numInstancesStr: process.env.NUM_INSTANCES,\n    maxWavesStr: process.env.MAX_WAVES,\n  };\n\n  if (!lambdaName || !numInstancesStr) {\n    throw new Error(\"Env vars 'LAMBDA_NAME' and 'NUM_INSTANCES' are required.\");\n  }\n\n  const numInstances = parseInt(numInstancesStr, 10);\n  // Default to 5 waves if MAX_WAVES is not set\n  const maxWaves = parseInt(maxWavesStr || \"5\", 10);\n\n  let totalInvocations = 0;\n  let wavesCompleted = 0;\n  const uniqueInstanceIds = new Set<string>();\n\n  for (let i = 1; i <= maxWaves; i++) {\n    wavesCompleted = i;\n\n    // Calculate how many more instances are needed\n    const neededCount = numInstances - uniqueInstanceIds.size;\n    if (neededCount <= 0) {\n      console.log(\"Target met. No more waves needed.\");\n      break;\n    }\n\n    console.log(`--- Wave ${i} of ${maxWaves} ---`);\n    const newIds = await invokeBatch(lambdaName, numInstances);\n    totalInvocations += numInstances;\n\n    newIds.forEach((id) => uniqueInstanceIds.add(id));\n\n    console.log(\n      `Wave ${i} complete. Found ${uniqueInstanceIds.size} of ${numInstances} unique instances.`,\n    );\n  }\n\n  console.log(\n    `Warming complete. Found ${uniqueInstanceIds.size} unique instances from ${totalInvocations} total invocations over ${wavesCompleted} waves.`,\n  );\n\n  return {\n    statusCode: 200,\n    body: JSON.stringify({\n      targetInstances: numInstances,\n      warmedInstances: uniqueInstanceIds.size,\n      totalInvocations,\n      wavesCompleted,\n      instanceIds: [...uniqueInstanceIds],\n    }),\n  };\n};\n\nif (import.meta.url === `file://${process.argv[1]}`) {\n  process.env.LAMBDA_NAME = \"infra-core-api-lambda\";\n  process.env.NUM_INSTANCES = \"3\";\n  process.env.MAX_WAVES = \"3\"; // Configurable number of waves\n  console.log(await handler({}));\n}\n"],
+  "mappings": ";;;;;;AAAA,OAAS,gBAAAA,EAAc,iBAAAC,MAAqB,yBAC5C,OAAS,eAAAC,MAAmB,OAE5B,IAAMC,EAAe,IAAIH,EAAa,CAAC,CAAC,EAClCI,EAAc,IAAIF,EAKxB,eAAeG,EACbC,EACAC,EACsB,CACtB,GAAIA,GAAS,EACX,OAAO,IAAI,IAGb,QAAQ,IAAI,qBAAqBA,CAAK,4BAA4B,EAElE,IAAMC,EAAqB,MAAM,KAAK,CAAE,OAAQD,CAAM,EAAG,IAAM,CAC7D,IAAME,EAAU,IAAIR,EAAc,CAChC,aAAcK,EACd,QAAS,KAAK,UAAU,CAAE,OAAQ,QAAS,CAAC,CAC9C,CAAC,EACD,OAAOH,EAAa,KAAKM,CAAO,CAClC,CAAC,EAEKC,EAAU,MAAM,QAAQ,WAAWF,CAAkB,EACrDG,EAAmB,IAAI,IAE7B,OAAAD,EAAQ,QAASE,GAAW,CAC1B,GAAIA,EAAO,SAAW,aAAeA,EAAO,MAAM,QAChD,GAAI,CACF,IAAMC,EAAgBT,EAAY,OAAOQ,EAAO,MAAM,OAAO,EACvDE,EAAO,KAAK,MAAMD,CAAa,EACjCC,EAAK,YACPH,EAAiB,IAAIG,EAAK,UAAU,CAExC,OAASC,EAAG,CACV,QAAQ,MAAM,8CAA+CA,CAAC,CAChE,MACSH,EAAO,SAAW,YAC3B,QAAQ,MAAM,qBAAsBA,EAAO,OAAO,OAAO,CAE7D,CAAC,EAEMD,CACT,CAEO,IAAMK,EAAU,MAAOC,GAAc,CAC1C,GAAM,CAAE,WAAAX,EAAY,gBAAAY,EAAiB,YAAAC,CAAY,EAAI,CACnD,WAAY,QAAQ,IAAI,YACxB,gBAAiB,QAAQ,IAAI,cAC7B,YAAa,QAAQ,IAAI,SAC3B,EAEA,GAAI,CAACb,GAAc,CAACY,EAClB,MAAM,IAAI,MAAM,0DAA0D,EAG5E,IAAME,EAAe,SAASF,EAAiB,EAAE,EAE3CG,EAAW,SAASF,GAAe,IAAK,EAAE,EAE5CG,EAAmB,EACnBC,EAAiB,EACfC,EAAoB,IAAI,IAE9B,QAASC,EAAI,EAAGA,GAAKJ,EAAUI,IAAK,CAKlC,GAJAF,EAAiBE,EAGGL,EAAeI,EAAkB,MAClC,EAAG,CACpB,QAAQ,IAAI,mCAAmC,EAC/C,KACF,CAEA,QAAQ,IAAI,YAAYC,CAAC,OAAOJ,CAAQ,MAAM,EAC9C,IAAMK,EAAS,MAAMrB,EAAYC,EAAYc,CAAY,EACzDE,GAAoBF,EAEpBM,EAAO,QAASC,GAAOH,EAAkB,IAAIG,CAAE,CAAC,EAEhD,QAAQ,IACN,QAAQF,CAAC,oBAAoBD,EAAkB,IAAI,OAAOJ,CAAY,oBACxE,CACF,CAEA,eAAQ,IACN,2BAA2BI,EAAkB,IAAI,0BAA0BF,CAAgB,2BAA2BC,CAAc,SACtI,EAEO,CACL,WAAY,IACZ,KAAM,KAAK,UAAU,CACnB,gBAAiBH,EACjB,gBAAiBI,EAAkB,KACnC,iBAAAF,EACA,eAAAC,EACA,YAAa,CAAC,GAAGC,CAAiB,CACpC,CAAC,CACH,CACF,EAEI,YAAY,MAAQ,UAAU,QAAQ,KAAK,CAAC,CAAC,KAC/C,QAAQ,IAAI,YAAc,wBAC1B,QAAQ,IAAI,cAAgB,IAC5B,QAAQ,IAAI,UAAY,IACxB,QAAQ,IAAI,MAAMR,EAAQ,CAAC,CAAC,CAAC",
+  "names": ["LambdaClient", "InvokeCommand", "TextDecoder", "lambdaClient", "textDecoder", "invokeBatch", "lambdaName", "count", "invocationPromises", "command", "results", "foundInstanceIds", "result", "payloadString", "body", "e", "handler", "event", "numInstancesStr", "maxWavesStr", "numInstances", "maxWaves", "totalInvocations", "wavesCompleted", "uniqueInstanceIds", "i", "newIds", "id"]
+}

modules/lambda-warmer/main.tf

@@ -0,0 +1,109 @@
+terraform {
+  required_providers {
+    aws = {
+      source = "hashicorp/aws"
+    }
+  }
+}
+
+resource "aws_cloudwatch_log_group" "warmer_logs" {
+  name              = "/aws/lambda/${var.function_to_warm}-warmer"
+  retention_in_days = var.log_retention_days
+}
+
+resource "aws_iam_role" "warmer_iam_role" {
+  assume_role_policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Action = "sts:AssumeRole"
+        Effect = "Allow"
+        Sid    = ""
+        Principal = {
+          Service = "lambda.amazonaws.com"
+        }
+      },
+    ]
+  })
+}
+
+data "archive_file" "warmer_code" {
+  type        = "zip"
+  source_dir  = "${path.module}/lambda"
+  output_path = "${path.module}/lambda/function.zip"
+}
+
+resource "aws_lambda_function" "warmer_function" {
+  depends_on       = [aws_cloudwatch_log_group.warmer_logs]
+  architectures    = ["arm64"]
+  filename         = data.archive_file.warmer_code.output_path
+  runtime          = "nodejs22.x"
+  source_code_hash = data.archive_file.warmer_code.output_base64sha256
+  role             = aws_iam_role.warmer_iam_role
+  function_name    = "${var.function_to_warm}-warmer"
+  memory_size      = 256
+  timeout          = 15
+  environment {
+    variables = {
+      LAMBDA_NAME   = var.function_to_warm
+      NUM_INSTANCES = var.num_desired_warm_instances
+    }
+  }
+}
+
+resource "aws_lambda_alias" "warmer_function_alias" {
+  name             = "live"
+  description      = "Live environment alias"
+  function_name    = aws_lambda_function.warmer_function.arn
+  function_version = aws_lambda_function.warmer_function.version
+}
+
+resource "aws_iam_role_policy" "warmer_log_write_policy" {
+  name = "warmer_log_write_policy"
+  role = aws_iam_role.warmer_iam_role
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Action = [
+          "logs:CreateLogGroup",
+          "logs:CreateLogStream",
+          "logs:PutLogEvents"
+        ]
+        Effect   = "Allow"
+        Resource = aws_cloudwatch_log_group.warmer_logs.arn
+      },
+    ]
+  })
+}
+
+
+resource "aws_iam_role_policy" "warmer_lambda_invoke_policy" {
+  name = "warmer_lambda_invoke_policy"
+  role = aws_iam_role.warmer_iam_role
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Action = [
+          "lambda:InvokeFunction"
+        ]
+        Effect   = "Allow"
+        Resource = "arn:aws:lambda:${data.aws_region.current.name}:${data.aws_caller_identity.current.account_id}:function:${var.function_to_warm}"
+      },
+    ]
+  })
+}
+
+resource "aws_cloudwatch_event_rule" "warmer_schedule" {
+  description         = "Schedule to run warmer for ${var.function_to_warm}"
+  schedule_expression = var.invoke_rate_string
+  state               = "ENABLED"
+}
+
+resource "aws_lambda_permission" "warmer_lambda_permission" {
+  function_name = "${var.function_to_warm}-warmer"
+  action        = "lambda:InvokeFunction"
+  principal     = "events.amazonaws.com"
+  source_arn    = aws_cloudwatch_event_rule.warmer_schedule.arn
+}

modules/lambda-warmer/variables.tf

@@ -0,0 +1,23 @@
+variable "function_to_warm" {
+  type        = string
+  description = "Name of the lambda function to warm"
+}
+
+variable "log_retention_days" {
+  type        = number
+  description = "Number of days to retain lambda warmer logs."
+  default     = 7
+}
+
+variable "invoke_rate_string" {
+  type        = string
+  description = "EventBridge rate string for how often to call the lambda."
+  default     = "rate(4 minutes)"
+}
+
+
+variable "num_desired_warm_instances" {
+  type        = number
+  description = "Number of warm lambda instances desired"
+  default     = 3
+}