Concurrency option for warming up X instances of of a lambda.

README.md

@@ -233,6 +233,7 @@ if(context.custom.source === 'serverless-plugin-warmup'){
 * **source** (default `{ "source": "serverless-plugin-warmup" }`)
 * **sourceRaw** (default `false`)
 * **tags** (default to serverless default tags)
+* **concurrency** (default `1`)
 
 ```yml
 custom:
@@ -248,6 +249,7 @@ custom:
     prewarm: true // Run WarmUp immediately after a deploymentlambda
     source: '{ "source": "my-custom-payload" }'
     sourceRaw: true // Won't JSON.stringify() the source, may be necessary for Go/AppSync deployments
+    concurrency: 2 // How many concurrent instances of a lambda to warm - can be overridden per lambda by setting warmupConcurrency in the function definition.
     tags:
       Project: foo
       Owner: bar    

src/index.js

@@ -115,7 +115,8 @@ class WarmUP {
       schedule: ['rate(5 minutes)'],
       timeout: 10,
       source: JSON.stringify({ source: 'serverless-plugin-warmup' }),
-      prewarm: false
+      prewarm: false,
+      concurrency: 1
     }
 
     /** Set global custom options */
@@ -174,6 +175,11 @@ class WarmUP {
     if (typeof this.custom.warmup.prewarm === 'boolean') {
       this.warmup.prewarm = this.custom.warmup.prewarm
     }
+
+    /** Concurrency */
+    if (typeof this.custom.warmup.concurrency === 'number') {
+      this.warmup.concurrency = this.custom.warmup.concurrency
+    }
   }
 
   /**
@@ -213,6 +219,7 @@ class WarmUP {
   createWarmer () {
     /** Get functions */
     const allFunctions = this.serverless.service.getAllFunctions()
+    let functionConcurrency = {}
 
     /** Filter functions for warm up */
     return BbPromise.filter(allFunctions, (functionName) => {
@@ -227,7 +234,15 @@ class WarmUP {
         : this.warmup.default
 
       /** Function needs to be warm */
-      return enable(functionConfig)
+      if (enable(functionConfig)) {
+        const concurrency = functionObject.hasOwnProperty('warmupConcurrency') && typeof functionObject.warmupConcurrency === 'number'
+          ? functionObject.warmupConcurrency : this.warmup.concurrency
+        functionConcurrency[functionName] = concurrency
+
+        return true
+      } else {
+        return false
+      }
     }).then((functionNames) => {
       /** Skip writing if no functions need to be warm */
       if (!functionNames.length) {
@@ -237,7 +252,7 @@ class WarmUP {
       }
 
       /** Write warm up function */
-      return this.createWarmUpFunctionArtifact(functionNames)
+      return this.createWarmUpFunctionArtifact(functionNames, functionConcurrency)
     }).then((skip) => {
       /** Add warm up function to service */
       if (skip !== true) {
@@ -250,20 +265,23 @@ class WarmUP {
    * @description Write warm up ES6 function
    *
    * @param {Array} functionNames - Function names
+   * @param (Dictionary) functionConcurrency - Mapping of concurrency level by function name
    *
    * @fulfil {} — Warm up function created
    * @reject {Error} Warm up error
    *
    * @return {Promise}
    * */
-  createWarmUpFunctionArtifact (functionNames) {
+  createWarmUpFunctionArtifact (functionNames, functionConcurrency) {
     /** Log warmup start */
     this.serverless.cli.log('WarmUP: setting ' + functionNames.length + ' lambdas to be warm')
 
     /** Get function names */
+    let fullFunctionConcurrency = {}
     functionNames = functionNames.map((functionName) => {
       const functionObject = this.serverless.service.getFunction(functionName)
       this.serverless.cli.log('WarmUP: ' + functionObject.name)
+      fullFunctionConcurrency[functionObject.name] = functionConcurrency[functionName]
       return functionObject.name
     })
 
@@ -274,26 +292,48 @@ const aws = require("aws-sdk");
 aws.config.region = "${this.options.region}";
 const lambda = new aws.Lambda();
 const functionNames = ${JSON.stringify(functionNames)};
+const functionConcurrency = ${JSON.stringify(fullFunctionConcurrency)};
 module.exports.warmUp = async (event, context, callback) => {
   console.log("Warm Up Start");
   const invokes = await Promise.all(functionNames.map(async (functionName) => {
-    const params = {
-      ClientContext: "${Buffer.from(`{"custom":${this.warmup.source}}`).toString('base64')}",
-      FunctionName: functionName,
-      InvocationType: "RequestResponse",
-      LogType: "None",
-      Qualifier: process.env.SERVERLESS_ALIAS || "$LATEST",
-      Payload: '${this.warmup.source}'
-    };
-
-    try {
-      const data = await lambda.invoke(params).promise();
-      console.log(\`Warm Up Invoke Success: \${functionName}\`, data);
-      return true;
-    } catch (e) {
-      console.log(\`Warm Up Invoke Error: \${functionName}\`, e);
-      return false;
+    let concurrency = functionConcurrency[functionName] > 0 ? functionConcurrency[functionName] : 1;
+    let source = ${this.warmup.source};
+    let promises = [];
+
+    console.log(\`Warming up function: \${functionName} with concurrency: \${concurrency}\`);
+
+    for (let x = 0; x < concurrency; x++) {
+      source.concurrencyIndex = x;
+      const params = {
+        ClientContext: Buffer.from(JSON.stringify({"custom":source})).toString('base64'),
+        FunctionName: functionName,
+        InvocationType: "RequestResponse",
+        LogType: "None",
+        Qualifier: process.env.SERVERLESS_ALIAS || "$LATEST",
+        Payload: JSON.stringify(source)
+      };
+
+      promises.push(lambda.invoke(params).promise());
     }
+
+    let success = true;
+    await Promise.all(promises).then(function(data) {
+      if (data && data.length > 0 && data[0].Payload) {
+        let payload = JSON.parse(data[0].Payload);
+
+        if (payload && payload.statusCode && payload.statusCode !== 200) {
+          console.log(\`Warm Up Invoke Success: \${functionName}\`, data);
+          success = false;
+          return;
+        }
+      }
+
+      console.log(\`Warm Up Invoke Success: \${functionName}\`, data);
+    }, function(err) {
+      console.log(\`Warm Up Invoke Error: \${functionName}\`, err);
+      success = false;
+    });
+    return success;
   }));
 
   console.log(\`Warm Up Finished with \${invokes.filter(r => !r).length} invoke errors\`);