fix(aws-lambda): improved ssm and coldstarts (#2199)

refs https://jsw.ibm.com/browse/INSTA-66082

Author: kirrg001

packages/aws-lambda/src/ssm.js

@@ -11,6 +11,7 @@ let envValue = null;
 let errorFromAWS = null;
 let initTimeoutInMs = 0;
 let logger;
+let coldStart;
 
 module.exports.reset = () => {
   fetchedValue = null;
@@ -33,8 +34,9 @@ module.exports.validate = () => {
   return true;
 };
 
-module.exports.init = config => {
+module.exports.init = (config, _coldStart) => {
   logger = config.logger;
+  coldStart = _coldStart;
 
   // CASE: INSTANA_SSM_PARAM_NAME is not set, skip
   if (!exports.isUsed()) {
@@ -106,10 +108,15 @@ module.exports.waitAndGetInstanaKey = callback => {
   }
 
   const endInMs = Date.now();
-  const awsTimeoutInMs = process.env.INSTANA_AWS_SSM_TIMEOUT_IN_MS
+  let awsTimeoutInMs = process.env.INSTANA_AWS_SSM_TIMEOUT_IN_MS
     ? Number(process.env.INSTANA_AWS_SSM_TIMEOUT_IN_MS)
     : 1000;
 
+  // CASE: cold start will freeze the handler for a while and we have to allow more time to get the SSM response
+  if (coldStart) {
+    awsTimeoutInMs += 2000;
+  }
+
   // CASE: The time between SSM initialization and waitAndGetInstanaKey is too long to wait for the AWS response.
   //       See init fn - we fetch the key as early as possible.
   if (endInMs - initTimeoutInMs > awsTimeoutInMs) {

packages/aws-lambda/src/wrapper.js

@@ -313,7 +313,7 @@ function init(event, arnInfo, _config) {
   // After core init, because ssm requires require('@aws-sdk/client-ssm'), which triggers
   // the requireHook + shimmer. Any module which requires another external module has to be
   // initialized after the core.
-  ssm.init(config);
+  ssm.init(config, coldStart);
 
   spanBuffer.setIsFaaS(true);
   captureHeaders.init(config);

packages/aws-lambda/test/ssm_test.js

@@ -5,13 +5,14 @@
 'use strict';
 
 const mock = require('@instana/core/test/test_util/mockRequire');
+const { createFakeLogger, delay } = require('@instana/core/test/test_util');
 const sinon = require('sinon');
 const expect = require('chai').expect;
 const ssm = require('../src/ssm');
 
 let sendCommandMock;
 
-describe('Unit: ssm library', () => {
+describe('Unit: ssm library', function () {
   before(() => {
     process.env.AWS_REGION = 'a-region';
   });
@@ -38,7 +39,9 @@ describe('Unit: ssm library', () => {
     expect(ssm.isUsed()).to.be.true;
   });
 
-  describe('init & waitAndGetInstanaKey', () => {
+  describe('init & waitAndGetInstanaKey', function () {
+    this.timeout(5000);
+
     beforeEach(() => {
       sendCommandMock = sinon.stub();
 
@@ -61,7 +64,7 @@ describe('Unit: ssm library', () => {
 
     it('should not fetch AWS SSM value if ssm env is not set', () => {
       ssm.validate();
-      expect(ssm.init({ logger: sinon.stub() })).to.be.undefined;
+      expect(ssm.init({ logger: createFakeLogger() })).to.be.undefined;
       expect(ssm.isUsed()).to.be.false;
       expect(sendCommandMock.callCount).to.equal(0);
     });
@@ -70,7 +73,7 @@ describe('Unit: ssm library', () => {
       process.env.INSTANA_SSM_PARAM_NAME = 'hello instana agent key';
 
       ssm.validate();
-      expect(ssm.init({ logger: { debug: sinon.stub(), warn: sinon.stub() } })).to.be.undefined;
+      expect(ssm.init({ logger: createFakeLogger() })).to.be.undefined;
 
       expect(sendCommandMock.callCount).to.equal(1);
       expect(sendCommandMock.getCall(0).args[0].Name).to.equal('hello instana agent key');
@@ -99,7 +102,7 @@ describe('Unit: ssm library', () => {
         });
       });
 
-      expect(ssm.init({ logger: { debug: sinon.stub() } })).to.be.undefined;
+      expect(ssm.init({ logger: createFakeLogger() })).to.be.undefined;
 
       expect(sendCommandMock.callCount).to.equal(1);
       expect(sendCommandMock.getCall(0).args[0].Name).to.equal('hello instana agent key');
@@ -129,7 +132,7 @@ describe('Unit: ssm library', () => {
         });
       });
 
-      expect(ssm.init({ logger: { debug: sinon.stub() } })).to.be.undefined;
+      expect(ssm.init({ logger: createFakeLogger() })).to.be.undefined;
 
       expect(sendCommandMock.callCount).to.equal(1);
       expect(sendCommandMock.getCall(0).args[0].Name).to.equal('hello instana agent key');
@@ -151,9 +154,11 @@ describe('Unit: ssm library', () => {
       ssm.validate();
       expect(ssm.isUsed()).to.be.true;
 
+      let callsFakeCalled = false;
       sendCommandMock.callsFake(() => {
         return new Promise(resolve => {
           setTimeout(() => {
+            callsFakeCalled = true;
             resolve({
               Parameter: {
                 Value: 'instana-value'
@@ -163,7 +168,7 @@ describe('Unit: ssm library', () => {
         });
       });
 
-      expect(ssm.init({ logger: { debug: sinon.stub() } })).to.be.undefined;
+      expect(ssm.init({ logger: createFakeLogger() })).to.be.undefined;
       expect(sendCommandMock.callCount).to.equal(1);
 
       const interval = ssm.waitAndGetInstanaKey((err, value) => {
@@ -172,10 +177,87 @@ describe('Unit: ssm library', () => {
             '"hello instana agent key", because we have not received a response from AWS.'
         );
         expect(value).to.be.undefined;
-        callback();
+
+        const checkIfCalled = () => {
+          if (callsFakeCalled) {
+            callback();
+          } else {
+            setTimeout(checkIfCalled, 50);
+          }
+        };
+
+        checkIfCalled();
       });
 
       expect(interval).to.exist;
     });
+
+    it('coldstart is true', cb => {
+      process.env.INSTANA_SSM_PARAM_NAME = 'hello instana agent key';
+
+      ssm.validate();
+      expect(ssm.isUsed()).to.be.true;
+
+      let callsFakeCalled = false;
+      sendCommandMock.callsFake(() => {
+        return new Promise(resolve => {
+          setTimeout(() => {
+            callsFakeCalled = true;
+            resolve({
+              Parameter: {
+                Value: 'instana-value'
+              }
+            });
+          }, 2100);
+        });
+      });
+
+      ssm.init({ logger: createFakeLogger() }, true);
+      expect(sendCommandMock.callCount).to.equal(1);
+
+      delay(2000).then(() => {
+        ssm.waitAndGetInstanaKey(err => {
+          expect(err).to.not.exist;
+
+          const checkIfCalled = () => {
+            if (callsFakeCalled) {
+              cb();
+            } else {
+              setTimeout(checkIfCalled, 50);
+            }
+          };
+
+          checkIfCalled();
+        });
+      });
+    });
+
+    it('coldstart is false', cb => {
+      process.env.INSTANA_SSM_PARAM_NAME = 'hello instana agent key';
+
+      ssm.validate();
+      expect(ssm.isUsed()).to.be.true;
+
+      sendCommandMock.callsFake(() => {
+        return new Promise(resolve => {
+          setTimeout(() => {
+            resolve({
+              Parameter: {
+                Value: 'instana-value'
+              }
+            });
+          }, 900);
+        });
+      });
+
+      ssm.init({ logger: createFakeLogger() }, false);
+
+      delay(800).then(() => {
+        ssm.waitAndGetInstanaKey(err => {
+          expect(err).to.not.exist;
+          cb();
+        });
+      });
+    });
   });
 });