Skip to content

Commit 9f38f0b

Browse files
tschoffelentschoffelen
committed
feat: save auto trace status
1 parent 9de7a8d commit 9f38f0b

File tree

4 files changed

+128
-64
lines changed

4 files changed

+128
-64
lines changed

‎packages/api/package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,8 @@
1616
"@includable/serverless-middleware": "^2.2.0",
1717
"date-fns": "^3.6.0",
1818
"dynamodb-toolbox": "^1.3.8",
19-
"hono": "^4.5.10"
19+
"hono": "^4.5.10",
20+
"p-limit": "^6.1.0"
2021
},
2122
"devDependencies": {
2223
"esbuild": "^0.20.1",

‎packages/api/src/events/auto-trace.js

Lines changed: 108 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@ import {
77
ApiGatewayV2Client,
88
GetApisCommand,
99
} from "@aws-sdk/client-apigatewayv2";
10+
import pLimit from "p-limit";
1011
import { acquireLock, releaseLock } from "../lib/locks";
1112
import Logger from "../lib/logger";
13+
import { put } from "../lib/database";
1214

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

56+
const updateLambda = async (lambda, arnBase, edgeEndpoint) => {
57+
const updateFunctionConfigurationCommand =
58+
new UpdateFunctionConfigurationCommand({
59+
FunctionName: lambda.FunctionName,
60+
Layers: [
61+
...(lambda.Layers || [])
62+
.map((layer) => layer.Arn)
63+
.filter((arn) => !arn.startsWith(arnBase)),
64+
process.env.LAMBDA_LAYER_ARN,
65+
],
66+
Environment: {
67+
...(lambda.Environment || {}),
68+
Variables: {
69+
...(lambda.Environment?.Variables || {}),
70+
AUTO_TRACE_HOST: edgeEndpoint,
71+
AWS_LAMBDA_EXEC_WRAPPER: lambdaExecWrapper,
72+
},
73+
},
74+
});
75+
76+
const res = await new LambdaClient().send(updateFunctionConfigurationCommand);
77+
logger.info(res);
78+
};
79+
80+
const saveFunctionInfo = async (lambda, traceStatus) => {
81+
await put(
82+
{
83+
pk: `function#${process.env.AWS_REGION}#${lambda.FunctionName}`,
84+
sk: `function#${process.env.AWS_REGION}`,
85+
name: lambda.FunctionName,
86+
type: "function",
87+
runtime: lambda.Runtime,
88+
region: process.env.AWS_REGION,
89+
arn: lambda.FunctionArn,
90+
memoryAllocated: lambda.MemorySize,
91+
timeout: lambda.Timeout,
92+
traceStatus,
93+
},
94+
true,
95+
);
96+
};
97+
5498
export const autoTrace = async () => {
5599
// Check if we know our Lambda Layer ARN
56100
const arn = process.env.LAMBDA_LAYER_ARN;
@@ -71,69 +115,70 @@ export const autoTrace = async () => {
71115

72116
// List all the lambda functions in the AWS account
73117
const lambdas = await getAccountLambdas();
74-
75-
// Find all lambdas that need a layer added
76-
let lambdasWithoutLayer = lambdas.filter((lambda) => {
77-
const layers = lambda.Layers || [];
78-
const envVars = lambda.Environment?.Variables || {};
79-
80-
const isTraceStack = envVars.LAMBDA_LAYER_ARN === arn;
81-
const isUpdating = lambda.LastUpdateStatus === "InProgress";
82-
const hasDisableEnvVar = envVars.AUTO_TRACE_EXCLUDE;
83-
const hasOtherWrapper =
84-
envVars.AWS_LAMBDA_EXEC_WRAPPER &&
85-
envVars.AWS_LAMBDA_EXEC_WRAPPER !== lambdaExecWrapper;
86-
const hasSupportedRuntime = supportedRuntimes.includes(lambda.Runtime);
87-
const hasLayer = layers.find(({ Arn }) => Arn.startsWith(arnBase));
88-
const hasUpdate = layers.find(
89-
({ Arn }) => Arn.startsWith(arnBase) && Arn !== arn,
90-
);
91-
92-
return (
93-
(!hasLayer || hasUpdate) &&
94-
!hasDisableEnvVar &&
95-
!hasOtherWrapper &&
96-
!isUpdating &&
97-
!isTraceStack &&
98-
hasSupportedRuntime
99-
);
100-
});
101-
102-
logger.info(`Found ${lambdasWithoutLayer.length} lambdas to update`);
103-
104-
for (const lambda of lambdasWithoutLayer) {
105-
try {
106-
const updateFunctionConfigurationCommand =
107-
new UpdateFunctionConfigurationCommand({
108-
FunctionName: lambda.FunctionName,
109-
Layers: [
110-
...(lambda.Layers || [])
111-
.map((layer) => layer.Arn)
112-
.filter((arn) => !arn.startsWith(arnBase)),
113-
process.env.LAMBDA_LAYER_ARN,
114-
],
115-
Environment: {
116-
...(lambda.Environment || {}),
117-
Variables: {
118-
...(lambda.Environment?.Variables || {}),
119-
AUTO_TRACE_HOST: edgeEndpoint,
120-
AWS_LAMBDA_EXEC_WRAPPER: lambdaExecWrapper,
121-
},
122-
},
123-
});
124-
125-
const res = await new LambdaClient().send(
126-
updateFunctionConfigurationCommand,
127-
);
128-
logger.info(res);
129-
130-
logger.info(`✓ Updated ${lambda.FunctionName}`);
131-
// TODO: save function info in DynamoDB
132-
} catch (e) {
133-
logger.warn(`✗ Failed to update ${lambda.FunctionName}`);
134-
logger.warn(e);
135-
}
136-
}
118+
logger.info(`Found ${lambdas.length} lambdas in the account`);
119+
120+
// Update qualifying lambdas
121+
const limit = pLimit(4);
122+
await Promise.all(
123+
lambdas.map((lambda) =>
124+
limit(async () => {
125+
const layers = lambda.Layers || [];
126+
const envVars = lambda.Environment?.Variables || {};
127+
128+
const isTraceStack = envVars.LAMBDA_LAYER_ARN === arn;
129+
const isUpdating = lambda.LastUpdateStatus === "InProgress";
130+
const hasDisableEnvVar = envVars.AUTO_TRACE_EXCLUDE;
131+
const hasOtherWrapper =
132+
envVars.AWS_LAMBDA_EXEC_WRAPPER &&
133+
envVars.AWS_LAMBDA_EXEC_WRAPPER !== lambdaExecWrapper;
134+
const hasSupportedRuntime = supportedRuntimes.includes(lambda.Runtime);
135+
const hasLayer = layers.find(({ Arn }) => Arn.startsWith(arnBase));
136+
const hasUpdate = layers.find(
137+
({ Arn }) => Arn.startsWith(arnBase) && Arn !== arn,
138+
);
139+
140+
if (hasLayer && !hasUpdate) {
141+
logger.info(`- ${lambda.FunctionName} already has the latest layer`);
142+
await saveFunctionInfo(lambda, "enabled");
143+
return;
144+
}
145+
if (hasDisableEnvVar) {
146+
logger.info(`- ${lambda.FunctionName} has AUTO_TRACE_EXCLUDE`);
147+
await saveFunctionInfo(lambda, "excluded");
148+
return;
149+
}
150+
if (hasOtherWrapper) {
151+
logger.info(`- ${lambda.FunctionName} has a custom wrapper`);
152+
await saveFunctionInfo(lambda, "custom_wrapper");
153+
return;
154+
}
155+
if (isUpdating) {
156+
logger.info(`- ${lambda.FunctionName} is already updating`);
157+
await saveFunctionInfo(lambda, "update_in_progress");
158+
return;
159+
}
160+
if (isTraceStack) {
161+
logger.info(`- ${lambda.FunctionName} is part of TraceStack`);
162+
return;
163+
}
164+
if (!hasSupportedRuntime) {
165+
logger.info(`- ${lambda.FunctionName} has an unsupported runtime`);
166+
await saveFunctionInfo(lambda, "unsupported_runtime");
167+
return;
168+
}
169+
170+
try {
171+
await updateLambda(lambda, arnBase, edgeEndpoint);
172+
173+
logger.info(`✓ Updated ${lambda.FunctionName}`);
174+
await saveFunctionInfo(lambda, "enabled");
175+
} catch (e) {
176+
logger.warn(`✗ Failed to update ${lambda.FunctionName}`, e);
177+
await saveFunctionInfo(lambda, "error");
178+
}
179+
}),
180+
),
181+
);
137182

138183
await releaseLock("auto-trace");
139184
};

‎packages/api/src/routes/collector/index.js

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ app.post("/", async (c) => {
7272
arn: span.invokedArn,
7373
memoryAllocated: span.memoryAllocated,
7474
timeout: span.maxFinishTime - span.started,
75+
traceStatus: 'enabled',
7576
},
7677
true,
7778
);

‎yarn.lock

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5614,6 +5614,7 @@ __metadata:
56145614
dynamodb-toolbox: "npm:^1.3.8"
56155615
esbuild: "npm:^0.20.1"
56165616
hono: "npm:^4.5.10"
5617+
p-limit: "npm:^6.1.0"
56175618
serverless: "npm:^3.38.0"
56185619
serverless-domain-manager: "npm:^7.3.6"
56195620
serverless-esbuild: "npm:^1.51.0"
@@ -14011,6 +14012,15 @@ __metadata:
1401114012
languageName: node
1401214013
linkType: hard
1401314014

14015+
"p-limit@npm:^6.1.0":
14016+
version: 6.1.0
14017+
resolution: "p-limit@npm:6.1.0"
14018+
dependencies:
14019+
yocto-queue: "npm:^1.1.1"
14020+
checksum: 10c0/40af29461206185a81bdc971ed499d97ceb344114fd21420db95debd9c979b6c02d66a41c321246d09245a51e68410e13df92622cc8c0130f87c6bd81a15d777
14021+
languageName: node
14022+
linkType: hard
14023+
1401414024
"p-locate@npm:^2.0.0":
1401514025
version: 2.0.0
1401614026
resolution: "p-locate@npm:2.0.0"
@@ -18348,6 +18358,13 @@ __metadata:
1834818358
languageName: node
1834918359
linkType: hard
1835018360

18361+
"yocto-queue@npm:^1.1.1":
18362+
version: 1.1.1
18363+
resolution: "yocto-queue@npm:1.1.1"
18364+
checksum: 10c0/cb287fe5e6acfa82690acb43c283de34e945c571a78a939774f6eaba7c285bacdf6c90fbc16ce530060863984c906d2b4c6ceb069c94d1e0a06d5f2b458e2a92
18365+
languageName: node
18366+
linkType: hard
18367+
1835118368
"yoctocolors-cjs@npm:^2.1.2":
1835218369
version: 2.1.2
1835318370
resolution: "yoctocolors-cjs@npm:2.1.2"

0 commit comments

Comments
 (0)