feat: prometheus instrumentation

Author: n1ru4l

packages/services/api/package.json

@@ -13,6 +13,7 @@
   "devDependencies": {
     "@aws-sdk/client-s3": "3.723.0",
     "@aws-sdk/s3-request-presigner": "3.723.0",
+    "@bentocache/plugin-prometheus": "0.2.0",
     "@date-fns/utc": "2.1.0",
     "@graphql-hive/core": "workspace:*",
     "@graphql-inspector/core": "5.1.0-alpha-20231208113249-34700c8a",

packages/services/api/src/create.ts

@@ -54,6 +54,7 @@ import { IdTranslator } from './modules/shared/providers/id-translator';
 import { Logger } from './modules/shared/providers/logger';
 import { Mutex } from './modules/shared/providers/mutex';
 import { PG_POOL_CONFIG } from './modules/shared/providers/pg-pool';
+import { PrometheusConfig } from './modules/shared/providers/prometheus-config';
 import { HivePubSub, PUB_SUB_CONFIG } from './modules/shared/providers/pub-sub';
 import { REDIS_INSTANCE } from './modules/shared/providers/redis';
 import { S3_CONFIG, type S3Config } from './modules/shared/providers/s3-config';
@@ -119,6 +120,7 @@ export function createRegistry({
   organizationOIDC,
   pubSub,
   appDeploymentsEnabled,
+  prometheus,
 }: {
   logger: Logger;
   storage: Storage;
@@ -164,6 +166,7 @@ export function createRegistry({
   organizationOIDC: boolean;
   pubSub: HivePubSub;
   appDeploymentsEnabled: boolean;
+  prometheus: null | Record<string, unknown>;
 }) {
   const s3Config: S3Config = [
     {
@@ -323,6 +326,12 @@ export function createRegistry({
       scope: Scope.Operation,
       deps: [CONTEXT],
     },
+    {
+      provide: PrometheusConfig,
+      useFactory() {
+        return new PrometheusConfig(!!prometheus);
+      },
+    },
   ];
 
   if (emailsEndpoint) {

packages/services/api/src/modules/auth/index.ts

@@ -2,6 +2,7 @@ import { createModule } from 'graphql-modules';
 import { AuditLogManager } from '../audit-logs/providers/audit-logs-manager';
 import { AuthManager } from './providers/auth-manager';
 import { OrganizationAccess } from './providers/organization-access';
+import { OrganizationAccessTokenValidationCache } from './providers/organization-access-token-validation-cache';
 import { ProjectAccess } from './providers/project-access';
 import { TargetAccess } from './providers/target-access';
 import { UserManager } from './providers/user-manager';
@@ -20,5 +21,6 @@ export const authModule = createModule({
     ProjectAccess,
     TargetAccess,
     AuditLogManager,
+    OrganizationAccessTokenValidationCache,
   ],
 });

packages/services/api/src/modules/auth/lib/organization-access-token-strategy.ts

@@ -1,23 +1,11 @@
 import * as crypto from 'node:crypto';
-import { BentoCache, bentostore } from 'bentocache';
-import { memoryDriver } from 'bentocache/build/src/drivers/memory';
 import { type FastifyReply, type FastifyRequest } from '@hive/service-common';
 import * as OrganizationAccessKey from '../../organization/lib/organization-access-key';
 import { OrganizationAccessTokensCache } from '../../organization/providers/organization-access-tokens-cache';
 import { Logger } from '../../shared/providers/logger';
+import { OrganizationAccessTokenValidationCache } from '../providers/organization-access-token-validation-cache';
 import { AuthNStrategy, AuthorizationPolicyStatement, Session } from './authz';
 
-const cache = new BentoCache({
-  default: 'organizationAccessTokenValidation',
-  stores: {
-    organizationAccessTokenValidation: bentostore().useL1Layer(
-      memoryDriver({
-        maxItems: 10_000,
-      }),
-    ),
-  },
-});
-
 function hashToken(token: string) {
   return crypto.createHash('sha256').update(token).digest('hex');
 }
@@ -50,12 +38,18 @@ export class OrganizationAccessTokenSession extends Session {
 export class OrganizationAccessTokenStrategy extends AuthNStrategy<OrganizationAccessTokenSession> {
   private logger: Logger;
 
-  private cache: OrganizationAccessTokensCache;
+  private organizationAccessTokenCache: OrganizationAccessTokensCache;
+  private organizationAccessTokenValidationCache: OrganizationAccessTokenValidationCache;
 
-  constructor(deps: { logger: Logger; cache: OrganizationAccessTokensCache }) {
+  constructor(deps: {
+    logger: Logger;
+    organizationAccessTokensCache: OrganizationAccessTokensCache;
+    organizationAccessTokenValidationCache: OrganizationAccessTokenValidationCache;
+  }) {
     super();
     this.logger = deps.logger.child({ module: 'OrganizationAccessTokenStrategy' });
-    this.cache = deps.cache;
+    this.organizationAccessTokenCache = deps.organizationAccessTokensCache;
+    this.organizationAccessTokenValidationCache = deps.organizationAccessTokenValidationCache;
   }
 
   async parse(args: {
@@ -89,14 +83,16 @@ export class OrganizationAccessTokenStrategy extends AuthNStrategy<OrganizationA
       return null;
     }
 
-    const organizationAccessToken = await this.cache.get(result.accessKey.id);
+    const organizationAccessToken = await this.organizationAccessTokenCache.get(
+      result.accessKey.id,
+    );
     if (!organizationAccessToken) {
       return null;
     }
 
     // let's hash it so we do not store the plain private key in memory
     const key = hashToken(accessToken);
-    const isHashMatch = await cache.getOrSetForever({
+    const isHashMatch = await this.organizationAccessTokenValidationCache.getOrSetForever({
       factory: () =>
         OrganizationAccessKey.verify(result.accessKey.privateKey, organizationAccessToken.hash),
       key,

packages/services/api/src/modules/auth/providers/organization-access-token-validation-cache.ts

@@ -0,0 +1,40 @@
+import { BentoCache, bentostore } from 'bentocache';
+import { memoryDriver } from 'bentocache/build/src/drivers/memory';
+import { Injectable, Scope } from 'graphql-modules';
+import { prometheusPlugin } from '@bentocache/plugin-prometheus';
+import { PrometheusConfig } from '../../shared/providers/prometheus-config';
+
+/**
+ * Cache for performant OrganizationAccessToken lookups.
+ */
+@Injectable({
+  scope: Scope.Singleton,
+  global: true,
+})
+export class OrganizationAccessTokenValidationCache {
+  private cache: BentoCache<{ store: ReturnType<typeof bentostore> }>;
+
+  constructor(prometheusConfig: PrometheusConfig) {
+    this.cache = new BentoCache({
+      default: 'organizationAccessTokenValidation',
+      plugins: prometheusConfig.isEnabled
+        ? [
+            prometheusPlugin({
+              prefix: 'bentocache_organization_access_token_validation',
+            }),
+          ]
+        : undefined,
+      stores: {
+        organizationAccessTokenValidation: bentostore().useL1Layer(
+          memoryDriver({
+            maxItems: 10_000,
+            prefix: 'bentocache:organization-access-token-validation',
+          }),
+        ),
+      },
+    });
+  }
+
+  getOrSetForever: typeof this.cache.getOrSetForever = (...args) =>
+    this.cache.getOrSetForever(...args);
+}

packages/services/api/src/modules/organization/providers/organization-access-tokens-cache.ts

@@ -4,8 +4,10 @@ import { redisDriver } from 'bentocache/build/src/drivers/redis';
 import { Inject, Injectable, Scope } from 'graphql-modules';
 import Redis from 'ioredis';
 import type { DatabasePool } from 'slonik';
+import { prometheusPlugin } from '@bentocache/plugin-prometheus';
 import { Logger } from '../../shared/providers/logger';
 import { PG_POOL_CONFIG } from '../../shared/providers/pg-pool';
+import { PrometheusConfig } from '../../shared/providers/prometheus-config';
 import { REDIS_INSTANCE } from '../../shared/providers/redis';
 import { findById, type OrganizationAccessToken } from './organization-access-tokens';
 
@@ -24,10 +26,18 @@ export class OrganizationAccessTokensCache {
     @Inject(REDIS_INSTANCE) redis: Redis,
     @Inject(PG_POOL_CONFIG) pool: DatabasePool,
     logger: Logger,
+    prometheusConfig: PrometheusConfig,
   ) {
     this.findById = findById({ pool, logger });
     this.cache = new BentoCache({
       default: 'organizationAccessTokens',
+      plugins: prometheusConfig.isEnabled
+        ? [
+            prometheusPlugin({
+              prefix: 'bentocache_organization_access_tokens',
+            }),
+          ]
+        : undefined,
       stores: {
         organizationAccessTokens: bentostore()
           .useL1Layer(

packages/services/api/src/modules/shared/providers/prometheus-config.ts

@@ -0,0 +1,12 @@
+import { Injectable, Scope } from 'graphql-modules';
+
+@Injectable({
+  scope: Scope.Singleton,
+})
+export class PrometheusConfig {
+  constructor(private _isEnabled = false) {}
+
+  get isEnabled() {
+    return this._isEnabled;
+  }
+}

packages/services/server/src/index.ts

@@ -58,6 +58,7 @@ import { AuthN } from '../../api/src/modules/auth/lib/authz';
 import { OrganizationAccessTokenStrategy } from '../../api/src/modules/auth/lib/organization-access-token-strategy';
 import { SuperTokensUserAuthNStrategy } from '../../api/src/modules/auth/lib/supertokens-strategy';
 import { TargetAccessTokenStrategy } from '../../api/src/modules/auth/lib/target-access-token-strategy';
+import { OrganizationAccessTokenValidationCache } from '../../api/src/modules/auth/providers/organization-access-token-validation-cache';
 import { OrganizationAccessTokensCache } from '../../api/src/modules/organization/providers/organization-access-tokens-cache';
 import { createContext, internalApiRouter } from './api';
 import { asyncStorage } from './async-storage';
@@ -404,6 +405,7 @@ export async function main() {
       supportConfig: env.zendeskSupport,
       pubSub,
       appDeploymentsEnabled: env.featureFlags.appDeploymentsEnabled,
+      prometheus: env.prometheus,
     });
 
     const authN = new AuthN({
@@ -421,7 +423,10 @@ export async function main() {
         (logger: Logger) =>
           new OrganizationAccessTokenStrategy({
             logger,
-            cache: registry.injector.get(OrganizationAccessTokensCache),
+            organizationAccessTokensCache: registry.injector.get(OrganizationAccessTokensCache),
+            organizationAccessTokenValidationCache: registry.injector.get(
+              OrganizationAccessTokenValidationCache,
+            ),
           }),
         (logger: Logger) =>
           new TargetAccessTokenStrategy({