Add metric to track how many times is each key requested

package.json

@@ -40,7 +40,7 @@
 		"@typescript-eslint/parser": "6.21.0",
 		"commander": "12.1.0",
 		"dotenv": "16.4.0",
-		"esbuild": "0.19.12",
+		"esbuild": "0.21.5",
 		"eslint": "8.56.0",
 		"eslint-config-prettier": "9.1.0",
 		"eslint-config-typescript": "3.0.0",

src/bindings.ts

@@ -4,6 +4,7 @@ export interface Bindings {
 	// variables and secrets
 	DIRECTORY_CACHE_MAX_AGE_SECONDS: string;
 	ENVIRONMENT: string;
+	SERVICE: string;
 	LOGGING_SHIM_TOKEN: string;
 	SENTRY_ACCESS_CLIENT_ID: string;
 	SENTRY_ACCESS_CLIENT_SECRET: string;

src/context/logging.ts

@@ -23,6 +23,7 @@ interface SentryOptions {
 
 	sampleRate?: number;
 	coloName?: string;
+	service?: string;
 }
 
 export class FlexibleLogger implements Logger {
@@ -58,12 +59,14 @@ export class SentryLogger implements Logger {
 	context: Context;
 	request: Request;
 	environment: string;
+	service: string;
 	sampleRate: number;
 
 	constructor(environment: string, options: SentryOptions) {
 		this.environment = environment;
 		this.context = options.context;
 		this.request = options.request;
+		this.service = options.service || '';
 
 		this.sentry = new Toucan({
 			dsn: options.dsn,
@@ -80,6 +83,7 @@ export class SentryLogger implements Logger {
 			},
 		});
 		this.sentry.setTag('coloName', options.coloName);
+		this.sentry.setTag('service', this.service);
 
 		// default sample rate
 		this.sampleRate = 1;

src/context/metrics.ts

@@ -1,17 +1,25 @@
-import { CounterType, RegistryType } from 'promjs';
+import { CounterType, Labels, RegistryType } from 'promjs';
 import { Registry } from 'promjs/registry';
+import { Bindings } from '../bindings';
 export const METRICS_ENDPOINT = 'https://workers-logging.cfdata.org/prometheus';
 
 export interface RegistryOptions {
 	bearerToken?: string;
 }
 
+export interface DefaultLabels {
+	env: string;
+	service: string;
+}
+
 /**
  * A wrapper around the promjs registry to manage registering and publishing metrics
  */
 export class MetricsRegistry {
-	registry: RegistryType;
+	env: Bindings;
+
 	options: RegistryOptions;
+	registry: RegistryType;
 
 	directoryCacheMissTotal: CounterType;
 	erroredRequestsTotal: CounterType;
@@ -22,46 +30,65 @@ export class MetricsRegistry {
 	signedTokenTotal: CounterType;
 	r2RequestsTotal: CounterType;
 
-	constructor(options: RegistryOptions) {
+	constructor(env: Bindings, options: RegistryOptions) {
+		this.env = env;
 		this.options = options;
 		this.registry = new Registry();
 
-		this.directoryCacheMissTotal = this.registry.create(
+		this.directoryCacheMissTotal = this.create(
 			'counter',
 			'directory_cache_miss_total',
 			'Number of requests for private token issuer directory which are not served by the cache.'
 		);
-		this.erroredRequestsTotal = this.registry.create(
+		this.erroredRequestsTotal = this.create(
 			'counter',
 			'errored_requests_total',
 			'Errored requests served to eyeball'
 		);
-		this.issuanceRequestTotal = this.registry.create(
+		this.issuanceRequestTotal = this.create(
 			'counter',
 			'issuance_request_total',
 			'Number of requests for private token issuance.'
 		);
-		this.keyRotationTotal = this.registry.create(
+		this.keyRotationTotal = this.create(
 			'counter',
 			'key_rotation_total',
 			'Number of key rotation performed.'
 		);
-		this.keyClearTotal = this.registry.create(
+		this.keyClearTotal = this.create(
 			'counter',
 			'key_clear_total',
 			'Number of key clear performed.'
 		);
-		this.requestsTotal = this.registry.create('counter', 'requests_total', 'total requests');
-		this.signedTokenTotal = this.registry.create(
+		this.requestsTotal = this.create('counter', 'requests_total', 'total requests');
+		this.signedTokenTotal = this.create(
 			'counter',
 			'signed_token_total',
 			'Number of issued signed private tokens.'
 		);
-		this.r2RequestsTotal = this.registry.create(
-			'counter',
-			'r2_requests_total',
-			'Number of accesses to R2'
-		);
+		this.r2RequestsTotal = this.create('counter', 'r2_requests_total', 'Number of accesses to R2');
+	}
+
+	private create(type: 'counter', name: string, help?: string): CounterType {
+		const counter = this.registry.create(type, name, help);
+		const defaultLabels: DefaultLabels = { env: this.env.ENVIRONMENT, service: this.env.SERVICE };
+		return new Proxy(counter, {
+			get(target, prop, receiver) {
+				if (['collect', 'get', 'inc', 'reset'].includes(prop.toString())) {
+					return function (labels?: Labels) {
+						const mergedLabels = { ...defaultLabels, ...labels };
+						return Reflect.get(target, prop, receiver)?.call(target, mergedLabels);
+					};
+				}
+				if (['add', 'set'].includes(prop.toString())) {
+					return function (value: number, labels?: Labels) {
+						const mergedLabels = { ...defaultLabels, ...labels };
+						return Reflect.get(target, prop, receiver)?.call(target, value, mergedLabels);
+					};
+				}
+				return Reflect.get(target, prop, receiver);
+			},
+		});
 	}
 
 	/**

src/errors.ts

@@ -11,7 +11,7 @@ function shouldSendToSentry(error: Error): boolean {
 export async function handleError(ctx: Context, error: Error) {
 	console.error(error.stack);
 
-	ctx.metrics.erroredRequestsTotal.inc({ env: ctx.env.ENVIRONMENT });
+	ctx.metrics.erroredRequestsTotal.inc();
 
 	const status = (error as HTTPError).status ?? 500;
 	const message = error.message || 'Server Error';

src/index.ts

@@ -30,7 +30,7 @@ interface StorageMetadata extends Record<string, string> {
 }
 
 export const handleTokenRequest = async (ctx: Context, request: Request) => {
-	ctx.metrics.issuanceRequestTotal.inc({ env: ctx.env.ENVIRONMENT });
+	ctx.metrics.issuanceRequestTotal.inc();
 	const contentType = request.headers.get('content-type');
 	if (!contentType || contentType !== MediaType.PRIVATE_TOKEN_REQUEST) {
 		throw new HeaderNotDefinedError(`"Content-Type" must be "${MediaType.PRIVATE_TOKEN_REQUEST}"`);
@@ -43,7 +43,8 @@ export const handleTokenRequest = async (ctx: Context, request: Request) => {
 		throw new Error('Invalid token type');
 	}
 
-	const key = await ctx.cache.ISSUANCE_KEYS.get(tokenRequest.truncatedTokenKeyId.toString());
+	const keyID = tokenRequest.truncatedTokenKeyId.toString();
+	const key = await ctx.cache.ISSUANCE_KEYS.get(keyID);
 
 	if (key === null) {
 		throw new Error('Issuer not initialised');
@@ -79,7 +80,7 @@ export const handleTokenRequest = async (ctx: Context, request: Request) => {
 	const domain = new URL(request.url).host;
 	const issuer = new Issuer(BlindRSAMode.PSS, domain, sk, pk, { supportsRSARAW: true });
 	const signedToken = await issuer.issue(tokenRequest);
-	ctx.metrics.signedTokenTotal.inc({ env: ctx.env.ENVIRONMENT });
+	ctx.metrics.signedTokenTotal.inc({ key_id: keyID });
 
 	return new Response(signedToken.serialize(), {
 		headers: { 'content-type': MediaType.PRIVATE_TOKEN_RESPONSE },
@@ -107,7 +108,7 @@ export const handleTokenDirectory = async (ctx: Context, request: Request) => {
 		}
 		return cachedResponse;
 	}
-	ctx.metrics.directoryCacheMissTotal.inc({ env: ctx.env.ENVIRONMENT });
+	ctx.metrics.directoryCacheMissTotal.inc();
 
 	const keys = await ctx.cache.ISSUANCE_KEYS.list({ include: ['customMetadata'] });
 
@@ -145,7 +146,7 @@ export const handleTokenDirectory = async (ctx: Context, request: Request) => {
 };
 
 export const handleRotateKey = async (ctx: Context, _request?: Request) => {
-	ctx.metrics.keyRotationTotal.inc({ env: ctx.env.ENVIRONMENT });
+	ctx.metrics.keyRotationTotal.inc();
 
 	// Generate a new type 2 Issuer key
 	let publicKeyEnc: string;
@@ -195,7 +196,7 @@ export const handleRotateKey = async (ctx: Context, _request?: Request) => {
 };
 
 const handleClearKey = async (ctx: Context, _request?: Request) => {
-	ctx.metrics.keyClearTotal.inc({ env: ctx.env.ENVIRONMENT });
+	ctx.metrics.keyClearTotal.inc();
 	const keys = await ctx.env.ISSUANCE_KEYS.list();
 
 	let latestKey: R2Object = keys.objects[0];
@@ -242,7 +243,7 @@ export default {
 			env,
 			ectx.waitUntil.bind(ectx),
 			new ConsoleLogger(),
-			new MetricsRegistry({ bearerToken: env.LOGGING_SHIM_TOKEN })
+			new MetricsRegistry(env, { bearerToken: env.LOGGING_SHIM_TOKEN })
 		);
 		const date = new Date(event.scheduledTime);
 		const isRotation = date.getUTCDate() === 1;

src/router.ts

@@ -82,7 +82,7 @@ export class Router {
 
 	private buildContext(request: Request, env: Bindings, ectx: ExecutionContext): Context {
 		// Prometheus Registry should be unique per request
-		const metrics = new MetricsRegistry({ bearerToken: env.LOGGING_SHIM_TOKEN });
+		const metrics = new MetricsRegistry(env, { bearerToken: env.LOGGING_SHIM_TOKEN });
 
 		// Use a flexible reporter, so that it uses console.log when debugging, and Core Sentry when in production
 		let logger: Logger;
@@ -100,6 +100,7 @@ export class Router {
 				accessClientId: env.SENTRY_ACCESS_CLIENT_ID,
 				accessClientSecret: env.SENTRY_ACCESS_CLIENT_SECRET,
 				release: RELEASE,
+				service: env.SERVICE,
 				sampleRate: sentrySampleRate,
 				coloName: request?.cf?.colo as string,
 			});
@@ -120,7 +121,7 @@ export class Router {
 		ectx: ExecutionContext
 	): Promise<Response> {
 		const ctx = this.buildContext(request, env, ectx);
-		ctx.metrics.requestsTotal.inc({ env: ctx.env.ENVIRONMENT });
+		ctx.metrics.requestsTotal.inc();
 		const rawPath = new URL(request.url).pathname;
 		const path = this.normalisePath(rawPath);
 

test/localhost-e2e.test.ts

@@ -1,41 +1,55 @@
 import { spawn } from 'node:child_process';
+import fetch from 'node-fetch';
 import { testE2E } from './e2e/issuer';
 
 const ISSUER_URL = 'localhost:8787';
 
 describe('e2e on localhost', () => {
-	let serverPID: number | undefined;
+	let serverProcess: ReturnType<typeof spawn> | undefined;
 
 	beforeAll(async () => {
 		try {
-			// start server as an independant process with npm run dev
-			const serverProcess = spawn('npm', ['run', 'dev'], { stdio: 'inherit', detached: true });
-			serverPID = serverProcess.pid;
-			console.log('Creating server with PID: ', serverPID);
+			// Start server as an independent process with npm run dev
+			serverProcess = spawn('npm', ['run', 'dev'], { stdio: 'inherit', detached: true });
+			console.log('Creating server with PID:', serverProcess.pid);
 
-			// check the server is online
+			// Check the server is online
 			const backoffInMs = 100;
-			while (true) {
+			const maxRetries = 200; // 20 seconds total with 100ms backoff
+			let retries = 0;
+			while (retries < maxRetries) {
 				try {
 					await fetch(`http://${ISSUER_URL}/`);
 					console.log('Server is up');
 					break;
 				} catch (e) {
+					retries++;
 					await new Promise(resolve => setTimeout(resolve, backoffInMs));
 				}
 			}
+			if (retries === maxRetries) {
+				throw new Error('Server did not start within the expected time');
+			}
 		} catch (err) {
-			console.log('Server failure: ', err);
+			console.log('Server failure:', err);
 		}
-	}, 10 * 1000);
+	}, 30 * 1000); // Increase timeout to 30 seconds
 
 	afterAll(() => {
-		process.kill(-serverPID!);
-		console.log('Server is down');
+		if (serverProcess && serverProcess.pid) {
+			try {
+				process.kill(-serverProcess.pid);
+				console.log('Server is down');
+			} catch (error) {
+				console.error(`Failed to kill server process: ${error}`);
+			}
+		} else {
+			console.log('Server process was not started or already terminated');
+		}
 	});
 
 	it('should issue a token that is valid', async () => {
-		// provision new keys
+		// Provision new keys
 		const response = await fetch(`http://${ISSUER_URL}/admin/rotate`, {
 			method: 'POST',
 		});

test/mocks.ts

@@ -55,7 +55,7 @@ export interface MockContextOptions {
 
 export const getContext = (options: MockContextOptions): Context => {
 	const logger = options.logger ?? new ConsoleLogger();
-	const metrics = options.metrics ?? new MetricsRegistry({});
+	const metrics = options.metrics ?? new MetricsRegistry(options.env, {});
 	const waitUntilFunc = options.waitUntilFunc || options.ectx.waitUntil.bind(options.ectx);
 	return new Context(options.env, waitUntilFunc, logger, metrics);
 };