Add parameter to power-tune only cold starts

lambda/cleaner.js

@@ -7,13 +7,25 @@ const utils = require('./utils');
  */
 module.exports.handler = async(event, context) => {
 
-    const {lambdaARN, powerValues} = event;
+    const {
+        lambdaARN,
+        num,
+        powerValues,
+        onlyColdStarts,
+    } = extractDataFromInput(event);
 
     validateInput(lambdaARN, powerValues); // may throw
 
     const ops = powerValues.map(async(value) => {
-        const alias = 'RAM' + value;
-        await cleanup(lambdaARN, alias); // may throw
+        let baseAlias = 'RAM' + value;
+        if (onlyColdStarts) {
+            for (let n of utils.range(num)){
+                let alias = utils.buildAliasString(baseAlias, onlyColdStarts, n);
+                await cleanup(lambdaARN, alias); // may throw
+            }
+        } else {
+            await cleanup(lambdaARN, baseAlias); // may throw
+        }
     });
 
     // run everything in parallel and wait until completed
@@ -22,6 +34,15 @@ module.exports.handler = async(event, context) => {
     return 'OK';
 };
 
+const extractDataFromInput = (event) => {
+    return {
+        lambdaARN: event.lambdaARN,
+        num: parseInt(event.num, 10),
+        powerValues: event.powerValues,
+        onlyColdStarts: !!event.onlyColdStarts,
+    };
+};
+
 const validateInput = (lambdaARN, powerValues) => {
     if (!lambdaARN) {
         throw new Error('Missing or empty lambdaARN');

lambda/executor.js

@@ -26,6 +26,7 @@ module.exports.handler = async(event, context) => {
         preProcessorARN,
         postProcessorARN,
         discardTopBottom,
+        onlyColdStarts,
     } = await extractDataFromInput(event);
 
     validateInput(lambdaARN, value, num); // may throw
@@ -47,9 +48,9 @@ module.exports.handler = async(event, context) => {
     const payloads = utils.generatePayloads(num, payload);
 
     if (enableParallel) {
-        results = await runInParallel(num, lambdaARN, lambdaAlias, payloads, preProcessorARN, postProcessorARN);
+        results = await runInParallel(num, lambdaARN, lambdaAlias, payloads, preProcessorARN, postProcessorARN, onlyColdStarts);
     } else {
-        results = await runInSeries(num, lambdaARN, lambdaAlias, payloads, preProcessorARN, postProcessorARN);
+        results = await runInSeries(num, lambdaARN, lambdaAlias, payloads, preProcessorARN, postProcessorARN, onlyColdStarts);
     }
 
     // get base cost for Lambda
@@ -104,14 +105,16 @@ const extractDataFromInput = async(event) => {
         preProcessorARN: input.preProcessorARN,
         postProcessorARN: input.postProcessorARN,
         discardTopBottom: discardTopBottom,
+        onlyColdStarts: !!input.onlyColdStarts,
     };
 };
 
-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);
+        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)}`);
@@ -123,11 +126,13 @@ const runInParallel = async(num, lambdaARN, lambdaAlias, payloads, preARN, postA
     return results;
 };
 
-const runInSeries = async(num, lambdaARN, lambdaAlias, payloads, preARN, postARN) => {
+const runInSeries = async(num, lambdaARN, lambdaAlias, payloads, preARN, postARN, onlyColdStarts) => {
     const results = [];
+    // reminder: always start from 0 (same as utils.range)
     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);
+        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,26 +8,51 @@ 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, envVars} = await utils.getLambdaPower(lambdaARN);
 
     // 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);
+        let baseAlias = 'RAM' + value;
+        if (onlyColdStarts) {
+            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
+                envVars.LambdaPowerTuningForceColdStart = alias;
+                await utils.createPowerConfiguration(lambdaARN, value, alias, envVars);
+            }
+        } else {
+            await utils.createPowerConfiguration(lambdaARN, value, baseAlias, envVars);
+        }
     }
 
-    await utils.setLambdaPower(lambdaARN, initialPower);
+    delete envVars.LambdaPowerTuningForceColdStart;
+    // restore power and env variables to initial state
+    await utils.setLambdaPower(lambdaARN, power, envVars);
 
     return powerValues;
 };
 
+const extractDataFromInput = (event) => {
+    return {
+        lambdaARN: event.lambdaARN,
+        num: parseInt(event.num, 10),
+        powerValues: extractPowerValues(event),
+        onlyColdStarts: !!event.onlyColdStarts,
+    };
+};
+
 const extractPowerValues = (event) => {
     var powerValues = event.powerValues; // could be undefined
 

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,9 +77,9 @@ 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
         await utils.waitForFunctionUpdate(lambdaARN);
@@ -122,7 +130,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 || {},
+    };
 };
 
 /**
@@ -145,11 +157,12 @@ module.exports.getLambdaArchitecture = async(lambdaARN) => {
 /**
  * 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();
@@ -462,7 +475,7 @@ module.exports.computeAverageDuration = (durations, discardTopBottom) => {
         // not an error, but worth logging
         // this happens when you have less than 5 invocations
         // (only happens if dryrun or in tests)
-        console.log("not enough results to discard");
+        console.log('not enough results to discard');
     }
 
     const newN = durations.length - 2 * toBeDiscarded;

test/unit/test-lambda.js

@@ -29,7 +29,9 @@ var setLambdaPowerCounter,
     publishLambdaVersionCounter,
     createLambdaAliasCounter,
     updateLambdaAliasCounter,
-    waitForFunctionUpdateCounter;
+    waitForFunctionUpdateCounter,
+    deleteLambdaAliasCounter,
+    deleteLambdaVersionCounter;
 
 // utility to invoke handler (success case)
 const invokeForSuccess = async(handler, event) => {
@@ -84,6 +86,8 @@ describe('Lambda Functions', async() => {
         createLambdaAliasCounter = 0;
         updateLambdaAliasCounter = 0;
         waitForFunctionUpdateCounter = 0;
+        deleteLambdaAliasCounter = 0;
+        deleteLambdaVersionCounter = 0;
 
         sandBox.stub(utils, 'regionFromARN')
             .callsFake((arn) => {
@@ -102,7 +106,10 @@ describe('Lambda Functions', async() => {
         sandBox.stub(utils, 'getLambdaPower')
             .callsFake(async() => {
                 getLambdaPowerCounter++;
-                return 1024;
+                return {
+                    power: 1024,
+                    envVars: {},
+                };
             });
         setLambdaPowerStub = sandBox.stub(utils, 'setLambdaPower')
             .callsFake(async() => {
@@ -229,6 +236,24 @@ describe('Lambda Functions', async() => {
             expect(waitForFunctionUpdateCounter).to.be(powerValues.length);
         });
 
+        it('should create N*num aliases and versions when onlyColdStarts', async() => {
+            await invokeForSuccess(handler, {
+                lambdaARN: 'arnOK',
+                num: 5,
+                onlyColdStarts: true,
+            });
+
+            const total = powerValues.length * 5;
+
+            // +1 because it will also reset power to its initial value
+            expect(setLambdaPowerCounter).to.be(total + 1);
+
+            expect(getLambdaPowerCounter).to.be(1);
+            expect(publishLambdaVersionCounter).to.be(total);
+            expect(createLambdaAliasCounter).to.be(total);
+            expect(waitForFunctionUpdateCounter).to.be(total);
+        });
+
         it('should explode if something goes wrong during power set', async() => {
             setLambdaPowerStub && setLambdaPowerStub.restore();
             setLambdaPowerStub = sandBox.stub(utils, 'setLambdaPower')
@@ -285,11 +310,13 @@ describe('Lambda Functions', async() => {
             deleteLambdaAliasStub && deleteLambdaAliasStub.restore();
             deleteLambdaAliasStub = sandBox.stub(utils, 'deleteLambdaAlias')
                 .callsFake(async() => {
+                    deleteLambdaAliasCounter++;
                     return 'OK';
                 });
             deleteLambdaVersionStub && deleteLambdaVersionStub.restore();
             deleteLambdaVersionStub = sandBox.stub(utils, 'deleteLambdaVersion')
                 .callsFake(async() => {
+                    deleteLambdaVersionCounter++;
                     return 'OK';
                 });
         });
@@ -300,6 +327,23 @@ describe('Lambda Functions', async() => {
             await invokeForSuccess(handler, eventOK);
         });
 
+        it('should delete all versions and aliases', async() => {
+            await invokeForSuccess(handler, eventOK);
+            expect(deleteLambdaAliasCounter).to.be(3);
+            expect(deleteLambdaVersionCounter).to.be(3);
+        });
+
+        it('should delete all versions and aliases, when onlyColdStarts', async() => {
+            await invokeForSuccess(handler, {
+                lambdaARN: 'arnOK',
+                powerValues: ['128', '256', '512'],
+                num: 10,
+                onlyColdStarts: true,
+            });
+            expect(deleteLambdaAliasCounter).to.be(30);
+            expect(deleteLambdaVersionCounter).to.be(30);
+        });
+
         it('should work fine even if the version does not exist', async() => {
             deleteLambdaVersionStub && deleteLambdaVersionStub.restore();
             deleteLambdaVersionStub = sandBox.stub(utils, 'deleteLambdaVersion')

test/unit/test-utils.js

@@ -24,7 +24,13 @@ const sandBox = sinon.createSandbox();
 
 // AWS SDK mocks
 AWS.mock('Lambda', 'getAlias', {});
-AWS.mock('Lambda', 'getFunctionConfiguration', {MemorySize: 1024, State: 'Active', LastUpdateStatus: 'Successful', Architectures: ['x86_64']});
+AWS.mock('Lambda', 'getFunctionConfiguration', {
+    MemorySize: 1024,
+    State: 'Active',
+    LastUpdateStatus: 'Successful',
+    Architectures: ['x86_64'],
+    Environment: {Variables: {TEST: 'OK'}},
+});
 AWS.mock('Lambda', 'updateFunctionConfiguration', {});
 AWS.mock('Lambda', 'publishVersion', {});
 AWS.mock('Lambda', 'deleteFunction', {});
@@ -106,9 +112,26 @@ describe('Lambda Utils', () => {
     });
 
     describe('getLambdaPower', () => {
-        it('should return the memory value', async() => {
+        it('should return the power value and env vars', async() => {
+            const value = await utils.getLambdaPower('arn:aws:lambda:us-east-1:XXX:function:YYY');
+            expect(value.power).to.be(1024);
+            expect(value.envVars).to.be.an('object');
+            expect(value.envVars.TEST).to.be('OK');
+        });
+
+        it('should return the power value and env vars even when empty env', async() => {
+            AWS.remock('Lambda', 'getFunctionConfiguration', {
+                MemorySize: 1024,
+                State: 'Active',
+                LastUpdateStatus: 'Successful',
+                Architectures: ['x86_64'],
+                Environment: null, // this is null if no vars are set
+            });
+
             const value = await utils.getLambdaPower('arn:aws:lambda:us-east-1:XXX:function:YYY');
-            expect(value).to.be(1024);
+            expect(value.power).to.be(1024);
+            expect(value.envVars).to.be.an('object');
+            expect(value.envVars.TEST).to.be(undefined);
         });
     });