feat: move trace storage to S3

tschoffelen · tschoffelen · commit 2c4b9fd9c52f · 2025-03-30T10:44:20.000+01:00
diff --git a/packages/api/package.json b/packages/api/package.json
@@ -14,9 +14,11 @@
     "@aws-sdk/client-cloudwatch-logs": "^3.658.1",
     "@aws-sdk/client-dynamodb": "^3.637.0",
     "@aws-sdk/client-lambda": "^3.637.0",
+    "@aws-sdk/client-s3": "^3.777.0",
     "@aws-sdk/lib-dynamodb": "^3.637.0",
     "@hono/node-server": "^1.12.2",
     "@includable/serverless-middleware": "^2.2.0",
+    "@thi.ng/ksuid": "^3.2.83",
     "bcryptjs": "^2.4.3",
     "date-fns": "^3.6.0",
     "dynamodb-toolbox": "^1.3.8",
diff --git a/packages/api/serverless.yml b/packages/api/serverless.yml
@@ -8,7 +8,7 @@ plugins:
 
 params:
   default:
-    DATA_RETENTION_DAYS: ${env:DATA_RETENTION_DAYS, 30}
+    DATA_RETENTION_DAYS: ${env:DATA_RETENTION_DAYS, 14}
     CUSTOM_DOMAIN: ${env:CUSTOM_DOMAIN, ""}
     HAS_CUSTOM_DOMAIN: ${strToBool(${env:HAS_CUSTOM_DOMAIN, "false"})}
     TRACER_TOKEN: ${env:TRACER_TOKEN, ""}
@@ -22,6 +22,7 @@ custom:
     endpointType: regional
     apiType: http
   logRetentionInDays: 14
+  bucketName: ${self:service}-storage-${sls:stage} # TODO: make this a unique name
 
 package:
   individually: true
@@ -52,21 +53,10 @@ provider:
     TRACER_TOKEN: ${param:TRACER_TOKEN}
     LAMBDA_LAYER_ARN: !Ref TracerLambdaLayer
     AUTO_TRACE_EXCLUDE: 1
+    STORAGE_BUCKET_NAME: ${self:custom.bucketName}
   httpApi:
     payload: "2.0"
     cors: true
-    # authorizers:
-    #   WebAuthorizer:
-    #     identitySource: $request.header.Authorization
-    #     issuerUrl:
-    #       Fn::Join:
-    #         - ""
-    #         - - "https://cognito-idp."
-    #           - "${opt:region, self:provider.region}"
-    #           - ".amazonaws.com/"
-    #           - Ref: CognitoUserPool
-    #     audience:
-    #       Ref: CognitoUserPoolClient
   lambdaHashingVersion: 20201221
   iam:
     role:
@@ -99,6 +89,15 @@ provider:
             - apigateway:GET
           Resource:
             - "arn:aws:apigateway:${aws:region}::/apis"
