feat(telemetry): implement MetricsSerializer for metric persistence and locking (#419)

Author: unknownzjc

packages/telemetry/package.json

@@ -18,6 +18,7 @@
 		"lodash": "^4.17.21",
 		"posthog-node": "^5.0.0",
 		"prom-client": "^15.1.3",
+		"proper-lockfile": "^4.1.2",
 		"zod": "^3.25.61"
 	},
 	"devDependencies": {
@@ -26,6 +27,7 @@
 		"@types/async-retry": "^1.4.9",
 		"@types/lodash": "^4.17.20",
 		"@types/node": "20.x",
+		"@types/proper-lockfile": "^4.1.4",
 		"@types/vscode": "^1.84.0",
 		"vitest": "^3.2.3"
 	}

packages/telemetry/src/prometheus/MetricsSerializer.ts

@@ -0,0 +1,71 @@
+import { Registry, Counter, Histogram, Metric, Gauge } from "prom-client"
+import * as fs from "fs/promises"
+import path from "path"
+import { ILogger } from "../../../../src/utils/logger"
+
+interface SerializedMetric {
+	name: string
+	help: string
+	type: string
+	values: {
+		value: number
+		labels: Record<string, string | number>
+		metricName?: string
+	}[]
+	aggregator: string
+}
+
+export class MetricsSerializer {
+	constructor(private logger: ILogger) {}
+
+	public async save(registry: Registry, filePath: string): Promise<void> {
+		try {
+			const metrics = await registry.getMetricsAsJSON()
+			await fs.mkdir(path.dirname(filePath), { recursive: true })
+			await fs.writeFile(filePath, JSON.stringify(metrics, null, 2))
+			this.logger.debug(`[MetricsSerializer] Metrics saved to ${filePath}`)
+		} catch (error) {
+			this.logger.error(`[MetricsSerializer] Failed to save metrics: ${error}`)
+		}
+	}
+
+	public async load(registry: Registry, filePath: string): Promise<void> {
+		try {
+			const data = await fs.readFile(filePath, "utf-8")
+			const metrics: SerializedMetric[] = JSON.parse(data)
+
+			for (const metricData of metrics) {
+				const metric: Metric<string> | undefined = registry.getSingleMetric(metricData.name)
+				if (!metric) {
+					continue
+				}
+
+				if (metric instanceof Counter) {
+					for (const item of metricData.values) {
+						metric.inc(item.labels, item.value)
+					}
+				} else if (metric instanceof Histogram) {
+					const sumValue = metricData.values.find((v) => v.metricName?.endsWith("_sum"))
+
+					if (sumValue && sumValue.value > 0) {
+						metric.observe(sumValue.labels, sumValue.value)
+						this.logger.debug(
+							`[MetricsSerializer] Histogram sum for ${metricData.name} loaded with value: ${sumValue.value}`,
+						)
+					}
+				} else if (metric instanceof Gauge) {
+					for (const item of metricData.values) {
+						metric.set(item.labels, item.value)
+					}
+				}
+			}
+			this.logger.debug(`[MetricsSerializer] Metrics loaded into registry from ${filePath}`)
+		} catch (error) {
+			if (error instanceof Error && "code" in error && error.code === "ENOENT") {
+				this.logger.debug(`[MetricsSerializer] No metrics file at ${filePath}. Starting with a fresh registry.`)
+			} else {
+				this.logger.error(`[MetricsSerializer] Failed to load metrics: ${error}`)
+			}
+		}
+	}
+}

packages/telemetry/src/prometheus/PrometheusTelemetryClient.ts

@@ -2,30 +2,86 @@ import { Pushgateway, Registry } from "prom-client"
 import delay from "delay"
 import crypto from "crypto"
 import retry from "async-retry"
+import path from "path"
+import * as lockfile from "proper-lockfile"
+import * as fs from "fs/promises"
+import * as os from "os"
 
-import { type TelemetryEvent } from "@roo-code/types"
+import { type TelemetryEvent, TelemetryEventName } from "@roo-code/types"
 
 import { BaseTelemetryClient } from "../BaseTelemetryClient"
 import MetricsRecorder from "./metrics"
 import { createLogger, ILogger } from "../../../../src/utils/logger"
 import { getWorkspacePath } from "../../../../src/utils/path"
 import { Package } from "../../../../src/shared/package"
 import { ClineProvider } from "../../../../src/core/webview/ClineProvider"
