Skip to content

Commit

Permalink
feat: save auto trace status
Browse files Browse the repository at this point in the history
  • Loading branch information
@tschoffelen
tschoffelen committed Sep 7, 2024
1 parent 9de7a8d commit 9f38f0b
Show file tree
Hide file tree
Showing 4 changed files with 128 additions and 64 deletions.
3 changes: 2 additions & 1 deletion packages/api/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@
"@includable/serverless-middleware": "^2.2.0",
"date-fns": "^3.6.0",
"dynamodb-toolbox": "^1.3.8",
"hono": "^4.5.10"
"hono": "^4.5.10",
"p-limit": "^6.1.0"
},
"devDependencies": {
"esbuild": "^0.20.1",
Expand Down
171 changes: 108 additions & 63 deletions packages/api/src/events/auto-trace.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,10 @@ import {
ApiGatewayV2Client,
GetApisCommand,
} from "@aws-sdk/client-apigatewayv2";
import pLimit from "p-limit";
import { acquireLock, releaseLock } from "../lib/locks";
import Logger from "../lib/logger";
import { put } from "../lib/database";

const supportedRuntimes = ["nodejs16.x", "nodejs18.x", "nodejs20.x"];
const lambdaExecWrapper = "/opt/nodejs/tracer_wrapper";
Expand Down Expand Up @@ -51,6 +53,48 @@ const getApiEndpoint = async () => {
return item.ApiEndpoint?.replace("https://", "");
};

const updateLambda = async (lambda, arnBase, edgeEndpoint) => {
const updateFunctionConfigurationCommand =
new UpdateFunctionConfigurationCommand({
FunctionName: lambda.FunctionName,
Layers: [
...(lambda.Layers || [])
.map((layer) => layer.Arn)
.filter((arn) => !arn.startsWith(arnBase)),
process.env.LAMBDA_LAYER_ARN,
],
Environment: {
...(lambda.Environment || {}),
Variables: {
...(lambda.Environment?.Variables || {}),
AUTO_TRACE_HOST: edgeEndpoint,
AWS_LAMBDA_EXEC_WRAPPER: lambdaExecWrapper,
},
},
});

const res = await new LambdaClient().send(updateFunctionConfigurationCommand);
logger.info(res);
};

const saveFunctionInfo = async (lambda, traceStatus) => {
await put(
{
pk: `function#${process.env.AWS_REGION}#${lambda.FunctionName}`,
sk: `function#${process.env.AWS_REGION}`,
name: lambda.FunctionName,
type: "function",
runtime: lambda.Runtime,
region: process.env.AWS_REGION,
arn: lambda.FunctionArn,
memoryAllocated: lambda.MemorySize,
timeout: lambda.Timeout,
traceStatus,
},
true,
);
};

export const autoTrace = async () => {
// Check if we know our Lambda Layer ARN
const arn = process.env.LAMBDA_LAYER_ARN;
Expand All @@ -71,69 +115,70 @@ export const autoTrace = async () => {

// List all the lambda functions in the AWS account
const lambdas = await getAccountLambdas();

// Find all lambdas that need a layer added
let lambdasWithoutLayer = lambdas.filter((lambda) => {
const layers = lambda.Layers || [];
const envVars = lambda.Environment?.Variables || {};

const isTraceStack = envVars.LAMBDA_LAYER_ARN === arn;
const isUpdating = lambda.LastUpdateStatus === "InProgress";
const hasDisableEnvVar = envVars.AUTO_TRACE_EXCLUDE;
const hasOtherWrapper =
envVars.AWS_LAMBDA_EXEC_WRAPPER &&
envVars.AWS_LAMBDA_EXEC_WRAPPER !== lambdaExecWrapper;
const hasSupportedRuntime = supportedRuntimes.includes(lambda.Runtime);
const hasLayer = layers.find(({ Arn }) => Arn.startsWith(arnBase));
const hasUpdate = layers.find(
({ Arn }) => Arn.startsWith(arnBase) && Arn !== arn,
);

return (
(!hasLayer || hasUpdate) &&
!hasDisableEnvVar &&
!hasOtherWrapper &&
!isUpdating &&
!isTraceStack &&
hasSupportedRuntime
);
});

logger.info(`Found ${lambdasWithoutLayer.length} lambdas to update`);

for (const lambda of lambdasWithoutLayer) {
try {
const updateFunctionConfigurationCommand =
new UpdateFunctionConfigurationCommand({
FunctionName: lambda.FunctionName,
Layers: [
...(lambda.Layers || [])
.map((layer) => layer.Arn)
.filter((arn) => !arn.startsWith(arnBase)),
process.env.LAMBDA_LAYER_ARN,
],
Environment: {
...(lambda.Environment || {}),
Variables: {
...(lambda.Environment?.Variables || {}),
AUTO_TRACE_HOST: edgeEndpoint,
AWS_LAMBDA_EXEC_WRAPPER: lambdaExecWrapper,
},
},
});

const res = await new LambdaClient().send(
updateFunctionConfigurationCommand,
);
logger.info(res);

logger.info(`✓ Updated ${lambda.FunctionName}`);
// TODO: save function info in DynamoDB
} catch (e) {
logger.warn(`✗ Failed to update ${lambda.FunctionName}`);
logger.warn(e);
}
}
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 hasDisableEnvVar = envVars.AUTO_TRACE_EXCLUDE;
const hasOtherWrapper =
envVars.AWS_LAMBDA_EXEC_WRAPPER &&
envVars.AWS_LAMBDA_EXEC_WRAPPER !== lambdaExecWrapper;
const hasSupportedRuntime = supportedRuntimes.includes(lambda.Runtime);
const hasLayer = layers.find(({ Arn }) => Arn.startsWith(arnBase));
const hasUpdate = layers.find(
({ Arn }) => Arn.startsWith(arnBase) && Arn !== arn,
);

if (hasLayer && !hasUpdate) {
logger.info(`- ${lambda.FunctionName} already has the latest layer`);
await saveFunctionInfo(lambda, "enabled");
return;
}
if (hasDisableEnvVar) {
logger.info(`- ${lambda.FunctionName} has AUTO_TRACE_EXCLUDE`);
await saveFunctionInfo(lambda, "excluded");
return;
}
if (hasOtherWrapper) {
logger.info(`- ${lambda.FunctionName} has a custom wrapper`);
await saveFunctionInfo(lambda, "custom_wrapper");
return;
}
if (isUpdating) {
logger.info(`- ${lambda.FunctionName} is already updating`);
await saveFunctionInfo(lambda, "update_in_progress");
return;
}
if (isTraceStack) {
logger.info(`- ${lambda.FunctionName} is part of TraceStack`);
return;
}
if (!hasSupportedRuntime) {
logger.info(`- ${lambda.FunctionName} has an unsupported runtime`);
await saveFunctionInfo(lambda, "unsupported_runtime");
return;
}

try {
await updateLambda(lambda, arnBase, edgeEndpoint);

logger.info(`✓ Updated ${lambda.FunctionName}`);
await saveFunctionInfo(lambda, "enabled");
} catch (e) {
logger.warn(`✗ Failed to update ${lambda.FunctionName}`, e);
await saveFunctionInfo(lambda, "error");
}
}),
),
);

