fix(sqs): fix missing async context in recent aws-sdk/client-sqs version

.gitignore

@@ -7,3 +7,4 @@ test-duration-breakdown
 coverage
 /db2
 instana-*.tgz
+.vscode

packages/collector/test/tracing/cloud/aws-sdk/v2/sqs/receiveMessage.js

@@ -131,6 +131,8 @@ async function receiveAsync() {
     const data = await sqsPromise;
 
     return new Promise(resolve => {
+      // TODO: This could be simplified by wrapping the SQS.prototype.receiveMessage method.
+      //       See v3/sqs.js
       instana.sdk.runInAsyncContext(sqsPromise.instanaAsyncContext, async () => {
         if (data && data.error) {
           log('receive message data error', data.error);

packages/collector/test/tracing/cloud/aws-sdk/v3/sqs/receiver.js

@@ -87,6 +87,7 @@ app.get('/warn-logs', (req, res) => {
 async function runAsPromise(isV2Style = false) {
   const command = new awsSdk3.ReceiveMessageCommand(receiveParams);
   numberOfReceiveMessageAttempts++;
+
   const promise = isV2Style ? sqsv2.receiveMessage(receiveParams) : sqs.send(command);
   let span;
 

packages/collector/test/tracing/cloud/aws-sdk/v3/sqs/test_definition.js

@@ -32,6 +32,7 @@ const sendSnsNotificationToSqsQueue = require('./sendNonInstrumented').sendSnsNo
 
 const sendingMethods = ['v3', 'cb', 'v2'];
 const receivingMethods = ['v3', 'cb', 'v2'];
+
 const getNextSendMethod = require('@instana/core/test/test_util/circular_list').getCircularList(sendingMethods);
 const getNextReceiveMethod = require('@instana/core/test/test_util/circular_list').getCircularList(receivingMethods);
 

packages/core/src/tracing/instrumentation/cloud/aws-sdk/v2/sqs.js

@@ -309,6 +309,7 @@ function instrumentedReceiveMessage(ctx, originalReceiveMessage, originalArgs) {
     // promise use case
     if (typeof awsRequest.promise === 'function' && typeof originalCallback !== 'function') {
       const originalPromiseFn = awsRequest.promise;
+
       awsRequest.promise = cls.ns.bind(function () {
         const promise = originalPromiseFn.apply(awsRequest, arguments);
 

packages/core/src/tracing/instrumentation/cloud/aws-sdk/v3/index.js

@@ -32,6 +32,9 @@ let onFileLoaded = false;
 exports.init = function init() {
   sqsConsumer.init();
 
+  // NOTE: each aws product can have it's own init fn to wrap or unwrap specific functions
+  awsProducts.forEach(awsProduct => awsProduct.init && awsProduct.init(requireHook, shimmer));
+
   /**
    * @aws-sdk/smithly-client >= 3.36.0 changed how the dist structure gets delivered
    * https://github.com/aws/aws-sdk-js-v3/blob/main/packages/smithy-client/CHANGELOG.md#3360-2021-10-08

packages/core/src/tracing/instrumentation/cloud/aws-sdk/v3/sqs.js

@@ -39,6 +39,31 @@ const operations = Object.keys(operationsInfo);
 const SPAN_NAME = 'sqs';
 
 class InstanaAWSSQS extends InstanaAWSProduct {
+  init(requireHook, shimmer) {
+    requireHook.onFileLoad(/@aws-sdk\/client-sqs\/dist-cjs\/SQS\.js/, function (module) {
+      shimmer.wrap(module.SQS.prototype, 'receiveMessage', function (originalReceiveMsgFn) {
+        return function instanaReceiveMessage() {
+          return cls.ns.runAndReturn(() => {
+            const ctx = cls.getAsyncContext();
+            this._instanaCtx = ctx;
+            const promise = originalReceiveMsgFn.apply(this, arguments);
+
+            promise.then = cls.ns.bind(promise.then);
+
+            if (promise.catch) {
+              promise.catch = cls.ns.bind(promise.catch);
+            }
+            if (promise.finally) {
+              promise.finally = cls.ns.bind(promise.finally);
+            }
+
+            return promise;
+          });
+        };
+      });
+    });
+  }
+
   instrumentedSmithySend(ctx, originalSend, smithySendArgs) {
     const commandName = smithySendArgs[0].constructor.name;
     const operation = operationsInfo[commandName];
@@ -166,6 +191,8 @@ class InstanaAWSSQS extends InstanaAWSProduct {
     const command = smithySendArgs[0];
     const sendMessageInput = command.input;
 
+    // Note: we pass in ctx._instanaCtx as the second parameter to runAndReturn which will run the function in the
+    // context belonging to this SQS promise.
     return cls.ns.runAndReturn(() => {
       const self = this;
       const span = cls.startSpan(SPAN_NAME, ENTRY);
@@ -237,6 +264,7 @@ class InstanaAWSSQS extends InstanaAWSProduct {
               setImmediate(() => this.finishSpan(null, span));
             } else if (data && data.Messages && data.Messages.length > 0) {
               const messages = data.Messages;
+
               let tracingAttributes = readTracingAttributes(messages[0].MessageAttributes);
               if (!hasTracingAttributes(tracingAttributes)) {
                 tracingAttributes = readTracingAttributesFromSns(messages[0].Body);
@@ -283,7 +311,7 @@ class InstanaAWSSQS extends InstanaAWSProduct {
 
         return request;
       }
-    });
+    }, ctx._instanaCtx);
   }
 
   buildSpanData(operation, sendMessageInput) {

packages/core/test/config.js

@@ -18,5 +18,9 @@ exports.getTestTimeout = () => {
   if (isCI()) {
     return 30000;
   }
+  // NOTE: Otherwise mocha will interrupt the debugging session quickly.
+  if (process.env.VSCODE_DEBUG === 'true') {
+    return 30000;
+  }
   return 5000;
 };