feat: track lmdb stats

Author: alexghr

yarn-project/archiver/src/archiver/archiver.ts

@@ -176,7 +176,7 @@ export class Archiver implements ArchiveSource {
       config.l1Contracts.registryAddress,
       archiverStore,
       config.archiverPollingIntervalMS ?? 10_000,
-      new ArchiverInstrumentation(telemetry),
+      new ArchiverInstrumentation(telemetry, () => archiverStore.estimateSize()),
       { l1StartBlock, l1GenesisTime, epochDuration, slotDuration, ethereumSlotDuration },
     );
     await archiver.start(blockUntilSynced);
@@ -271,9 +271,6 @@ export class Archiver implements ArchiveSource {
       // the chain locally before we start unwinding stuff. This can be optimized by figuring out
       // up to which point we're pruning, and then requesting L2 blocks up to that point only.
       await this.handleEpochPrune(provenBlockNumber, currentL1BlockNumber);
-
-      const storeSizes = this.store.estimateSize();
-      this.instrumentation.recordDBMetrics(storeSizes);
     }
   }
 

yarn-project/archiver/src/archiver/instrumentation.ts

@@ -5,6 +5,7 @@ import {
   type Gauge,
   type Histogram,
   LmdbMetrics,
+  LmdbStatsCallback,
   Metrics,
   type TelemetryClient,
   type UpDownCounter,
@@ -23,7 +24,7 @@ export class ArchiverInstrumentation {
 
   private log = createDebugLogger('aztec:archiver:instrumentation');
 
-  constructor(private telemetry: TelemetryClient) {
+  constructor(private telemetry: TelemetryClient, lmdbStats?: LmdbStatsCallback) {
     const meter = telemetry.getMeter('Archiver');
     this.blockHeight = meter.createGauge(Metrics.ARCHIVER_BLOCK_HEIGHT, {
       description: 'The height of the latest block processed by the archiver',
@@ -72,13 +73,10 @@ export class ArchiverInstrumentation {
         name: Metrics.ARCHIVER_DB_NUM_ITEMS,
         description: 'Num items in the archiver database',
       },
+      lmdbStats,
     );
   }
 
-  public recordDBMetrics(metrics: { mappingSize: number; numItems: number; actualSize: number }) {
-    this.dbMetrics.recordDBMetrics(metrics);
-  }
-
   public isEnabled(): boolean {
     return this.telemetry.isEnabled();
   }

yarn-project/p2p/src/mem_pools/instrumentation.ts

@@ -3,6 +3,7 @@ import {
   Attributes,
   type Histogram,
   LmdbMetrics,
+  LmdbStatsCallback,
   Metrics,
   type TelemetryClient,
   type UpDownCounter,
@@ -58,7 +59,7 @@ export class PoolInstrumentation<PoolObject extends Gossipable> {
 
   private defaultAttributes;
 
-  constructor(telemetry: TelemetryClient, name: PoolName) {
+  constructor(telemetry: TelemetryClient, name: PoolName, dbStats?: LmdbStatsCallback) {
     const meter = telemetry.getMeter(name);
     this.defaultAttributes = { [Attributes.POOL_NAME]: name };
 
@@ -98,13 +99,10 @@ export class PoolInstrumentation<PoolObject extends Gossipable> {
         name: Metrics.MEMPOOL_DB_NUM_ITEMS,
         description: 'Num items in database for the Tx mempool',
       },
+      dbStats,
     );
   }
 
-  public recordDBMetrics(metrics: { mappingSize: number; numItems: number; actualSize: number }) {
-    this.dbMetrics.recordDBMetrics(metrics);
-  }
-
   public recordSize(poolObject: PoolObject) {
     this.objectSize.record(poolObject.getSize());
   }

yarn-project/p2p/src/mem_pools/tx_pool/aztec_kv_tx_pool.ts

@@ -37,7 +37,7 @@ export class AztecKVTxPool implements TxPool {
 
     this.#store = store;
     this.#log = log;
-    this.#metrics = new PoolInstrumentation(telemetry, PoolName.TX_POOL);
+    this.#metrics = new PoolInstrumentation(telemetry, PoolName.TX_POOL, () => store.estimateSize());
   }
 
   public markAsMined(txHashes: TxHash[], blockNumber: number): Promise<void> {
@@ -53,8 +53,6 @@ export class AztecKVTxPool implements TxPool {
       }
       this.#metrics.recordRemovedObjects(deleted, 'pending');
       this.#metrics.recordAddedObjects(txHashes.length, 'mined');
-      const storeSizes = this.#store.estimateSize();
-      this.#metrics.recordDBMetrics(storeSizes);
     });
   }
 

yarn-project/prover-client/src/proving_broker/factory.ts

@@ -11,7 +11,7 @@ export async function createAndStartProvingBroker(
   client: TelemetryClient,
 ): Promise<ProvingBroker> {
   const database = config.proverBrokerDataDirectory
-    ? new KVBrokerDatabase(AztecLmdbStore.open(config.proverBrokerDataDirectory))
+    ? new KVBrokerDatabase(AztecLmdbStore.open(config.proverBrokerDataDirectory), client)
     : new InMemoryBrokerDatabase();
 
   const broker = new ProvingBroker(database, client, {

yarn-project/prover-client/src/proving_broker/proving_broker.test.ts

@@ -18,7 +18,7 @@ describe.each([
   () => ({ database: new InMemoryBrokerDatabase(), cleanup: undefined }),
   () => {
     const store = openTmpStore(true);
-    const database = new KVBrokerDatabase(store);
+    const database = new KVBrokerDatabase(store, new NoopTelemetryClient());
     const cleanup = () => store.close();
     return { database, cleanup };
   },

yarn-project/prover-client/src/proving_broker/proving_broker_database/persisted.ts

@@ -1,14 +1,29 @@
 import { type ProofUri, ProvingJob, type ProvingJobId, ProvingJobSettledResult } from '@aztec/circuit-types';
 import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';
 import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
+import { LmdbMetrics, Metrics, TelemetryClient } from '@aztec/telemetry-client';
 
 import { type ProvingBrokerDatabase } from '../proving_broker_database.js';
 
 export class KVBrokerDatabase implements ProvingBrokerDatabase {
   private jobs: AztecMap<ProvingJobId, string>;
   private jobResults: AztecMap<ProvingJobId, string>;
-
-  constructor(private store: AztecKVStore) {
+  private metrics: LmdbMetrics;
+
+  constructor(private store: AztecKVStore, client: TelemetryClient) {
+    this.metrics = new LmdbMetrics(
+      client.getMeter('KVBrokerDatabase'),
+      {
+        name: Metrics.PROVING_QUEUE_DB_MAP_SIZE,
+        description: 'Database map size for the proving broker',
+      },
+      {
+        name: Metrics.PROVING_QUEUE_DB_USED_SIZE,
+        description: 'Database used size for the proving broker',
+      },
+      { name: Metrics.PROVING_QUEUE_DB_NUM_ITEMS, description: 'Number of items in the broker database' },
+      () => store.estimateSize(),
+    );
     this.jobs = store.openMap('proving_jobs');
     this.jobResults = store.openMap('proving_job_results');
   }

yarn-project/telemetry-client/src/lmdb_metrics.ts

@@ -1,38 +1,47 @@
-import { type Gauge, type Meter, type Metrics, ValueType } from './telemetry.js';
+import { BatchObservableResult, type Meter, type Metrics, ObservableGauge, ValueType } from './telemetry.js';
 
 export type LmdbMetricDescriptor = {
   name: Metrics;
   description: string;
 };
 
+export type LmdbStatsCallback = () => { mappingSize: number; numItems: number; actualSize: number };
+
 export class LmdbMetrics {
-  private dbMapSize: Gauge;
-  private dbUsedSize: Gauge;
-  private dbNumItems: Gauge;
+  private dbMapSize: ObservableGauge;
+  private dbUsedSize: ObservableGauge;
+  private dbNumItems: ObservableGauge;
 
   constructor(
     meter: Meter,
     dbMapSizeDescriptor: LmdbMetricDescriptor,
     dbUsedSizeDescriptor: LmdbMetricDescriptor,
     dbNumItemsDescriptor: LmdbMetricDescriptor,
+    private getStats?: LmdbStatsCallback,
   ) {
-    this.dbMapSize = meter.createGauge(dbMapSizeDescriptor.name, {
+    this.dbMapSize = meter.createObservableGauge(dbMapSizeDescriptor.name, {
       description: dbMapSizeDescriptor.description,
       valueType: ValueType.INT,
     });
-    this.dbUsedSize = meter.createGauge(dbUsedSizeDescriptor.name, {
+    this.dbUsedSize = meter.createObservableGauge(dbUsedSizeDescriptor.name, {
       description: dbUsedSizeDescriptor.description,
       valueType: ValueType.INT,
     });
-    this.dbNumItems = meter.createGauge(dbNumItemsDescriptor.name, {
+    this.dbNumItems = meter.createObservableGauge(dbNumItemsDescriptor.name, {
       description: dbNumItemsDescriptor.description,
       valueType: ValueType.INT,
     });
-  }
 
-  public recordDBMetrics(metrics: { mappingSize: number; numItems: number; actualSize: number }) {
-    this.dbMapSize.record(metrics.mappingSize);
-    this.dbNumItems.record(metrics.actualSize);
-    this.dbUsedSize.record(metrics.actualSize);
+    meter.addBatchObservableCallback(this.recordDBMetrics, [this.dbMapSize, this.dbUsedSize, this.dbNumItems]);
   }
+
+  private recordDBMetrics = (observable: BatchObservableResult) => {
+    if (!this.getStats) {
+      return;
+    }
+    const metrics = this.getStats();
+    observable.observe(this.dbMapSize, metrics.mappingSize);
+    observable.observe(this.dbNumItems, metrics.numItems);
+    observable.observe(this.dbUsedSize, metrics.actualSize);
+  };
 }

yarn-project/telemetry-client/src/metrics.ts

@@ -86,6 +86,9 @@ export const PROVING_QUEUE_RETRIED_JOBS = 'aztec.proving_queue.retried_jobs';
 export const PROVING_QUEUE_TIMED_OUT_JOBS = 'aztec.proving_queue.timed_out_jobs';
 export const PROVING_QUEUE_JOB_WAIT = 'aztec.proving_queue.job_wait';
 export const PROVING_QUEUE_JOB_DURATION = 'aztec.proving_queue.job_duration';
+export const PROVING_QUEUE_DB_NUM_ITEMS = 'aztec.proving_queue.db.num_items';
+export const PROVING_QUEUE_DB_MAP_SIZE = 'aztec.proving_queue.db.map_size';
+export const PROVING_QUEUE_DB_USED_SIZE = 'aztec.proving_queue.db.used_size';
 
 export const PROVING_AGENT_IDLE = 'aztec.proving_queue.agent.idle';
 

yarn-project/telemetry-client/src/telemetry.ts

@@ -1,6 +1,9 @@
 import {
   type AttributeValue,
+  BatchObservableCallback,
   type MetricOptions,
+  Observable,
+  type BatchObservableResult as OtelBatchObservableResult,
   type Gauge as OtelGauge,
   type Histogram as OtelHistogram,
   type ObservableGauge as OtelObservableGauge,
@@ -33,6 +36,7 @@ export type UpDownCounter = OtelUpDownCounter<Attributes>;
 export type ObservableGauge = OtelObservableGauge<Attributes>;
 export type ObservableUpDownCounter = OtelObservableUpDownCounter<Attributes>;
 export type ObservableResult = OtelObservableResult<Attributes>;
+export type BatchObservableResult = OtelBatchObservableResult<Attributes>;
 
 export { Tracer };
 
@@ -55,6 +59,16 @@ export interface Meter {
    */
   createObservableGauge(name: Metrics, options?: MetricOptions): ObservableGauge;
 
+  addBatchObservableCallback(
+    callback: BatchObservableCallback<Attributes>,
+    observables: Observable<Attributes>[],
+  ): void;
+
+  removeBatchObservableCallback(
+    callback: BatchObservableCallback<Attributes>,
+    observables: Observable<Attributes>[],
+  ): void;
+
   /**
    * Creates a new histogram instrument. A histogram is a metric that samples observations (usually things like request durations or response sizes) and counts them in configurable buckets.
    * @param name - The name of the histogram