Add parameter to power-tune only cold starts

lambda/cleaner.js

@@ -8,13 +8,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
@@ -23,12 +25,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

@@ -20,6 +20,7 @@ module.exports.handler = async(event, context) => {
         preProcessorARN,
         postProcessorARN,
         discardTopBottom,
+        onlyColdStarts,
         sleepBetweenRunsMs,
         disablePayloadLogs,
     } = await extractDataFromInput(event);
@@ -35,8 +36,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
@@ -49,12 +51,13 @@ module.exports.handler = async(event, context) => {
         payloads: payloads,
         preARN: preProcessorARN,
         postARN: postProcessorARN,
+        onlyColdStarts: onlyColdStarts,
         sleepBetweenRunsMs: sleepBetweenRunsMs,
         disablePayloadLogs: disablePayloadLogs,
     };
 
     // wait if the function/alias state is Pending
-    if (isPending) {
+    if (isPending && !onlyColdStarts) {
         await utils.waitForAliasActive(lambdaARN, lambdaAlias);
         console.log('Alias active');
     }
@@ -97,7 +100,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);
@@ -128,16 +135,22 @@ const extractDataFromInput = async(event) => {
         preProcessorARN: input.preProcessorARN,
         postProcessorARN: input.postProcessorARN,
         discardTopBottom: discardTopBottom,
+        onlyColdStarts: !!input.onlyColdStarts,
         sleepBetweenRunsMs: sleepBetweenRunsMs,
         disablePayloadLogs: !!input.disablePayloadLogs,
     };
 };
 
-const runInParallel = async({num, lambdaARN, lambdaAlias, payloads, preARN, postARN, disablePayloadLogs}) => {
+const runInParallel = async({num, lambdaARN, lambdaAlias, payloads, preARN, postARN, disablePayloadLogs, 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, disablePayloadLogs);
+        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, disablePayloadLogs);
         // invocation errors return 200 and contain FunctionError and Payload
         if (invocationResults.FunctionError) {
             let errorMessage = 'Invocation error (running in parallel)';
@@ -150,11 +163,15 @@ const runInParallel = async({num, lambdaARN, lambdaAlias, payloads, preARN, post
     return results;
 };
 
-const runInSeries = async({num, lambdaARN, lambdaAlias, payloads, preARN, postARN, sleepBetweenRunsMs, disablePayloadLogs}) => {
+const runInSeries = async({num, lambdaARN, lambdaAlias, payloads, preARN, postARN, sleepBetweenRunsMs, disablePayloadLogs, onlyColdStarts}) => {
     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, disablePayloadLogs);
+        if (onlyColdStarts){
+            await utils.waitForAliasActive(lambdaARN, aliasToInvoke);
+        }
+        const {invocationResults, actualPayload} = await utils.invokeLambdaWithProcessors(lambdaARN, aliasToInvoke, payloads[i], preARN, postARN, disablePayloadLogs);
         // invocation errors return 200 and contain FunctionError and Payload
         if (invocationResults.FunctionError) {
             let errorMessage = 'Invocation error (running in series)';

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 initConfigurations = [];
 
     // 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){
+            initConfigurations.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
+                initConfigurations.push({powerValue: powerValue, alias: alias});
+            }
+        }
     }
+    // Publish another version to revert the Lambda Function to its original configuration
+    initConfigurations.push({powerValue: power});
+
+    const returnObj = {
+        initConfigurations: initConfigurations,
+        iterator: {
+            index: 0,
+            count: initConfigurations.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 iterator for initialization');
+    }
+    const iterator = event.powerValues.iterator;
+    if (!(iterator.index >= 0 && iterator.index < iterator.count)){
+        throw new Error(`Invalid iterator index: ${iterator.index}`);
+    }
+    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 init configuration: ${currConfig}`);
+    }
+    return {iterator, aliases, currConfig, lambdaARN};
+}

lambda/utils.js

@@ -31,6 +31,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 = [];
@@ -76,14 +84,18 @@ 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);
@@ -142,7 +154,11 @@ module.exports.getLambdaPower = async(lambdaARN) => {
     };
     const lambda = utils.lambdaClientFromARN(lambdaARN);
     const config = await lambda.send(new GetFunctionConfigurationCommand(params));
-    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 || {},
+    };
 };
 
 /**
@@ -177,11 +193,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.send(new UpdateFunctionConfigurationCommand(params));
@@ -610,12 +627,20 @@ module.exports.regionFromARN = (arn) => {
     return arn.split(':')[3];
 };
 
+let client;
 module.exports.lambdaClientFromARN = (lambdaARN) => {
     const region = this.regionFromARN(lambdaARN);
-    return new LambdaClient({
-        region,
-        requestTimeout: 15 * 60 * 1000,
-    });
+    // create a client only once
+    if (typeof client === 'undefined'){
+        // set Max Retries to 20, increase the retry delay to 500
+        client = new LambdaClient({
+            region,
+            maxAttempts: 20,
+            requestTimeout: 15 * 60 * 1000
+        })
+
+    }
+    return client;
 };
 
 /**