+        - Effect: Allow
+          Action:
+            - s3:GetObject
+            - s3:PutObject
+            - s3:DeleteObject
+            - s3:ListBucket
+          Resource:
+            - arn:aws:s3:::${self:custom.bucketName}
+            - arn:aws:s3:::${self:custom.bucketName}/*
 
 layers:
   tracer:
@@ -109,12 +108,14 @@ layers:
       - nodejs16.x
       - nodejs18.x
       - nodejs20.x
+      - nodejs22.x
 
 functions:
   main:
     handler: src/index.handler
     timeout: 900
-    memorySize: 2048
+    memorySize: 1024
+    reservedConcurrency: 1
     url: true
     events:
       - httpApi:
@@ -127,3 +128,4 @@ functions:
 
 resources:
   - ${file(./src/infrastructure/database/table.yml)}
+  - ${file(./src/infrastructure/bucket.yml)}
diff --git a/packages/api/src/infrastructure/bucket.yml b/packages/api/src/infrastructure/bucket.yml
@@ -0,0 +1,10 @@
+Resources:
+  StorageBucket:
+    Type: AWS::S3::Bucket
+    Properties:
+      BucketName: ${self:custom.bucketName}
+      LifecycleConfiguration:
+        Rules:
+          - Id: ExpireObjectsRule
+            Status: Enabled
+            ExpirationInDays: ${param:DATA_RETENTION_DAYS}
diff --git a/packages/api/src/lib/invocations.js b/packages/api/src/lib/invocations.js
@@ -1,4 +1,4 @@
-import { put } from "./database";
+import { store } from "./storage";
 
 const getStatusCodeFromSpan = (span) => {
   if (!span.return_value) return null;
@@ -36,16 +36,11 @@ const getResultSummaryFromSpan = (span) => {
   return mainStatus;
 };
 
-export const saveInvocation = async (span) => {
-  return put(
-    {
-      ...span,
-      statusCode: getStatusCodeFromSpan(span),
-      resultSummary: getResultSummaryFromSpan(span),
-      pk: `function#${span.region}#${span.name}`,
-      sk: `invocation#${span.started}#${span.id}`,
-      type: "invocation",
-    },
-    true,
-  );
+export const saveInvocation = async (span, allSpans) => {
+  return store(["invocations", span.region, span.name, span.started, span.id], {
+    ...span,
+    statusCode: getStatusCodeFromSpan(span),
+    resultSummary: getResultSummaryFromSpan(span),
+    spans: allSpans,
+  });
 };
diff --git a/packages/api/src/lib/storage.js b/packages/api/src/lib/storage.js
@@ -0,0 +1,19 @@
+import { PutObjectCommand, S3Client } from "@aws-sdk/client-s3";
+
+const s3 = new S3Client();
+
+export const store = async (key, data) => {
+  const datePrefix = new Date().toISOString().split("T")[0];
+  const bucketName = process.env.STORAGE_BUCKET_NAME;
+
+  if (Array.isArray(key)) key = key.join("/");
+
+  const params = {
+    Bucket: bucketName,
+    Key: `${datePrefix}/${key}`,
+    Body: JSON.stringify(data),
+    ContentType: "application/json",
+  };
+
+  return s3.send(new PutObjectCommand(params));
+};
diff --git a/packages/api/src/routes/collector/index.js b/packages/api/src/routes/collector/index.js
@@ -1,139 +1,113 @@
 import { Hono } from "hono";
+import { defULID } from "@thi.ng/ksuid";
 
-import { getExpiryTime, put, update } from "../../lib/database";
 import { saveHourlyStat } from "../../lib/stats";
 import { getErrorKey } from "../../lib/errors";
-import { groupSpans } from "../../lib/spans";
 import { saveInvocation } from "../../lib/invocations";
+import { store } from "../../lib/storage";
 
 const app = new Hono();
 
+const id = defULID();
+const getId = () => {
+  return id.next();
+};
+
+// Holds transactions for which we don't know the invocation ID yet
+// 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();
 
-  const groupedItems = {};
-
   for (const span of body) {
     if (process.env.TRACER_TOKEN && span.token !== process.env.TRACER_TOKEN) {
       console.log(`Invalid token: ${span.token}`);
       continue;
     }
 
-    const pk = `transaction#${span.transactionId || span.transaction_id}`;
-    if (!groupedItems[pk]) groupedItems[pk] = [];
+    const groupKey = span.transactionId || span.transaction_id;
+    if (!transactionCache[groupKey]) transactionCache[groupKey] = [];
 
-    if (span.type === "log") {
-      // save log span
-      groupedItems[pk].push({
-        ...span,
-        transactionId: span.transactionId || span.transaction_id,
-        type: "log",
-      });
-      continue;
-    }
+    transactionCache[groupKey].push(span);
+  }
+
+  // Check transactions cache to see if there's any transactions we can flush
+  for (const [transactionId, spans] of Object.entries(transactionCache)) {
+    const invocationEndedSpan = spans.find(
+      (span) =>
+        span.type === "function" && span.ended && !span.id.includes("_started"),
+    );
+
+    if (!invocationEndedSpan) {
+      console.log(
+        `No invocation ended span found for transaction ${spans[0].transactionId}`,
+      );
+
+      // TODO: if we are close to running out of time, we should flush the transaction cache anyway
 
-    // ignore started spans and enrichments
-    if (span.id.endsWith("_started") || span.type === "enrichment") {
       continue;
+    } else {
+      console.log(
+        "Flushing transaction cache for",
+        invocationEndedSpan.transactionId,
+      );
     }
 
-    // save transaction span
-    groupedItems[pk].push({
-      ...span,
-      type: "span",
-      spanType: span.type,
-    });
-
-    if (
-      span.type === "function" &&
-      span.ended &&
-      !span.id.includes("_started")
-    ) {
-      // save function invocation details
-      await saveInvocation(span);
-
-      // save error
-      if (span.error) {
-        const errorKey = getErrorKey(span.error);
-        await update({
-          Key: {
-            pk: `function#${span.region}#${span.name}`,
-            sk: `error#${errorKey}`,
-          },
-          UpdateExpression: `SET #error = :error, #lastInvocation = :lastInvocation, #lastSeen = :lastSeen, #expires = :expires, #type = :type, #name = :name, #region = :region`,
-          ExpressionAttributeValues: {
-            ":error": span.error,
-            ":lastInvocation": `${span.started}/${span.id}`,
-            ":lastSeen": new Date(span.ended).toISOString(),
-            ":expires": getExpiryTime(),
-            ":type": "error",
-            ":name": span.name,
-            ":region": span.region,
-          },
-          ExpressionAttributeNames: {
-            "#error": "error",
-            "#lastInvocation": "lastInvocation",
-            "#lastSeen": "lastSeen",
-            "#expires": "_expires",
-            "#type": "type",
-            "#name": "name",
-            "#region": "region",
-          },
-        });
-        await saveHourlyStat(span.region, span.name + ".error." + errorKey, 1);
-      }
-
-      // save function meta data
-      try {
-        await update({
-          Key: {
-            pk: `function#${span.region}#${span.name}`,
-            sk: `function#${span.region}`,
-          },
-          UpdateExpression: `SET lastInvocation = :lastInvocation, memoryAllocated = :memoryAllocated, #timeout = :timeout, traceStatus = :traceStatus, #expires = :expires`,
-          ExpressionAttributeValues: {
-            ":lastInvocation": span.started,
-            ":memoryAllocated": span.memoryAllocated,
-            ":timeout": span.maxFinishTime - span.started,
-            ":traceStatus": "enabled",
-            ":expires": getExpiryTime(),
-          },
-          ExpressionAttributeNames: {
-            "#timeout": "timeout",
-            "#expires": "_expires",
-          },
-        });
-      } catch (e) {
-        console.log(e);
-      }
-
-      // save stats
-      const duration = span.ended - span.started;
-      await saveHourlyStat(span.region, span.name + ".invocations", 1);
-      await saveHourlyStat(span.region, span.name + ".duration", duration);
-      await saveHourlyStat("global", "invocations", 1);
-      if (span.error) {
-        await saveHourlyStat(span.region, span.name + ".errors", 1);
-        await saveHourlyStat("global", "errors", 1);
-      }
+    // save function invocation details
+    await saveInvocation(invocationEndedSpan, spans);
+
+    // const duration = invocationEndedSpan.ended - invocationEndedSpan.started;
+    await saveHourlyStat(
+      invocationEndedSpan.region,
+      invocationEndedSpan.name + ".invocations",
+      1,
+    );
+    // await saveHourlyStat(
+    //   invocationEndedSpan.region,
+    //   invocationEndedSpan.name + ".duration",
+    //   duration,
+    // );
+    await saveHourlyStat("global", "invocations", 1);
+    if (invocationEndedSpan.error) {
+      await saveHourlyStat(
+        invocationEndedSpan.region,
+        invocationEndedSpan.name + ".errors",
+        1,
+      );
+      await saveHourlyStat("global", "errors", 1);
     }
-  }
 
-  const itemsToSave = [];
-  for (let [pk, items] of Object.entries(groupedItems)) {
-    items = groupSpans(items);
-    while (items.length) {
-      const chunk = items.splice(0, 100);
-      itemsToSave.push({
-        pk,
-        sk: `spans#${chunk[0].started || chunk[0].sending_time}#${chunk[0].id}`,
-        type: "spans",
-        spans: chunk,
-      });
+    // save error
+    if (invocationEndedSpan.error) {
+      const errorKey = getErrorKey(invocationEndedSpan.error);
+      await store(
+        [
+          "errors",
+          invocationEndedSpan.region,
+          invocationEndedSpan.name,
+          errorKey,
+          getId(),
+        ],
+        {
+          error: invocationEndedSpan.error,
+          lastInvocation: `${invocationEndedSpan.started}/${invocationEndedSpan.id}`,
+          lastSeen: new Date(invocationEndedSpan.ended).toISOString(),
+          type: "error",
+          name: invocationEndedSpan.name,
+          region: invocationEndedSpan.region,
+        },
+      );
+      await saveHourlyStat(
+        invocationEndedSpan.region,
+        invocationEndedSpan.name + ".error." + errorKey,
+        1,
+      );
     }
-  }
 
-  await Promise.all(itemsToSave.map((item) => put(item), true));
+    // Delete the transaction from the cache
+    delete transactionCache[transactionId];
+  }
 
   return c.json({ success: true });
 });
diff --git a/yarn.lock b/yarn.lock