await releaseLock("auto-trace");
};
1 change: 1 addition & 0 deletions packages/api/src/routes/collector/index.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ app.post("/", async (c) => {
arn: span.invokedArn,
memoryAllocated: span.memoryAllocated,
timeout: span.maxFinishTime - span.started,
traceStatus: 'enabled',
},
true,
);
Expand Down
17 changes: 17 additions & 0 deletions yarn.lock
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5614,6 +5614,7 @@ __metadata:
dynamodb-toolbox: "npm:^1.3.8"
esbuild: "npm:^0.20.1"
hono: "npm:^4.5.10"
p-limit: "npm:^6.1.0"
serverless: "npm:^3.38.0"
serverless-domain-manager: "npm:^7.3.6"
serverless-esbuild: "npm:^1.51.0"
Expand Down Expand Up @@ -14011,6 +14012,15 @@ __metadata:
languageName: node
linkType: hard

"p-limit@npm:^6.1.0":
version: 6.1.0
resolution: "p-limit@npm:6.1.0"
dependencies:
yocto-queue: "npm:^1.1.1"
checksum: 10c0/40af29461206185a81bdc971ed499d97ceb344114fd21420db95debd9c979b6c02d66a41c321246d09245a51e68410e13df92622cc8c0130f87c6bd81a15d777
languageName: node
linkType: hard

"p-locate@npm:^2.0.0":
version: 2.0.0
resolution: "p-locate@npm:2.0.0"
Expand Down Expand Up @@ -18348,6 +18358,13 @@ __metadata:
languageName: node
linkType: hard

"yocto-queue@npm:^1.1.1":
version: 1.1.1
resolution: "yocto-queue@npm:1.1.1"
checksum: 10c0/cb287fe5e6acfa82690acb43c283de34e945c571a78a939774f6eaba7c285bacdf6c90fbc16ce530060863984c906d2b4c6ceb069c94d1e0a06d5f2b458e2a92
languageName: node
linkType: hard

"yoctocolors-cjs@npm:^2.1.2":
version: 2.1.2
resolution: "yoctocolors-cjs@npm:2.1.2"
Expand Down

0 comments on commit 9f38f0b

Please sign in to comment.