Add parameter to power-tune only cold starts

lambda/cleaner.js

@@ -7,13 +7,15 @@ const utils = require('./utils');
  */
 module.exports.handler = async(event, context) => {
 
-    const {lambdaARN, powerValues} = event;
+    const {
+        lambdaARN,
+        aliases,
+    } = extractDataFromInput(event);
 
-    validateInput(lambdaARN, powerValues); // may throw
+    validateInput(lambdaARN, aliases); // may throw
 
-    const ops = powerValues.map(async(value) => {
-        const alias = 'RAM' + value;
-        await cleanup(lambdaARN, alias); // may throw
+    const ops = aliases.map(async(alias) => {
+        await cleanup(lambdaARN, alias);
     });
 
     // run everything in parallel and wait until completed
@@ -22,12 +24,19 @@ module.exports.handler = async(event, context) => {
     return 'OK';
 };
 
-const validateInput = (lambdaARN, powerValues) => {
+const extractDataFromInput = (event) => {
+    return {
+        lambdaARN: event.lambdaARN,
+        aliases: event.powerValues.aliases,
+    };
+};
+
+const validateInput = (lambdaARN, aliases) => {
     if (!lambdaARN) {
         throw new Error('Missing or empty lambdaARN');
     }
-    if (!powerValues || !powerValues.length) {
-        throw new Error('Missing or empty power values');
+    if (!aliases || !aliases.length) {
+        throw new Error('Missing or empty alias values');
     }
 };
 

lambda/executor.js

@@ -26,6 +26,7 @@ module.exports.handler = async(event, context) => {
         preProcessorARN,
         postProcessorARN,
         discardTopBottom,
+        onlyColdStarts,
         sleepBetweenRunsMs,
     } = await extractDataFromInput(event);
 
@@ -40,8 +41,9 @@ module.exports.handler = async(event, context) => {
     const lambdaAlias = 'RAM' + value;
     let results;
 
-    // fetch architecture from $LATEST
-    const {architecture, isPending} = await utils.getLambdaConfig(lambdaARN, lambdaAlias);
+    let aliasToInvoke = utils.buildAliasString(lambdaAlias, onlyColdStarts, 0);
+    const {architecture, isPending} = await utils.getLambdaConfig(lambdaARN, aliasToInvoke);
+
     console.log(`Detected architecture type: ${architecture}, isPending: ${isPending}`);
 
     // pre-generate an array of N payloads
@@ -54,11 +56,12 @@ module.exports.handler = async(event, context) => {
         payloads: payloads,
         preARN: preProcessorARN,
         postARN: postProcessorARN,
+        onlyColdStarts: onlyColdStarts,
         sleepBetweenRunsMs: sleepBetweenRunsMs,
     };
 
     // wait if the function/alias state is Pending
-    if (isPending) {
+    if (isPending && !onlyColdStarts) {
         await utils.waitForAliasActive(lambdaARN, lambdaAlias);
         console.log('Alias active');
     }
@@ -101,7 +104,11 @@ const extractDiscardTopBottomValue = (event) => {
     // extract discardTopBottom used to trim values from average duration
     let discardTopBottom = event.discardTopBottom;
     if (typeof discardTopBottom === 'undefined') {
-        discardTopBottom = 0.2;
+        if (event.onlyColdStarts){
+            discardTopBottom = 0;
+        } else {
+            discardTopBottom = 0.2;
+        }
     }
     // discardTopBottom must be between 0 and 0.4
     return Math.min(Math.max(discardTopBottom, 0.0), 0.4);
@@ -132,15 +139,21 @@ const extractDataFromInput = async(event) => {
         preProcessorARN: input.preProcessorARN,
         postProcessorARN: input.postProcessorARN,
         discardTopBottom: discardTopBottom,
+        onlyColdStarts: !!input.onlyColdStarts,
         sleepBetweenRunsMs: sleepBetweenRunsMs,
     };
 };
 
-const runInParallel = async({num, lambdaARN, lambdaAlias, payloads, preARN, postARN}) => {
+const runInParallel = async({num, lambdaARN, lambdaAlias, payloads, preARN, postARN, onlyColdStarts}) => {
     const results = [];
     // run all invocations in parallel ...
     const invocations = utils.range(num).map(async(_, i) => {
-        const {invocationResults, actualPayload} = await utils.invokeLambdaWithProcessors(lambdaARN, lambdaAlias, payloads[i], preARN, postARN);
+        let aliasToInvoke = utils.buildAliasString(lambdaAlias, onlyColdStarts, i);
+        if (onlyColdStarts){
+            await utils.waitForAliasActive(lambdaARN, aliasToInvoke);
+            console.log(`${aliasToInvoke} is active`);
+        }
+        const {invocationResults, actualPayload} = await utils.invokeLambdaWithProcessors(lambdaARN, aliasToInvoke, payloads[i], preARN, postARN);
         // invocation errors return 200 and contain FunctionError and Payload
         if (invocationResults.FunctionError) {
             throw new Error(`Invocation error (running in parallel): ${invocationResults.Payload} with payload ${JSON.stringify(actualPayload)}`);
@@ -152,11 +165,15 @@ const runInParallel = async({num, lambdaARN, lambdaAlias, payloads, preARN, post
     return results;
 };
 
-const runInSeries = async({num, lambdaARN, lambdaAlias, payloads, preARN, postARN, sleepBetweenRunsMs}) => {
+const runInSeries = async({num, lambdaARN, lambdaAlias, payloads, preARN, postARN, onlyColdStarts, sleepBetweenRunsMs }) => {
     const results = [];
     for (let i = 0; i < num; i++) {
+        let aliasToInvoke = utils.buildAliasString(lambdaAlias, onlyColdStarts, i);
         // run invocations in series
-        const {invocationResults, actualPayload} = await utils.invokeLambdaWithProcessors(lambdaARN, lambdaAlias, payloads[i], preARN, postARN);
+        if (onlyColdStarts){
+            await utils.waitForAliasActive(lambdaARN, aliasToInvoke);
+        }
+        const {invocationResults, actualPayload} = await utils.invokeLambdaWithProcessors(lambdaARN, aliasToInvoke, payloads[i], preARN, postARN);
         // invocation errors return 200 and contain FunctionError and Payload
         if (invocationResults.FunctionError) {
             throw new Error(`Invocation error (running in series): ${invocationResults.Payload} with payload ${JSON.stringify(actualPayload)}`);

lambda/initializer.js

@@ -8,24 +8,58 @@ const defaultPowerValues = process.env.defaultPowerValues.split(',');
  */
 module.exports.handler = async(event, context) => {
 
-    const {lambdaARN, num} = event;
-    const powerValues = extractPowerValues(event);
+    const {
+        lambdaARN,
+        num,
+        powerValues,
+        onlyColdStarts,
+    } = extractDataFromInput(event);
 
     validateInput(lambdaARN, num); // may throw
 
     // fetch initial $LATEST value so we can reset it later
-    const initialPower = await utils.getLambdaPower(lambdaARN);
+    const {power} = await utils.getLambdaPower(lambdaARN);
+
+    let lambdaFunctionsToSet = [];
 
     // reminder: configuration updates must run sequentially
     // (otherwise you get a ResourceConflictException)
-    for (let value of powerValues){
-        const alias = 'RAM' + value;
-        await utils.createPowerConfiguration(lambdaARN, value, alias);
+    for (let powerValue of powerValues){
+        const baseAlias = 'RAM' + powerValue;
+        if (!onlyColdStarts){
+            lambdaFunctionsToSet.push({powerValue: powerValue, alias: baseAlias});
+        } else {
+            for (let n of utils.range(num)){
+                let alias = utils.buildAliasString(baseAlias, onlyColdStarts, n);
+                // here we inject a custom env variable to force the creation of a new version
+                // even if the power is the same, which will force a cold start
+                lambdaFunctionsToSet.push({powerValue: powerValue, alias: alias});
+            }
+        }
     }
+    // Publish another version to revert the Lambda Function to its original configuration
+    lambdaFunctionsToSet.push({powerValue: power});
+
+    const returnObj = {
+        initConfigurations: lambdaFunctionsToSet,
+        iterator: {
+            index: 0,
+            count: lambdaFunctionsToSet.length,
+            continue: true,
+        },
+        powerValues: powerValues,
+    };
+    return returnObj;
+};
 
-    await utils.setLambdaPower(lambdaARN, initialPower);
 
-    return powerValues;
+const extractDataFromInput = (event) => {
+    return {
+        lambdaARN: event.lambdaARN,
+        num: parseInt(event.num, 10),
+        powerValues: extractPowerValues(event),
+        onlyColdStarts: !!event.onlyColdStarts,
+    };
 };
 
 const extractPowerValues = (event) => {

lambda/optimizer.js

@@ -26,7 +26,8 @@ module.exports.handler = async(event, context) => {
         await utils.setLambdaPower(lambdaARN, optimalValue);
     } else {
         // create/update alias
-        await utils.createPowerConfiguration(lambdaARN, optimalValue, autoOptimizeAlias);
+        const {envVars} = await utils.getLambdaPower(lambdaARN);
+        await utils.createPowerConfiguration(lambdaARN, optimalValue, autoOptimizeAlias, envVars);
     }
 
     return 'OK';

lambda/publisher.js

@@ -0,0 +1,52 @@
+'use strict';
+
+const utils = require('./utils');
+
+
+module.exports.handler = async(event, context) => {
+    const {iterator, aliases, currConfig, lambdaARN} = validateInputs(event);
+    const {envVars} = await utils.getLambdaPower(lambdaARN);
+    // Alias may not exist when we are reverting the Lambda function to its original configuration
+    if (typeof currConfig.alias !== 'undefined'){
+        envVars.LambdaPowerTuningForceColdStart = currConfig.alias;
+    } else {
+        delete envVars.LambdaPowerTuningForceColdStart;
+    }
+
+    // publish version & assign alias (if present)
+    await utils.createPowerConfiguration(lambdaARN, currConfig.powerValue, currConfig.alias, envVars);
+    if (typeof currConfig.alias !== 'undefined') {
+        // keep track of all aliases
+        aliases.push(currConfig.alias);
+    }
+
+    // update iterator
+    iterator.index++;
+    iterator.continue = (iterator.index < iterator.count);
+    if (!iterator.continue) {
+        delete event.powerValues.initConfigurations;
+    }
+    event.powerValues.aliases = aliases;
+    return event.powerValues;
+};
+function validateInputs(event) {
+    if (!event.lambdaARN) {
+        throw new Error('Missing or empty lambdaARN');
+    }
+    const lambdaARN = event.lambdaARN;
+    if (!(event.powerValues && event.powerValues.iterator && event.powerValues.initConfigurations)){
+        throw new Error('Invalid input');
+    }
+    const iterator = event.powerValues.iterator;
+    if (!(iterator.index >= 0 && iterator.index < iterator.count)){
+        throw new Error('Invalid iterator input');
+    }
+    const initConfigurations = event.powerValues.initConfigurations;
+    const aliases = event.powerValues.aliases || [];
+    const currIdx = iterator.index;
+    const currConfig = initConfigurations[currIdx];
+    if (!(currConfig && currConfig.powerValue)){
+        throw new Error('Invalid configuration');
+    }
+    return {iterator, aliases, currConfig, lambdaARN};
+}

lambda/utils.js

@@ -25,6 +25,14 @@ module.exports.lambdaBaseCost = (region, architecture) => {
     return this.baseCostForRegion(priceMap, region);
 };
 
+module.exports.buildAliasString = (baseAlias, onlyColdStarts, index) => {
+    let alias = baseAlias;
+    if (onlyColdStarts) {
+        alias += `-${index}`;
+    }
+    return alias;
+};
+
 module.exports.allPowerValues = () => {
     const increment = 64;
     const powerValues = [];
@@ -69,28 +77,27 @@ module.exports.verifyAliasExistance = async(lambdaARN, alias) => {
 /**
  * Update power, publish new version, and create/update alias.
  */
-module.exports.createPowerConfiguration = async(lambdaARN, value, alias) => {
+module.exports.createPowerConfiguration = async(lambdaARN, value, alias, envVars) => {
     try {
-        await utils.setLambdaPower(lambdaARN, value);
+        await utils.setLambdaPower(lambdaARN, value, envVars);
 
-        // wait for functoin update to complete
+        // wait for function update to complete
         await utils.waitForFunctionUpdate(lambdaARN);
 
         const {Version} = await utils.publishLambdaVersion(lambdaARN);
+        if (typeof alias === 'undefined'){
+            console.log('No alias defined');
+            return;
+        }
         const aliasExists = await utils.verifyAliasExistance(lambdaARN, alias);
         if (aliasExists) {
             await utils.updateLambdaAlias(lambdaARN, alias, Version);
         } else {
             await utils.createLambdaAlias(lambdaARN, alias, Version);
         }
     } catch (error) {
-        if (error.message && error.message.includes('Alias already exists')) {
-            // shouldn't happen, but nothing we can do in that case
-            console.log('OK, even if: ', error);
-        } else {
-            console.log('error during config creation for value ' + value);
-            throw error; // a real error :)
-        }
+        console.log('error during config creation for value ' + value);
+        throw error; // a real error :)
     }
 };
 
@@ -126,7 +133,7 @@ module.exports.waitForAliasActive = async(lambdaARN, alias) => {
         },
     };
     const lambda = utils.lambdaClientFromARN(lambdaARN);
-    return lambda.waitFor('functionActive', params).promise();
+    return lambda.waitFor('functionActiveV2', params).promise();
 };
 
 /**
@@ -140,7 +147,11 @@ module.exports.getLambdaPower = async(lambdaARN) => {
     };
     const lambda = utils.lambdaClientFromARN(lambdaARN);
     const config = await lambda.getFunctionConfiguration(params).promise();
-    return config.MemorySize;
+    return {
+        power: config.MemorySize,
+        // we need to fetch env vars only to add a new one and force a cold start
+        envVars: (config.Environment || {}).Variables || {},
+    };
 };
 
 /**
@@ -175,11 +186,12 @@ module.exports.getLambdaConfig = async(lambdaARN, alias) => {
 /**
  * Update a given Lambda Function's memory size (always $LATEST version).
  */
-module.exports.setLambdaPower = (lambdaARN, value) => {
+module.exports.setLambdaPower = (lambdaARN, value, envVars) => {
     console.log('Setting power to ', value);
     const params = {
         FunctionName: lambdaARN,
         MemorySize: parseInt(value, 10),
+        Environment: {Variables: envVars},
     };
     const lambda = utils.lambdaClientFromARN(lambdaARN);
     return lambda.updateFunctionConfiguration(params).promise();
@@ -543,9 +555,15 @@ module.exports.regionFromARN = (arn) => {
     return arn.split(':')[3];
 };
 
+let client;
 module.exports.lambdaClientFromARN = (lambdaARN) => {
     const region = this.regionFromARN(lambdaARN);
-    return new AWS.Lambda({region});
+    // create a client only once
+    if (typeof client === 'undefined'){
+        // set Max Retries to 20, increase the retry delay to 500
+        client = new AWS.Lambda({region: region, maxRetries: 20, retryDelayOptions: {base: 500}});
+    }
+    return client;
 };
 
 /**