feat(controlplane): add service-level tracing spans to Connect-RPC handlers (#2771)

Author: thisisnithin

controlplane/.eslintrc

@@ -1,6 +1,6 @@
 {
   "extends": ["eslint-config-unjs", "plugin:require-extensions/recommended"],
-  "plugins": ["require-extensions"],
+  "plugins": ["require-extensions", "local-rules"],
   "rules": {
     "space-before-function-paren": 0,
     "arrow-parens": 0,
@@ -25,5 +25,6 @@
     "no-useless-constructor": 0,
     "unicorn/prefer-ternary": 0,
     "unicorn/no-nested-ternary": 0,
+    "local-rules/no-arrow-in-traced": "error",
   },
 }

controlplane/eslint-local-rules.cjs

@@ -0,0 +1,71 @@
+'use strict';
+
+function isTraced(node) {
+  if (node.decorators && node.decorators.some((d) => d.expression && d.expression.name === 'traced')) {
+    return true;
+  }
+
+  let container = node.parent;
+  let classNode = node;
+  // export default class wraps the ClassDeclaration in ExportDefaultDeclaration
+  if (container && container.type === 'ExportDefaultDeclaration') {
+    classNode = container;
+    container = container.parent;
+  }
+  if (!container || !container.body) {
+    return false;
+  }
+  const siblings = container.body;
+  const idx = siblings.indexOf(classNode);
+  for (let i = idx + 1; i < siblings.length; i++) {
+    const sibling = siblings[i];
+    if (
+      sibling.type === 'ExpressionStatement' &&
+      sibling.expression &&
+      sibling.expression.type === 'CallExpression' &&
+      sibling.expression.callee &&
+      sibling.expression.callee.name === 'traced' &&
+      sibling.expression.arguments.length === 1 &&
+      sibling.expression.arguments[0].name === node.id?.name
+    ) {
+      return true;
+    }
+    if (sibling.type === 'ClassDeclaration' || sibling.type === 'FunctionDeclaration') {
+      break;
+    }
+  }
+  return false;
+}
+
+module.exports = {
+  'no-arrow-in-traced': {
+    meta: {
+      type: 'problem',
+      docs: {
+        description: 'Disallow arrow function class fields in @traced classes',
+      },
+      messages: {
+        noArrowInTraced:
+          'Arrow function class fields are not traced by the @traced decorator. Convert to a regular method.',
+      },
+    },
+    create(context) {
+      return {
+        ClassDeclaration(node) {
+          if (!isTraced(node)) {
+            return;
+          }
+          for (const member of node.body.body) {
+            if (
+              member.type === 'PropertyDefinition' &&
+              member.value &&
+              member.value.type === 'ArrowFunctionExpression'
+            ) {
+              context.report({ node: member, messageId: 'noArrowInTraced' });
+            }
+          }
+        },
+      };
+    },
+  },
+};

controlplane/package.json

@@ -28,6 +28,7 @@
     "lint": "eslint --cache --ext .ts,.mjs,.cjs . && prettier -c src",
     "lint:fix": "eslint --cache --fix --ext .ts,.mjs,.cjs . && pnpm format",
     "format": "prettier -w .",
+    "sentry:spotlight": "spotlight",
     "migrate": "pnpm db:migrate && pnpm ch:migrate",
     "drizzle:up": "drizzle-kit up",
     "ch:down": "dbmate -d \"./clickhouse/migrations\" down",
@@ -106,6 +107,7 @@
     "@bufbuild/protobuf": "^1.9.0",
     "@bufbuild/protoc-gen-es": "^1.9.0",
     "@connectrpc/protoc-gen-connect-es": "^1.4.0",
+    "@spotlightjs/spotlight": "^4.10.0",
     "@types/cookie": "^0.6.0",
     "@types/ejs": "^3.1.5",
     "@types/eslint": "^9.6.1",
@@ -118,6 +120,7 @@
     "del-cli": "^5.1.0",
     "drizzle-kit": "^0.26.2",
     "eslint-config-unjs": "^0.2.1",
+    "eslint-plugin-local-rules": "^3.0.2",
     "eslint-plugin-require-extensions": "^0.1.3",
     "msw": "^2.2.11",
     "pino-pretty": "^10.3.1",

controlplane/src/core/auth-utils.ts

@@ -26,7 +26,7 @@ import { OrganizationGroupRepository } from './repositories/OrganizationGroupRep
 import { DefaultNamespace, NamespaceRepository } from './repositories/NamespaceRepository.js';
 import Keycloak from './services/Keycloak.js';
 import { IPlatformWebhookService } from './webhooks/PlatformWebhookService.js';
-
+import { traced } from './tracing.js';
 export type AuthUtilsOptions = {
   webBaseUrl: string;
   webErrorPath: string;
@@ -711,3 +711,5 @@ export default class AuthUtils {
     return 0;
   }
 }
+
+traced(AuthUtils);

controlplane/src/core/blobstorage/dual.ts

@@ -1,3 +1,4 @@
+import { traced } from '../tracing.js';
 import type { BlobObject, BlobStorage } from './index.js';
 
 /**
@@ -6,6 +7,7 @@ import type { BlobObject, BlobStorage } from './index.js';
  * - Writes and deletes go to both stores concurrently; both must succeed.
  * - Reads try the primary first, falling back to the secondary on failure.
  */
+@traced
 export class DualBlobStorage implements BlobStorage {
   constructor(
     private primary: BlobStorage,

controlplane/src/core/blobstorage/s3.ts

@@ -7,6 +7,7 @@ import {
   PutObjectCommand,
   S3Client,
 } from '@aws-sdk/client-s3';
+import { traced } from '../tracing.js';
 import { BlobNotFoundError, BlobObject, type BlobStorage } from './index.js';
 
 const maxConcurrency = 10; // Maximum number of concurrent operations
@@ -26,6 +27,7 @@ export interface S3BlobStorageConfig {
 /**
  * Stores objects in S3 given an S3Client and a bucket name
  */
+@traced
 export class S3BlobStorage implements BlobStorage {
   private readonly useIndividualDeletes: boolean;
 

controlplane/src/core/build-server.ts

@@ -1,6 +1,7 @@
 import Fastify, { FastifyBaseLogger } from 'fastify';
 import { S3Client } from '@aws-sdk/client-s3';
 import { fastifyConnectPlugin } from '@connectrpc/connect-fastify';
+import * as Sentry from '@sentry/node';
 import { cors, createContextValues } from '@connectrpc/connect';
 import fastifyCors from '@fastify/cors';
 import { pino, stdTimeFunctions, LoggerOptions } from 'pino';
@@ -37,7 +38,13 @@ import { BillingRepository } from './repositories/BillingRepository.js';
 import { BillingService } from './services/BillingService.js';
 import { UserRepository } from './repositories/UserRepository.js';
 import { AIGraphReadmeQueue, createAIGraphReadmeWorker } from './workers/AIGraphReadmeWorker.js';
-import { fastifyLoggerId, createS3ClientConfig, extractS3BucketName, isGoogleCloudStorageUrl } from './util.js';
+import {
+  fastifyLoggerId,
+  sentrySpanId,
+  createS3ClientConfig,
+  extractS3BucketName,
+  isGoogleCloudStorageUrl,
+} from './util.js';
 import { ApiKeyRepository } from './repositories/ApiKeyRepository.js';
 import { createDeleteOrganizationWorker, DeleteOrganizationQueue } from './workers/DeleteOrganizationWorker.js';
 import {
@@ -531,6 +538,16 @@ export default async function build(opts: BuildConfig) {
     keycloakRealm: opts.keycloak.realm,
   });
 
+  // Capture the active Sentry span in preHandler (where OTEL context is still available)
+  // and store it on the request so Connect interceptors can use it as parentSpan.
+  fastify.addHook('preHandler', (req, _reply, done) => {
+    const span = Sentry.getActiveSpan();
+    if (span) {
+      (req.raw as any).__sentrySpan = span;
+    }
+    done();
+  });
+
   // Must be registered after custom fastify routes
   // Because it registers an all-catch route for connect handlers
 
@@ -567,7 +584,16 @@ export default async function build(opts: BuildConfig) {
       cdnBaseUrl: opts.cdnBaseUrl,
     }),
     contextValues(req) {
-      return createContextValues().set<FastifyBaseLogger>({ id: fastifyLoggerId, defaultValue: req.log }, req.log);
+      const values = createContextValues().set<FastifyBaseLogger>(
+        { id: fastifyLoggerId, defaultValue: req.log },
+        req.log,
+      );
+      // Read the parent span captured during the preHandler hook
+      const parentSpan = (req.raw as any).__sentrySpan;
+      if (parentSpan) {
+        values.set({ id: sentrySpanId, defaultValue: undefined }, parentSpan);
+      }
+      return values;
     },
     logLevel: opts.logger.level as pino.LevelWithSilent,
     // Avoid compression for small requests
@@ -578,6 +604,18 @@ export default async function build(opts: BuildConfig) {
     // We go with 32MiB to avoid allocating too much memory for large requests
     writeMaxBytes: 32 * 1024 * 1024,
     acceptCompression: [compressionBrotli, compressionGzip],
+    interceptors: [
+      (next) => (req) => {
+        const parentSpan = req.contextValues?.get({
+          id: sentrySpanId,
+          defaultValue: undefined,
+        });
+        if (parentSpan) {
+          return Sentry.withActiveSpan(parentSpan, () => next(req));
+        }
+        return next(req);
+      },
+    ],
   });
 
   await fastify.register(fastifyGracefulShutdown, {

controlplane/src/core/clickhouse/client/ClickHouseClient.ts

@@ -7,6 +7,7 @@ import pkg from 'stream-json';
 
 import { Observable, Subscriber } from 'rxjs';
 
+import { traced } from '../../tracing.js';
 import { ClickHouseCompressionMethod, ClickHouseDataFormat } from './enums/index.js';
 
 import { ClickHouseClientOptions } from './interfaces/index.js';
@@ -16,6 +17,7 @@ const { Parser } = pkg;
  * ClickHouse Client
  * Most of the code is taken from https://github.com/depyronick/clickhouse-client
  */
+@traced
 export class ClickHouseClient {
   /**
    * ClickHouse Endpoint without path and query

controlplane/src/core/composition/composer.ts

@@ -33,6 +33,7 @@ import { CacheWarmerRepository } from '../repositories/CacheWarmerRepository.js'
 import { NamespaceRepository } from '../repositories/NamespaceRepository.js';
 import { InspectorSchemaChange } from '../services/SchemaUsageTrafficInspector.js';
 import { SchemaCheckChangeAction } from '../../db/models.js';
+import { traced } from '../tracing.js';
 import {
   composeGraphsInWorker,
   DeserializedComposedGraph,
@@ -136,6 +137,7 @@ export type CheckSubgraph = {
   // will be used only for new subgraphs
   labels?: Label[];
 };
+@traced
 export class Composer {
   constructor(
     private logger: FastifyBaseLogger,

controlplane/src/core/repositories/ApiKeyRepository.ts

@@ -4,10 +4,12 @@ import { PostgresJsDatabase } from 'drizzle-orm/postgres-js';
 import * as schema from '../../db/schema.js';
 import { apiKeyPermissions, apiKeyResources, apiKeys, users } from '../../db/schema.js';
 import { APIKeyDTO } from '../../types/index.js';
+import { traced } from '../tracing.js';
 
 /**
  * Repository for organization related operations.
  */
+@traced
 export class ApiKeyRepository {
   constructor(private db: PostgresJsDatabase<typeof schema>) {}