diff --git a/packages/api/serverless.yml b/packages/api/serverless.yml index da45f77..601c81a 100644 --- a/packages/api/serverless.yml +++ b/packages/api/serverless.yml @@ -10,6 +10,7 @@ params: DATA_RETENTION_DAYS: ${env:DATA_RETENTION_DAYS, 30} CUSTOM_DOMAIN: ${env:CUSTOM_DOMAIN, ""} HAS_CUSTOM_DOMAIN: ${strToBool(${env:HAS_CUSTOM_DOMAIN, "false"})} + TRACER_TOKEN: ${env:TRACER_TOKEN, ""} custom: customDomain: @@ -45,6 +46,7 @@ provider: NODE_ENV: ${sls:stage} TABLE_NAME: ${self:service}-${sls:stage} CUSTOM_DOMAIN: ${param:CUSTOM_DOMAIN} + TRACER_TOKEN: ${param:TRACER_TOKEN} LAMBDA_LAYER_ARN: !Ref TracerLambdaLayer AUTO_TRACE_EXCLUDE: 1 httpApi: diff --git a/packages/api/src/events/auto-trace.js b/packages/api/src/events/auto-trace.js index 6d09c8f..dc841d7 100644 --- a/packages/api/src/events/auto-trace.js +++ b/packages/api/src/events/auto-trace.js @@ -7,6 +7,7 @@ import { ApiGatewayV2Client, GetApisCommand, } from "@aws-sdk/client-apigatewayv2"; +import { acquireLock, releaseLock } from "../lib/locks"; const supportedRuntimes = ["nodejs16.x", "nodejs18.x", "nodejs20.x"]; const lambdaExecWrapper = "/opt/nodejs/tracer_wrapper"; @@ -57,6 +58,13 @@ export const autoTrace = async () => { // Get our API Gateway endpoint for the collector const edgeEndpoint = await getApiEndpoint(); + // Make sure we lock so that only one process is updating lambdas + const lockAcquired = await acquireLock("auto-trace"); + if (!lockAcquired) { + console.log("Lock not acquired, skipping"); + return; + } + // List all the lambda functions in the AWS account const lambdas = await getAccountLambdas(); @@ -118,4 +126,6 @@ export const autoTrace = async () => { // TODO: save function info in DynamoDB } + + await releaseLock("auto-trace"); }; diff --git a/packages/api/src/index.js b/packages/api/src/index.js index 80eaf4e..a06a22d 100644 --- a/packages/api/src/index.js +++ b/packages/api/src/index.js @@ -5,11 +5,13 @@ import fs from "fs"; import collector from "./routes/collector"; import explore from "./routes/explore"; +import autoTraceRoute from "./routes/auto-trace"; import { autoTrace } from "./events/auto-trace"; const app = new Hono(); app.route("/api/spans", collector); app.route("/api/explore", explore); +app.route("/api/auto-trace", autoTraceRoute); let html = ""; app.use("/assets/*", serveStatic({ root: "./dist" })); diff --git a/packages/api/src/lib/database.js b/packages/api/src/lib/database.js index b72f28c..817815e 100644 --- a/packages/api/src/lib/database.js +++ b/packages/api/src/lib/database.js @@ -1,5 +1,6 @@ import { DynamoDBClient } from "@aws-sdk/client-dynamodb"; import { + DeleteCommand, DynamoDBDocumentClient, PutCommand, QueryCommand, @@ -22,7 +23,7 @@ export const time = () => Math.floor(Date.now() / 1000); export const getExpiryTime = () => time() + 86400 * Number(process.env.DATA_RETENTION_DAYS); -export const put = async (item, expires = false) => { +export const put = async (item, expires = false, options = {}) => { item = { ...item, _created: time(), @@ -37,6 +38,7 @@ export const put = async (item, expires = false) => { new PutCommand({ TableName: process.env.TABLE_NAME, Item: item, + ...options }), ); }; @@ -63,6 +65,15 @@ export const update = async ( ); }; +export const deleteItem = async (Key) => { + return await dynamo.send( + new DeleteCommand({ + TableName: process.env.TABLE_NAME, + Key, + }), + ); +} + export const queryAll = async ( /** @type {Omit} */ params, ) => { diff --git a/packages/api/src/lib/locks.js b/packages/api/src/lib/locks.js new file mode 100644 index 0000000..1754b42 --- /dev/null +++ b/packages/api/src/lib/locks.js @@ -0,0 +1,39 @@ +import { deleteItem, put } from "./database"; + +export const acquireLock = async (key, ttl = 900) => { + const lockKey = `lock#${key}`; + const expires = Math.floor(Date.now() / 1000) + ttl; + try { + await put( + { + pk: lockKey, + sk: lockKey, + _expires: expires, + }, + true, + { + ConditionExpression: "attribute_not_exists(pk)", + }, + ); + return true; + } catch (error) { + if (error.name === "ConditionalCheckFailedException") { + return false; + } + throw error; + } +}; + +export const releaseLock = async (key) => { + const lockKey = `lock#${key}`; + try { + await deleteItem({ + pk: lockKey, + sk: lockKey, + }); + return true; + } catch (e) { + console.log(e); + return false; + } +}; diff --git a/packages/api/src/routes/auto-trace/index.js b/packages/api/src/routes/auto-trace/index.js new file mode 100644 index 0000000..f83407a --- /dev/null +++ b/packages/api/src/routes/auto-trace/index.js @@ -0,0 +1,17 @@ +import { Hono } from "hono"; +import { autoTrace } from "../../events/auto-trace"; + +const app = new Hono(); + +app.post("/", async (c) => { + const body = await c.req.json(); + if (body.token !== process.env.TRACER_TOKEN) { + return c.json({ error: "Invalid token" }, 401); + } + + await autoTrace(); + + return c.json({ success: true }); +}); + +export default app; diff --git a/packages/api/src/routes/collector/index.js b/packages/api/src/routes/collector/index.js index a65474b..f013705 100644 --- a/packages/api/src/routes/collector/index.js +++ b/packages/api/src/routes/collector/index.js @@ -7,6 +7,10 @@ const app = new Hono(); app.post("/", async (c) => { const body = await c.req.json(); + if (process.env.TRACER_TOKEN && body.token !== process.env.TRACER_TOKEN) { + return c.json({ error: "Invalid token" }, 401); + } + for (const span of body) { console.log(span); diff --git a/packages/deploy-script/src/index.js b/packages/deploy-script/src/index.js index 6f88d77..03103c7 100755 --- a/packages/deploy-script/src/index.js +++ b/packages/deploy-script/src/index.js @@ -2,6 +2,7 @@ import chalk from "chalk"; import boxen from "boxen"; +import crypto from "crypto"; import inquirer from "inquirer"; import degit from "degit"; import child_process from "child_process"; @@ -12,6 +13,9 @@ import { GetApisCommand, } from "@aws-sdk/client-apigatewayv2"; +// TODO: use previous token if it exists +const tracerToken = crypto.randomBytes(16).toString("hex"); + const exec = (command, options = {}) => { const child = child_process.exec(command, { ...options, @@ -96,12 +100,19 @@ await cloner.clone("/tmp/trace-stack"); console.log(chalk.blue("Installing dependencies...")); await exec("yarn install", { cwd: "/tmp/trace-stack" }); -// Write .env file +// Write config files +console.log(chalk.blue("Writing tracer config file...")); +await writeFile( + "/tmp/trace-stack/packages/lambda-layer/config.json", + JSON.stringify({ token: tracerToken }), +); + console.log(chalk.blue("Writing .env file...")); await writeFile( "/tmp/trace-stack/packages/api/.env", `RETENTION_DAYS=${answers.RETENTION_DAYS}\n` + `CUSTOM_DOMAIN=${answers.CUSTOM_DOMAIN}\n` + + `TRACER_TOKEN=${tracerToken}\n` + `HAS_CUSTOM_DOMAIN=${answers.CUSTOM_DOMAIN ? "true" : "false"}\n`, ); @@ -109,12 +120,20 @@ await writeFile( console.log(chalk.blue("Deploying...")); await exec("yarn deploy", { cwd: "/tmp/trace-stack" }); +// Run auto-trace +const endpoint = await getApiEndpoint(); +await fetch(`https://${endpoint}/api/auto-trace`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ token: tracerToken }), +}); + // Create user console.log(chalk.blue("Creating user...")); // TODO: create user account in DDB // Done -const domain = answers.CUSTOM_DOMAIN || (await getApiEndpoint()); +const domain = answers.CUSTOM_DOMAIN || endpoint; console.log( "\n\n" + boxen( diff --git a/packages/lambda-layer/index.js b/packages/lambda-layer/index.js index ab47924..5cfb2f4 100644 --- a/packages/lambda-layer/index.js +++ b/packages/lambda-layer/index.js @@ -1,6 +1,16 @@ +let config; +try { + config = require("./config.json"); +} catch (e) { + config = { + token: "t_0000000000000000", + edgeHost: process.env.AUTO_TRACE_HOST, + }; +} + const tracer = require("@lumigo/tracer")({ - token: "t_0000000000000000", - edgeHost: process.env.AUTO_TRACE_HOST, + token: config.token, + edgeHost: process.env.AUTO_TRACE_HOST || config.edgeHost, }); const { load } = require("./lib/aws/aws-user-function.js");