chore: add un-trace command

Author: tschoffelen

package.json

@@ -11,6 +11,7 @@
   "scripts": {
     "deploy": "yarn build && yarn turbo deploy --env-mode=loose --color",
     "auto-trace": "cd packages/api && yarn auto-trace",
+    "un-trace": "cd packages/api && yarn un-trace",
     "build": "yarn turbo build --env-mode=loose --color",
     "test": "yarn turbo test --env-mode=loose --color",
     "start": "turbo dev --env-mode=loose --color",

packages/api/package.json

@@ -7,6 +7,7 @@
     "dev": "serverless offline start",
     "deploy": "serverless deploy --verbose",
     "auto-trace": "serverless invoke local -f main -d '{\"action\": \"auto-trace\"}'",
+    "un-trace": "serverless invoke local -f main -d '{\"action\": \"un-trace\"}'",
     "test": "jest"
   },
   "dependencies": {

packages/api/src/events/un-trace.js

@@ -0,0 +1,107 @@
+import {
+  GetFunctionCommand,
+  LambdaClient,
+  ListFunctionsCommand,
+  UpdateFunctionConfigurationCommand,
+} from "@aws-sdk/client-lambda";
+import pLimit from "p-limit";
+
+import Logger from "../lib/logger";
+
+const lambdaExecWrapper = "/opt/nodejs/tracer_wrapper";
+const logger = new Logger("auto-trace");
+
+const getAccountLambdas = async () => {
+  const lambdaClient = new LambdaClient();
+
+  let nextToken = null;
+  let lambdas = [];
+
+  do {
+    const listFunctionsCommand = new ListFunctionsCommand({
+      Marker: nextToken,
+      MaxItems: 50,
+    });
+    const { Functions, NextMarker } =
+      await lambdaClient.send(listFunctionsCommand);
+    lambdas = lambdas.concat(Functions);
+    nextToken = NextMarker;
+  } while (nextToken);
+
+  return lambdas;
+};
+
+const updateLambda = async (lambda, arnBase) => {
+  const command = new UpdateFunctionConfigurationCommand({
+    FunctionName: lambda.FunctionName,
+    Layers: [
+      ...(lambda.Layers || [])
+        .map((layer) => layer.Arn)
+        .filter((arn) => !arn.startsWith(arnBase)),
+    ],
+    Environment: {
+      ...(lambda.Environment || {}),
+      Variables: {
+        ...(lambda.Environment?.Variables || {}),
+        AUTO_TRACE_QUEUE_URL: undefined,
+        AUTO_TRACE_QUEUE_REGION: undefined,
+        AWS_LAMBDA_EXEC_WRAPPER: undefined,
+      },
+    },
+  });
+
+  await new LambdaClient().send(command);
+};
+
+const getLambdaTags = async (lambda) => {
+  // try not to run into the 100 requests per second limit
+  await new Promise((resolve) => setTimeout(resolve, 250));
+  const lambdaClient = new LambdaClient();
+  const { Tags } = await lambdaClient.send(
+    new GetFunctionCommand({
+      FunctionName: lambda.FunctionName,
+    }),
+  );
+
+  return Tags;
+};
+
+export const unTrace = async () => {
+  // Check if we know our Lambda Layer ARN
+  const arn = process.env.LAMBDA_LAYER_ARN;
+  const arnBase = arn?.substring(0, arn?.lastIndexOf(":") + 1);
+  if (!arn || !arnBase) {
+    throw new Error("LAMBDA_LAYER_ARN is not defined");
+  }
+
+  // List all the lambda functions in the AWS account
+  let lambdas = await getAccountLambdas();
+  logger.info(`Found ${lambdas.length} lambdas in the account`);
+
+  // Update qualifying lambdas
+  const limit = pLimit(4);
+  await Promise.all(
+    lambdas.map((lambda) =>
+      limit(async () => {
+        const layers = lambda.Layers || [];
+        const envVars = lambda.Environment?.Variables || {};
+
+        const isTraceStack = envVars.LAMBDA_LAYER_ARN === arn;
+        const isUpdating = lambda.LastUpdateStatus === "InProgress";
+        const hasLayer = layers.find(({ Arn }) => Arn.startsWith(arnBase));
+        const hasOtherWrapper =
+          envVars.AWS_LAMBDA_EXEC_WRAPPER &&
+          envVars.AWS_LAMBDA_EXEC_WRAPPER !== lambdaExecWrapper;
+
+        if (hasLayer && !isTraceStack && !isUpdating && !hasOtherWrapper) {
+          try {
+            await updateLambda(lambda, arnBase);
+            logger.info(`✓ Removed layer from ${lambda.FunctionName}`);
+          } catch (e) {
+            logger.warn(`✗ Failed to update ${lambda.FunctionName}`, e);
+          }
+        }
+      }),
+    ),
+  );
+};

packages/api/src/index.js

@@ -10,6 +10,7 @@ import authRoute from "./routes/auth";
 import usersRoute from "./routes/users";
 
 import { autoTrace } from "./events/auto-trace";
+import {unTrace} from "./events/un-trace";
 import { auth } from "./routes/auth/middleware";
 
 const app = new Hono();
@@ -38,6 +39,9 @@ export const handler = (event, context) => {
   if (event.action === "auto-trace") {
     return autoTrace();
   }
+  if (event.action === "un-trace") {
+    return unTrace();
+  }
 
   return httpApp(event, context);
 };

packages/lambda-layer/lib/utils/httpSpansAgent.js

@@ -1,5 +1,5 @@
 const { SQSClient, SendMessageCommand } = require("@aws-sdk/client-sqs");
-const { FetchHttpHandler } = require("@smithy/fetch-http-handler");
+// const { FetchHttpHandler } = require("@smithy/fetch-http-handler");
 
 /*
  * This module handles sending spans to an SQS queue.
@@ -16,7 +16,7 @@ exports.HttpSpansAgent = (() => {
     try {
       sqsClient = new SQSClient({
         region: process.env.AUTO_TRACE_QUEUE_REGION || "eu-west-1",
-        requestHandler: new FetchHttpHandler({}),
+        // requestHandler: new FetchHttpHandler({}),
       });
     } catch (error) {
       console.warn("Error initializing SQS client:", error);
@@ -46,15 +46,15 @@ exports.HttpSpansAgent = (() => {
 
       // This is a workaround to avoid the fetch function being overridden
       // @ts-ignore
-      if (fetch && fetch.__originalFetch) {
-        const wrappedFetch = fetch;
-        // @ts-ignore
-        fetch = wrappedFetch.__originalFetch;
-        setTimeout(() => {
-          // @ts-ignore
-          fetch = wrappedFetch;
-        }, 1);
-      }
+      // if (fetch && fetch.__originalFetch) {
+      //   const wrappedFetch = fetch;
+      //   // @ts-ignore
+      //   fetch = wrappedFetch.__originalFetch;
+      //   setTimeout(() => {
+      //     // @ts-ignore
+      //     fetch = wrappedFetch;
+      //   }, 1);
+      // }
 
       await sqsClient.send(command);
     } catch (error) {