Upgrade dd-trace, add APM tracing around key commands, fix tags should be attached to root spans (#3243)

server/commands/accountProvisioner.ts

@@ -5,6 +5,7 @@ import {
   EmailAuthenticationRequiredError,
   AuthenticationProviderDisabledError,
 } from "@server/errors";
+import { APM } from "@server/logging/tracing";
 import mailer from "@server/mailer";
 import { Collection, Team, User } from "@server/models";
 import teamCreator from "./teamCreator";
@@ -43,7 +44,7 @@ export type AccountProvisionerResult = {
   isNewUser: boolean;
 };
 
-export default async function accountProvisioner({
+async function accountProvisioner({
   ip,
   user: userParams,
   team: teamParams,
@@ -142,3 +143,5 @@ export default async function accountProvisioner({
     throw err;
   }
 }
+
+export default APM.traceFunction({})(accountProvisioner);

server/commands/collectionExporter.ts

@@ -1,7 +1,8 @@
+import { APM } from "@server/logging/tracing";
 import { Collection, Event, Team, User, FileOperation } from "@server/models";
 import { getAWSKeyForFileOp } from "@server/utils/s3";
 
-export default async function collectionExporter({
+async function collectionExporter({
   collection,
   team,
   user,
@@ -43,3 +44,5 @@ export default async function collectionExporter({
 
   return fileOperation;
 }
+
+export default APM.traceFunction({})(collectionExporter);

server/commands/collectionImporter.ts

@@ -6,6 +6,7 @@ import invariant from "invariant";
 import { values, keys } from "lodash";
 import { v4 as uuidv4 } from "uuid";
 import Logger from "@server/logging/logger";
+import { APM } from "@server/logging/tracing";
 import { Attachment, Event, Document, Collection, User } from "@server/models";
 import { parseOutlineExport, Item } from "@server/utils/zip";
 import { FileImportError } from "../errors";
@@ -17,7 +18,7 @@ type FileWithPath = File & {
   path: string;
 };
 
-export default async function collectionImporter({
+async function collectionImporter({
   file,
   type,
   user,
@@ -198,3 +199,5 @@ export default async function collectionImporter({
     attachments: values(attachments),
   };
 }
+
+export default APM.traceFunction({})(collectionImporter);

server/commands/documentImporter.ts

@@ -8,6 +8,7 @@ import TurndownService from "turndown";
 import utf8 from "utf8";
 import { MAX_TITLE_LENGTH } from "@shared/constants";
 import parseTitle from "@shared/utils/parseTitle";
+import { APM } from "@server/logging/tracing";
 import { User } from "@server/models";
 import dataURItoBuffer from "@server/utils/dataURItoBuffer";
 import { deserializeFilename } from "@server/utils/fs";
@@ -141,7 +142,7 @@ async function confluenceToMarkdown(file): Promise<string> {
   return html.replace(/<br>/g, " \\n ");
 }
 
-export default async function documentImporter({
+async function documentImporter({
   file,
   user,
   ip,
@@ -212,3 +213,5 @@ export default async function documentImporter({
     title,
   };
 }
+
+export default APM.traceFunction({})(documentImporter);

server/commands/documentMover.ts

@@ -1,6 +1,7 @@
 import invariant from "invariant";
 import { Transaction } from "sequelize";
 import { sequelize } from "@server/database/sequelize";
+import { APM } from "@server/logging/tracing";
 import {
   User,
   Document,
@@ -62,7 +63,7 @@ type Result = {
   collectionChanged: boolean;
 };
 
-export default async function documentMover({
+async function documentMover({
   user,
   document,
   collectionId,
@@ -238,3 +239,5 @@ export default async function documentMover({
   // we need to send all updated models back to the client
   return result;
 }
+
+export default APM.traceFunction({})(documentMover);

server/commands/teamCreator.ts

@@ -1,5 +1,6 @@
 import invariant from "invariant";
 import Logger from "@server/logging/logger";
+import { APM } from "@server/logging/tracing";
 import { Team, AuthenticationProvider } from "@server/models";
 import { isDomainAllowed } from "@server/utils/authentication";
 import { generateAvatarUrl } from "@server/utils/avatars";
@@ -22,7 +23,7 @@ type Props = {
   };
 };
 
-export default async function teamCreator({
+async function teamCreator({
   name,
   domain,
   subdomain,
@@ -125,3 +126,5 @@ export default async function teamCreator({
     isNewTeam: true,
   };
 }
+
+export default APM.traceFunction({})(teamCreator);

server/commands/teamPermanentDeleter.ts

@@ -1,6 +1,7 @@
 import { Transaction } from "sequelize";
 import { sequelize } from "@server/database/sequelize";
 import Logger from "@server/logging/logger";
+import { APM } from "@server/logging/tracing";
 import {
   ApiKey,
   Attachment,
@@ -20,7 +21,7 @@ import {
   Share,
 } from "@server/models";
 
-export default async function teamPermanentDeleter(team: Team) {
+async function teamPermanentDeleter(team: Team) {
   if (!team.deletedAt) {
     throw new Error(
       `Cannot permanently delete ${team.id} team. Please delete it and try again.`
@@ -203,3 +204,5 @@ export default async function teamPermanentDeleter(team: Team) {
     throw err;
   }
 }
+
+export default APM.traceFunction({})(teamPermanentDeleter);

server/index.ts

@@ -1,7 +1,7 @@
 /* eslint-disable import/order */
 import env from "./env";
 
-import "./tracing"; // must come before importing any instrumented module
+import "./logging/tracing"; // must come before importing any instrumented module
 
 import http from "http";
 import https from "https";

server/logging/tracing.ts

@@ -0,0 +1,20 @@
+import { init, tracer } from "@theo.gravity/datadog-apm";
+
+export * as APM from "@theo.gravity/datadog-apm";
+
+// If the DataDog agent is installed and the DD_API_KEY environment variable is
+// in the environment then we can safely attempt to start the DD tracer
+if (process.env.DD_API_KEY) {
+  init(
+    {
+      // SOURCE_COMMIT is used by Docker Hub
+      // SOURCE_VERSION is used by Heroku
+      version: process.env.SOURCE_COMMIT || process.env.SOURCE_VERSION,
+    },
+    {
+      useMock: process.env.NODE_ENV === "test",
+    }
+  );
+}
+
+export default tracer;

server/middlewares/authentication.ts

@@ -1,6 +1,6 @@
 import { Next } from "koa";
+import tracer, { APM } from "@server/logging/tracing";
 import { User, Team, ApiKey } from "@server/models";
-import tracer from "@server/tracing";
 import { getUserForJWT } from "@server/utils/jwt";
 import { AuthenticationError, UserSuspendedError } from "../errors";
 import { ContextWithState } from "../types";
@@ -100,12 +100,14 @@ export default function auth(
       ctx.state.user = user;
 
       if (tracer) {
-        const span = tracer.scope().active();
-        if (span !== null) {
-          span.setTag("request.userId", user.id);
-          span.setTag("request.teamId", user.teamId);
-          span.setTag("request.authType", ctx.state.authType);
-        }
+        APM.addTags(
+          {
+            "request.userId": user.id,
+            "request.teamId": user.teamId,
+            "request.authType": ctx.state.authType,
+          },
+          APM.getRootSpanFromRequestContext(ctx)
+        );
       }
     }
 

server/tracing.ts

@@ -1,13 +0,0 @@
-import tracer from "dd-trace";
-
-// If the DataDog agent is installed and the DD_API_KEY environment variable is
-// in the environment then we can safely attempt to start the DD tracer
-if (process.env.DD_API_KEY) {
-  tracer.init({
-    // SOURCE_COMMIT is used by Docker Hub
-    // SOURCE_VERSION is used by Heroku
-    version: process.env.SOURCE_COMMIT || process.env.SOURCE_VERSION,
-  });
-}
-
-export default tracer;