+import { MetricsSerializer } from "./MetricsSerializer"
 
 export class PrometheusTelemetryClient extends BaseTelemetryClient {
 	private endpoint: string
-	private registry: Registry
-	private metricsRecorder: MetricsRecorder
 	private logger: ILogger
+	private metricsSerializer: MetricsSerializer
+	private persistenceFilePath: string
+
 	constructor(endpoint: string, debug = false) {
-		super(undefined, debug)
+		super(
+			{
+				type: "include",
+				events: [
+					TelemetryEventName.CODE_ACCEPT,
+					TelemetryEventName.CODE_REJECT,
+					TelemetryEventName.CODE_TAB_COMPLETION,
+					TelemetryEventName.ERROR,
+				],
+			},
+			debug,
+		)
 		this.endpoint = endpoint
 		this.logger = createLogger(Package.outputChannel)
-		this.registry = new Registry()
-		this.metricsRecorder = new MetricsRecorder(this.registry)
+		this.metricsSerializer = new MetricsSerializer(this.logger)
+
+		const workspaceHash = this.hashWorkspaceDir()
+		const homeDir = os.homedir()
+		const persistenceDir = path.join(homeDir, ".costrict", "telemetry")
+		this.persistenceFilePath = path.join(persistenceDir, `metrics-${workspaceHash}.json`)
+
 		this.setupPush()
 		this.updateTelemetryState(true)
 	}
+
+	private async operateWithLock<T>(
+		operation: (recorder: MetricsRecorder, registry: Registry) => Promise<T>,
+	): Promise<T | undefined> {
+		try {
+			// Check if file exists first, only create directory and file if needed
+			try {
+				await fs.access(this.persistenceFilePath)
+			} catch (_) {
+				// File doesn't exist, create directory and file
+				await fs.mkdir(path.dirname(this.persistenceFilePath), { recursive: true })
+				await fs.writeFile(this.persistenceFilePath, "[]", "utf-8")
+			}
+
+			await lockfile.lock(this.persistenceFilePath)
+
+			const tempRegistry = new Registry()
+			// Recorder is created ONCE here to register metrics.
+			const tempRecorder = new MetricsRecorder(tempRegistry)
+
+			await this.metricsSerializer.load(tempRegistry, this.persistenceFilePath)
+
+			// We pass the recorder instance to the operation.
+			const result = await operation(tempRecorder, tempRegistry)
+
+			return result
+		} catch (error) {
+			this.logger.error(`[PrometheusTelemetryClient] Lock operation failed: ${error}`)
+			return undefined
+		} finally {
+			await lockfile.unlock(this.persistenceFilePath)
+		}
+	}
+
 	private hashWorkspaceDir() {
 		return crypto.createHash("sha256").update(getWorkspacePath()).digest("hex").toString().slice(0, 8)
 	}
@@ -34,40 +90,42 @@ export class PrometheusTelemetryClient extends BaseTelemetryClient {
 		setInterval(async () => {
 			try {
 				if (this.debug) {
-					this.logger.debug(`[PrometheusTelemetryClient#push] Pushing metrics`)
+					this.logger.debug(`[PrometheusTelemetryClient] Periodic push triggered.`)
 				}
 				await delay(Math.random() * 1000)
 				await this.pushAdd()
 			} catch (error) {
-				this.logger.error(`[PrometheusTelemetryClient#push] ${error}`)
+				this.logger.error(`[PrometheusTelemetryClient#setupPush] ${error}`)
 			}
 		}, times)
 	}
-	// Implement an immediate push method
+
 	public async pushAdd() {
-		const provider = this.providerRef?.deref() as unknown as ClineProvider
-		const { apiConfiguration } = await provider.getState()
-		const { zgsmAccessToken } = apiConfiguration
-		const client = new Pushgateway(
-			this.endpoint,
-			{
-				headers: {
-					Authorization: `Bearer ${zgsmAccessToken}`,
-				},
-			},
-			this.registry,
-		)
-		await retry(
-			async () => {
-				await client.pushAdd({
-					jobName: "costrict",
-					groupings: {
-						instance: this.hashWorkspaceDir(),
-					},
-				})
-			},
-			{ retries: 3 },
-		)
+		return this.operateWithLock(async (recorder, registry) => {
+			// recorder is unused here, which is perfectly fine.
+			if ((await registry.getMetricsAsJSON()).length === 0) {
+				this.logger.debug("[PrometheusTelemetryClient] No metrics to push.")
+				return
+			}
+			const provider = this.providerRef?.deref() as unknown as ClineProvider
+			const { apiConfiguration } = await provider.getState()
+			const { zgsmAccessToken } = apiConfiguration
+			const client = new Pushgateway(
+				this.endpoint,
+				{ headers: { Authorization: `Bearer ${zgsmAccessToken}` } },
+				registry,
+			)
+			await retry(
+				() =>
+					client.pushAdd({
+						jobName: "costrict",
+						groupings: { instance: this.hashWorkspaceDir() },
+					}),
+				{ retries: 3 },
+			)
+
+			this.logger.debug("[PrometheusTelemetryClient] Push successful.")
+		})
 	}
 
 	public override async capture(event: TelemetryEvent): Promise<void> {
@@ -77,21 +135,21 @@ export class PrometheusTelemetryClient extends BaseTelemetryClient {
 			}
 			return
 		}
-		try {
+
+		await this.operateWithLock(async (recorder, registry) => {
+			// Use the recorder instance passed from operateWithLock. NO `new` here.
 			const properties = await this.getEventProperties(event)
-			this.metricsRecorder.record({
-				event: event.event,
-				properties,
-			})
-			if (this.debug) {
-				this.logger.debug(`[PrometheusTelemetryClient#capture] ${event.event}`)
-			}
-		} catch (error) {
+			recorder.record({ event: event.event, properties })
+
+			// Save the modified state back to the file.
+			await this.metricsSerializer.save(registry, this.persistenceFilePath)
+
 			if (this.debug) {
-				this.logger.error(`[PrometheusTelemetryClient#capture] ${error}`)
+				this.logger.debug(`[PrometheusTelemetryClient#capture] Captured and persisted: ${event.event}`)
 			}
-		}
+		})
 	}
+
 	protected override async getEventProperties(event: TelemetryEvent) {
 		let providerProperties: TelemetryEvent["properties"] = {}
 		const { properties } = event
@@ -114,14 +172,12 @@ export class PrometheusTelemetryClient extends BaseTelemetryClient {
 		this.telemetryEnabled = true
 	}
 	public override shutdown(): Promise<void> {
-		return new Promise((resolve, reject) => {
+		return new Promise((resolve) => {
 			this.pushAdd()
-				.then(() => {
-					resolve()
-				})
 				.catch((error) => {
-					reject(error)
+					this.logger.error(`[PrometheusTelemetryClient#shutdown] Final push failed: ${error}`)
 				})
+				.finally(resolve)
 		})
 	}
 }