diff --git a/lambda/cleaner.js b/lambda/cleaner.js index 4c1d3c3..9b8a408 100644 --- a/lambda/cleaner.js +++ b/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'); diff --git a/lambda/executor.js b/lambda/executor.js index 0fff1eb..0a5469f 100644 --- a/lambda/executor.js +++ b/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)}`); diff --git a/lambda/initializer.js b/lambda/initializer.js index bfd6a3c..fa4be15 100644 --- a/lambda/initializer.js +++ b/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 diff --git a/lambda/utils.js b/lambda/utils.js index a03fcfa..7f0e58a 100644 --- a/lambda/utils.js +++ b/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; diff --git a/test/unit/test-lambda.js b/test/unit/test-lambda.js index 01499b5..ea57b39 100644 --- a/test/unit/test-lambda.js +++ b/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') diff --git a/test/unit/test-utils.js b/test/unit/test-utils.js index 98a7e92..ebeb5b3 100644 --- a/test/unit/test-utils.js +++ b/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); }); });