Refactor OTEL instrumentation to ensure it works correctly

Instrument Prisma.

README.md

@@ -116,7 +116,7 @@ below to generate a project of your choice.
 
 ```bash
 # Generate an standard JavaScript library.
-npx nx g @nrwl/workspace:lib
+npx nx g @nrwl/js:lib 
 # Generate an NestJS library.
 npx nx g @nx/nest:library
 # Generate an NestJS application.
@@ -147,25 +147,30 @@ the `nx.json`.
 The Armory Stack uses [OpenTelemetry](https://opentelemetry.io/docs/) (OTEL)
 for observability through traces and metrics.
 
+For details on how OTEL is instrumented, head over
+[@narval/open-telemetry](./packages/open-telemetry/).
+
 ### Setup
 
-1. Set environment variables:
+1. Start local dependencies:
 
 ```bash
-OTEL_METRIC_EXPORTER_URL=http://localhost:4318/v1/metrics
-OTEL_TRACE_EXPORTER_URL=http://localhost:4318/v1/traces
+make docker/otel/up
 ```
 
-> [!IMPORTANT]
-> OTEL module automatically disables if either URL is not set.
-
-2. Start local dependencies:
+2. Set environment variables:
 
 ```bash
-make docker/otel/up
+OTEL_SDK_DISABLED=false
+OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318
 ```
 
-3. Access Jaeger UI: http://localhost:16686
+> [!NOTE]
+> OTEL is disabled by default in development. 
+
+3. Restart the application.
+
+4. Access Jaeger UI: http://localhost:16686
 
 ### Naming Conventions
 

apps/armory/.env.default

@@ -34,5 +34,6 @@ POLICY_ENGINE_ADMIN_API_KEYS=engine-admin-api-key
 MANAGED_DATASTORE_BASE_URL=http://localhost:3005/v1/data
 
 # OpenTelemetry configuration.
-OTEL_SDK_DISABLED=false
+OTEL_SDK_DISABLED=true
 OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318
+OTEL_LOGS_EXPORTER=otlp

apps/armory/src/main.ts

@@ -1,12 +1,14 @@
+import { instrumentOpenTelemetry } from '@narval/open-telemetry'
+import { DiagLogLevel } from '@opentelemetry/api'
+
+// IMPORTANT: OpenTelemetry SDK must be registered before any other imports to
+// ensure proper instrumentation. The instrumentation packages patches Node.js
+// runtime - if NestFactory or other dependencies load first, they'll use the
+// unpatched runtime and won't be instrumented correctly.
+instrumentOpenTelemetry({ serviceName: 'auth', diagLogLevel: DiagLogLevel.DEBUG })
+
 import { ConfigService } from '@narval/config-module'
-import {
-  LoggerService,
-  buildOpenTelemetrySdk,
-  withApiVersion,
-  withCors,
-  withLogger,
-  withSwagger
-} from '@narval/nestjs-shared'
+import { LoggerService, withApiVersion, withCors, withLogger, withSwagger } from '@narval/nestjs-shared'
 import { ClassSerializerInterceptor, INestApplication, ValidationPipe } from '@nestjs/common'
 import { NestFactory, Reflector } from '@nestjs/core'
 import { json } from 'express'
@@ -68,11 +70,6 @@ const withGlobalFilters =
  * successfully bootstrapped.
  */
 async function bootstrap(): Promise<void> {
-  // IMPORTANT: For the @opentelemetry/auto-instrumentations-node package work
-  // correctly, OpenTelemetry MUST register before the NestFactory.create.
-  const openTelemetrySdk = buildOpenTelemetrySdk({ serviceName: 'auth' })
-  openTelemetrySdk.start()
-
   const application = await NestFactory.create(ArmoryModule, { bufferLogs: true })
   const configService = application.get<ConfigService<Config>>(ConfigService)
   const logger = application.get<LoggerService>(LoggerService)

apps/armory/src/managed-data-store/core/service/entity-data-store.service.ts

@@ -1,5 +1,7 @@
+import { MetricService, TraceService } from '@narval/nestjs-shared'
 import { EntityStore } from '@narval/policy-engine-shared'
-import { HttpStatus, Injectable, NotFoundException } from '@nestjs/common'
+import { HttpStatus, Inject, Injectable, NotFoundException } from '@nestjs/common'
+import { Counter } from '@opentelemetry/api'
 import { ClientService } from '../../../client/core/service/client.service'
 import { ClusterService } from '../../../policy-engine/core/service/cluster.service'
 import { SetEntityStoreResponse } from '../../http/rest/dto/set-entity-store.dto'
@@ -8,21 +10,45 @@ import { SignatureService } from './signature.service'
 
 @Injectable()
 export class EntityDataStoreService extends SignatureService {
+  private getCounter: Counter
+  private setCounter: Counter
+
   constructor(
     private entityDataStoreRepository: EntityDataStoreRepository,
     private clientService: ClientService,
-    private clusterService: ClusterService
+    private clusterService: ClusterService,
+    @Inject(TraceService) private traceService: TraceService,
+    @Inject(MetricService) private metricService: MetricService
   ) {
     super()
+
+    this.getCounter = this.metricService.createCounter('entity_data_store_get_count')
+    this.setCounter = this.metricService.createCounter('entity_data_store_set_count')
   }
 
   async getEntities(clientId: string): Promise<EntityStore | null> {
+    this.getCounter.add(1, { clientId })
+
+    const span = this.traceService.startSpan(`${EntityDataStoreService.name}.getEntities`, {
+      attributes: { clientId }
+    })
+
     const entityStore = await this.entityDataStoreRepository.getLatestDataStore(clientId)
 
-    return entityStore ? EntityStore.parse(entityStore.data) : null
+    const response = entityStore ? EntityStore.parse(entityStore.data) : null
+
+    span.end()
+
+    return response
   }
 
   async setEntities(clientId: string, payload: EntityStore) {
+    this.setCounter.add(1, { clientId })
+
+    const span = this.traceService.startSpan(`${EntityDataStoreService.name}.setEntities`, {
+      attributes: { clientId }
+    })
+
     const client = await this.clientService.findById(clientId)
 
     if (!client) {
@@ -47,10 +73,14 @@ export class EntityDataStoreService extends SignatureService {
 
     const success = await this.clusterService.sync(clientId)
 
-    return SetEntityStoreResponse.parse({
+    const response = SetEntityStoreResponse.parse({
       latestSync: { success },
       entity: EntityStore.parse(data),
       version
     })
+
+    span.end()
+
+    return response
   }
 }

apps/armory/src/managed-data-store/core/service/policy-data-store.service.ts

@@ -1,27 +1,53 @@
+import { MetricService, TraceService } from '@narval/nestjs-shared'
 import { PolicyStore } from '@narval/policy-engine-shared'
-import { HttpStatus, Injectable, NotFoundException } from '@nestjs/common'
+import { HttpStatus, Inject, Injectable, NotFoundException } from '@nestjs/common'
+import { Counter } from '@opentelemetry/api'
 import { ClientService } from '../../../client/core/service/client.service'
 import { ClusterService } from '../../../policy-engine/core/service/cluster.service'
 import { PolicyDataStoreRepository } from '../../persistence/repository/policy-data-store.repository'
 import { SignatureService } from './signature.service'
 
 @Injectable()
 export class PolicyDataStoreService extends SignatureService {
+  private getCounter: Counter
+  private setCounter: Counter
+
   constructor(
     private policyDataStoreRepository: PolicyDataStoreRepository,
     private clientService: ClientService,
-    private clusterService: ClusterService
+    private clusterService: ClusterService,
+    @Inject(TraceService) private traceService: TraceService,
+    @Inject(MetricService) private metricService: MetricService
   ) {
     super()
+
+    this.getCounter = this.metricService.createCounter('policy_data_store_get_count')
+    this.setCounter = this.metricService.createCounter('policy_data_store_set_count')
   }
 
   async getPolicies(clientId: string): Promise<PolicyStore | null> {
+    this.getCounter.add(1, { clientId })
+
+    const span = this.traceService.startSpan(`${PolicyDataStoreService.name}.getPolicies`, {
+      attributes: { clientId }
+    })
+
     const policyStore = await this.policyDataStoreRepository.getLatestDataStore(clientId)
 
-    return policyStore ? PolicyStore.parse(policyStore.data) : null
+    const response = policyStore ? PolicyStore.parse(policyStore.data) : null
+
+    span.end()
+
+    return response
   }
 
   async setPolicies(clientId: string, payload: PolicyStore) {
+    this.setCounter.add(1, { clientId })
+
+    const span = this.traceService.startSpan(`${PolicyDataStoreService.name}.setPolicies`, {
+      attributes: { clientId }
+    })
+
     const client = await this.clientService.findById(clientId)
 
     if (!client) {
@@ -46,10 +72,14 @@ export class PolicyDataStoreService extends SignatureService {
 
     const success = await this.clusterService.sync(clientId)
 
-    return {
+    const response = {
       latestSync: { success },
       policy: PolicyStore.parse(data),
       version
     }
+
+    span.end()
+
+    return response
   }
 }