feat: use SQS to send spans

tschoffelen · tschoffelen · commit 09953a7b4e64 · 2025-03-30T14:22:56.000+01:00
diff --git a/README.md b/README.md
@@ -34,7 +34,7 @@ Once you're done, it will display the URL for your new dashboard.
 ### Auto-tracing
 
 Once installed, all Lambda functions in your AWS account will automatically be traced,
-by adding a Lambda layer to them. You'll also notice a new environment variable called `AUTO_TRACE_HOST`.
+by adding a Lambda layer to them. You'll also notice a new environment variable called `AUTO_TRACE_QUEUE_URL`.
 
 If you wish to disable tracing for a specific function, you can add the environment variable `AUTO_TRACE_EXCLUDE`
 with a value of `1`. This will keep the Lambda layer from being added.
diff --git a/packages/api/serverless.yml b/packages/api/serverless.yml
@@ -54,6 +54,8 @@ provider:
     LAMBDA_LAYER_ARN: !Ref TracerLambdaLayer
     AUTO_TRACE_EXCLUDE: 1
     STORAGE_BUCKET_NAME: ${self:custom.bucketName}
+    QUEUE_URL: !Ref Queue
+    QUEUE_REGION: ${aws:region}
   httpApi:
     payload: "2.0"
     cors: true
@@ -93,7 +95,6 @@ provider:
           Action:
             - s3:GetObject
             - s3:PutObject
-            - s3:DeleteObject
             - s3:ListBucket
           Resource:
             - arn:aws:s3:::${self:custom.bucketName}
@@ -111,12 +112,21 @@ layers:
       - nodejs22.x
 
 functions:
+  collector:
+    handler: src/events/collector.handler
+    timeout: 300
+    memorySize: 2048
+    reservedConcurrency: 1
+    events:
+      - sqs:
+          arn: !GetAtt Queue.Arn
+          batchSize: 100
+          maximumBatchingWindow: 30
+
   main:
     handler: src/index.handler
     timeout: 900
     memorySize: 1024
-    reservedConcurrency: 1
-    url: true
     events:
       - httpApi:
           path: /{proxy+}
@@ -127,5 +137,6 @@ functions:
             action: auto-trace
 
 resources:
-  - ${file(./src/infrastructure/database/table.yml)}
-  - ${file(./src/infrastructure/bucket.yml)}
+  - ${file(./src/infrastructure/table.yml)}
+  - ${file(./src/infrastructure/bucket.yml)}
+  - ${file(./src/infrastructure/queue.yml)}
diff --git a/packages/api/src/events/auto-trace.js b/packages/api/src/events/auto-trace.js
@@ -1,6 +1,5 @@
 import {
   GetFunctionCommand,
-  GetFunctionUrlConfigCommand,
   LambdaClient,
   ListFunctionsCommand,
   UpdateFunctionConfigurationCommand,
@@ -11,7 +10,12 @@ import { acquireLock, releaseLock } from "../lib/locks";
 import Logger from "../lib/logger";
 import { put } from "../lib/database";
 
-const supportedRuntimes = ["nodejs16.x", "nodejs18.x", "nodejs20.x", "nodejs22.x"];
+const supportedRuntimes = [
+  "nodejs16.x",
+  "nodejs18.x",
+  "nodejs20.x",
+  "nodejs22.x",
+];
 const lambdaExecWrapper = "/opt/nodejs/tracer_wrapper";
 
 const logger = new Logger("auto-trace");
@@ -36,28 +40,7 @@ const getAccountLambdas = async () => {
   return lambdas;
 };
 
-const getEdgeEndpoint = async (lambdas) => {
-  const lambdaName = `${process.env.SERVICE}-${process.env.STAGE}-main`;
-
-  try {
-    const command = new GetFunctionUrlConfigCommand({
-      FunctionName: lambdaName,
-    });
-
-    const res = await new LambdaClient().send(command);
-
-    return res.FunctionUrl?.replace("https://", "").replace("/", "");
-  } catch (e) {
-    if (e.name === "ResourceNotFoundException") {
-      throw new Error(
-        `Lambda ${lambdaName} not found - could not determine edge endpoint.`,
-      );
-    }
-    throw e;
-  }
-};
-
-const updateLambda = async (lambda, arnBase, edgeEndpoint) => {
+const updateLambda = async (lambda, arnBase, queue, queueRegion) => {
   const command = new UpdateFunctionConfigurationCommand({
     FunctionName: lambda.FunctionName,
     Layers: [
@@ -70,7 +53,8 @@ const updateLambda = async (lambda, arnBase, edgeEndpoint) => {
       ...(lambda.Environment || {}),
       Variables: {
         ...(lambda.Environment?.Variables || {}),
-        AUTO_TRACE_HOST: edgeEndpoint,
+        AUTO_TRACE_QUEUE_URL: queue,
+        AUTO_TRACE_QUEUE_REGION: queueRegion,
         AWS_LAMBDA_EXEC_WRAPPER: lambdaExecWrapper,
       },
     },
@@ -127,7 +111,11 @@ export const autoTrace = async () => {
   }
 
   // Get our Lambda URL endpoint for the collector
-  const edgeEndpoint = await getEdgeEndpoint();
+  const queue = process.env.QUEUE_URL;
+  const queueRegion = process.env.QUEUE_REGION || "eu-west-1";
+  if (!queue) {
+    throw new Error("QUEUE_URL is not defined");
+  }
 
   // Make sure we lock so that only one process is updating lambdas
   const lockAcquired = await acquireLock("auto-trace");
@@ -137,7 +125,7 @@ export const autoTrace = async () => {
   }
 
   // List all the lambda functions in the AWS account
-  const lambdas = await getAccountLambdas();
+  let lambdas = await getAccountLambdas();
   logger.info(`Found ${lambdas.length} lambdas in the account`);
 
   // Update qualifying lambdas
@@ -151,7 +139,7 @@ export const autoTrace = async () => {
         const isTraceStack = envVars.LAMBDA_LAYER_ARN === arn;
         const isUpdating = lambda.LastUpdateStatus === "InProgress";
         const hasDisableEnvVar = envVars.AUTO_TRACE_EXCLUDE;
-        const hasWrongEndpoint = envVars.AUTO_TRACE_HOST !== edgeEndpoint;
+        const hasWrongEndpoint = envVars.AUTO_TRACE_QUEUE_URL !== queue;
         const hasOtherWrapper =
           envVars.AWS_LAMBDA_EXEC_WRAPPER &&
           envVars.AWS_LAMBDA_EXEC_WRAPPER !== lambdaExecWrapper;
@@ -192,7 +180,7 @@ export const autoTrace = async () => {
         }
 
         try {
-          await updateLambda(lambda, arnBase, edgeEndpoint);
+          await updateLambda(lambda, arnBase, queue, queueRegion);
 
           logger.info(`✓ Updated ${lambda.FunctionName}`);
           await saveFunctionInfo(lambda, "enabled");
diff --git a/packages/api/src/events/collector.js b/packages/api/src/events/collector.js
@@ -1,12 +1,9 @@
-import { Hono } from "hono";
 import { defULID } from "@thi.ng/ksuid";
 
-import { saveHourlyStat } from "../../lib/stats";
-import { getErrorKey } from "../../lib/errors";
-import { saveInvocation } from "../../lib/invocations";
-import { store } from "../../lib/storage";
-
-const app = new Hono();
+import { saveHourlyStat } from "../lib/stats";
+import { getErrorKey } from "../lib/errors";
+import { saveInvocation } from "../lib/invocations";
+import { store } from "../lib/storage";
 
 const id = defULID();
 const getId = () => {
@@ -17,19 +14,22 @@ const getId = () => {
 // This is used to avoid sending the transaction to the database before we know the invocation ID
 const transactionCache = {};
 
-app.post("/", async (c) => {
-  const body = await c.req.json();
+export const handler = async ({ Records }) => {
+  // Read all of the messages off the queue
+  for (const record of Records) {
+    const body = JSON.parse(JSON.parse(record.body));
 
-  for (const span of body) {
-    if (process.env.TRACER_TOKEN && span.token !== process.env.TRACER_TOKEN) {
-      console.log(`Invalid token: ${span.token}`);
-      continue;
-    }
+    for (const span of body) {
+      if (process.env.TRACER_TOKEN && span.token !== process.env.TRACER_TOKEN) {
+        console.log(`Invalid token: ${span.token}`);
+        continue;
+      }
 
-    const groupKey = span.transactionId || span.transaction_id;
-    if (!transactionCache[groupKey]) transactionCache[groupKey] = [];
+      const groupKey = span.transactionId || span.transaction_id;
+      if (!transactionCache[groupKey]) transactionCache[groupKey] = [];
 
-    transactionCache[groupKey].push(span);
+      transactionCache[groupKey].push(span);
+    }
   }
 
   // Check transactions cache to see if there's any transactions we can flush
@@ -108,8 +108,4 @@ app.post("/", async (c) => {
     // Delete the transaction from the cache
     delete transactionCache[transactionId];
   }
-
-  return c.json({ success: true });
-});
-
-export default app;
+};
diff --git a/packages/api/src/index.js b/packages/api/src/index.js
@@ -4,7 +4,6 @@ import { handle } from "hono/aws-lambda";
 import { secureHeaders } from "hono/secure-headers";
 import { serveStatic } from "@hono/node-server/serve-static";
 
-import collectorRoute from "./routes/collector";
 import exploreRoute from "./routes/explore";
 import autoTraceRoute from "./routes/auto-trace";
 import authRoute from "./routes/auth";
@@ -17,7 +16,6 @@ const app = new Hono();
 app.use(secureHeaders());
 
 app.route("/api/auth", authRoute);
-app.route("/api/spans", collectorRoute);
 app.route("/api/auto-trace", autoTraceRoute);
 
 app.use("/api/explore/*", auth);
diff --git a/packages/api/src/infrastructure/queue.yml b/packages/api/src/infrastructure/queue.yml
@@ -0,0 +1,23 @@
+Resources:
+  Queue:
+    Type: AWS::SQS::Queue
+    Properties:
+      QueueName: ${self:service}-queue-${sls:stage}
+      VisibilityTimeout: 300
+      MessageRetentionPeriod: 1209600
+  QueuePolicy:
+    Type: AWS::SQS::QueuePolicy
+    Properties:
+      Queues:
+        - !Ref Queue
+      PolicyDocument:
+        Version: "2012-10-17"
+        Statement:
+          - Effect: Allow
+            Principal:
+              AWS: "*"
+            Action: sqs:SendMessage
+            Resource: !GetAtt Queue.Arn
+            Condition:
+              StringEquals:
+                aws:PrincipalAccount: ${aws:accountId}
diff --git a/packages/api/src/infrastructure/table.yml b/packages/api/src/infrastructure/table.yml
diff --git a/packages/api/src/lib/spans.js b/packages/api/src/lib/spans.js
@@ -41,33 +41,20 @@ const removeGroupingKeys = (spans) => {
 const simplifySpans = (spans) => {
   return spans.map((span) => {
     const {
-      info,
-      vendor,
-      version,
-      runtime,
       token,
       region,
       type,
       transactionId,
       isMetadata,
-      memoryAllocated,
-      readiness,
       messageVersion,
       account,
       invokedArn,
       invokedVersion,
       lumigo_execution_tags_no_scrub,
       ...rest
     } = span;
-    const { traceId, tracer, logGroupName, logStreamName, ...restInfo } =
-      info || {};
-    return {
-      ...rest,
-      info: {
-        traceId: traceId?.Root,
-        ...restInfo,
-      },
-    };
+
+    return { ...rest };
   });
 };
 
@@ -80,12 +67,10 @@ export const groupSpans = (spans) => {
   for (const span of spans) {
     const latestSpan = groupedSpans[groupedSpans.length - 1] || {};
     const similarSpans = groupedSpans.filter(
-      (s) => s.groupingKey === span.groupingKey,
+      (s) => s.extendedGroupingKey === span.extendedGroupingKey,
     );
 
-    if (span.extendedGroupingKey === latestSpan.extendedGroupingKey) {
-      latestSpan.instances++;
-    } else if (similarSpans.length >= 25) {
+    if (similarSpans.length >= 20) {
       similarSpans[similarSpans.length - 1].instances++;
     } else {
       groupedSpans.push(span);
diff --git a/packages/deploy-script/README.md b/packages/deploy-script/README.md
@@ -34,7 +34,7 @@ Once you're done, it will display the URL for your new dashboard.
 ### Auto-tracing
 
 Once installed, all Lambda functions in your AWS account will automatically be traced,
-by adding a Lambda layer to them. You'll also notice a new environment variable called `AUTO_TRACE_HOST`.
+by adding a Lambda layer to them. You'll also notice a new environment variable called `AUTO_TRACE_QUEUE_URL`.
 
 If you wish to disable tracing for a specific function, you can add the environment variable `AUTO_TRACE_EXCLUDE`
 with a value of `1`. This will keep the Lambda layer from being added.
diff --git a/packages/lambda-layer/index.js b/packages/lambda-layer/index.js
@@ -4,13 +4,12 @@ try {
 } catch (e) {
   config = {
     token: "t_0000000000000000",
-    edgeHost: process.env.AUTO_TRACE_HOST,
   };
 }
 
 const tracer = require("@lumigo/tracer")({
   token: config.token,
-  edgeHost: process.env.AUTO_TRACE_HOST || config.edgeHost,
+  edgeHost: "localhost",
 });
 
 const verbose = process.env.TRACER_LOG_VERBOSE;
diff --git a/packages/lambda-layer/lib/utils/httpSpansAgent.js b/packages/lambda-layer/lib/utils/httpSpansAgent.js
@@ -0,0 +1,43 @@
+const { SQSClient, SendMessageCommand } = require("@aws-sdk/client-sqs");
+
+/*
+ * This module handles sending spans to an SQS queue.
+ * It initializes the SQS client and provides a function to post spans.
+ *
+ * It is dropped in as a replacement for the original HTTP spans agent:
+ * https://github.com/lumigo-io/lumigo-node/blob/master/src/httpSpansAgent.js
+ */
+
+exports.HttpSpansAgent = (() => {
+  let sqsClient;
+
+  const initAgent = () => {
+    sqsClient = new SQSClient({
+      region: process.env.AUTO_TRACE_QUEUE_REGION || "eu-west-1",
+    });
+  };
+
+  const cleanSessionInstance = () => {
+    sqsClient = null;
+  };
+
+  const postSpans = async (requestBody) => {
+    if (!sqsClient) {
+      initAgent();
+    }
+
+    try {
+      const params = {
+        QueueUrl: process.env.AUTO_TRACE_QUEUE_URL,
+        MessageBody: JSON.stringify(requestBody),
+      };
+
+      const command = new SendMessageCommand(params);
+      await sqsClient.send(command);
+    } catch (error) {
+      console.warn("Error sending trace spans:", error);
+    }
+  };
+
+  return { postSpans, cleanSessionInstance, initAgent };
+})();
diff --git a/packages/lambda-layer/package.json b/packages/lambda-layer/package.json
@@ -6,7 +6,8 @@
   "main": "index.js",
   "scripts": {
     "test": "yarn build && jest",
-    "build": "esbuild index.js --bundle --platform=node --minify --target=node16 --outfile=artifact/nodejs/dist/auto_tracer_wrapper.js"
+    "replace-agent": "rm node_modules/@lumigo/tracer/dist/httpSpansAgent.js && cp lib/utils/httpSpansAgent.js node_modules/@lumigo/tracer/dist/httpSpansAgent.js",
+    "build": "yarn replace-agent && esbuild index.js --bundle --platform=node --minify --target=node16 --outfile=artifact/nodejs/dist/auto_tracer_wrapper.js"
   },
   "files": [
     "nodejs/dist/index.js"
@@ -17,6 +18,7 @@
     "jest": "^29.7.0"
   },
   "dependencies": {
+    "@aws-sdk/client-sqs": "^3.777.0",
     "@lumigo/tracer": "^1.107.1",
     "axios": "^1.8.4",
     "node-fetch": "^3.3.2"
diff --git a/yarn.lock b/yarn.lock
