@aztec/p2p 0.0.0-test.1 → 0.0.1-commit.1142ef1

package/src/services/discv5/discV5_service.ts

@@ -3,16 +3,16 @@ import { sleep } from '@aztec/foundation/sleep';
 import { type ComponentsVersions, checkCompressedComponentVersion } from '@aztec/stdlib/versioning';
 import { OtelMetricsAdapter, type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
-import { Discv5, type Discv5EventEmitter } from '@chainsafe/discv5';
-import { ENR, SignableENR } from '@chainsafe/enr';
 import type { PeerId } from '@libp2p/interface';
 import { type Multiaddr, multiaddr } from '@multiformats/multiaddr';
+import { Discv5, type Discv5EventEmitter, type IDiscv5CreateOptions } from '@nethermindeth/discv5';
+import { ENR, SignableENR } from '@nethermindeth/enr';
 import EventEmitter from 'events';
 
 import type { P2PConfig } from '../../config.js';
+import { createNodeENR } from '../../enr/generate-enr.js';
 import { AZTEC_ENR_KEY, Discv5Event, PeerEvent } from '../../types/index.js';
 import { convertToMultiaddr } from '../../util.js';
-import { setAztecEnrKey } from '../../versioning.js';
 import { type PeerDiscoveryService, PeerDiscoveryState } from '../service.js';
 
 const delayBeforeStart = 2000; // 2sec
@@ -22,7 +22,7 @@ const delayBeforeStart = 2000; // 2sec
  */
 export class DiscV5Service extends EventEmitter implements PeerDiscoveryService {
   /** The Discv5 instance */
-  private discv5: Discv5 & Discv5EventEmitter;
+  private discv5: Discv5EventEmitter;
 
   /** This instance's ENR */
   private enr: SignableENR;
@@ -30,57 +30,74 @@ export class DiscV5Service extends EventEmitter implements PeerDiscoveryService
   /** Version identifiers. */
   private versions: ComponentsVersions;
 
-  /** UDP listen addr */
-  private listenMultiAddrUdp: Multiaddr;
-
   private currentState = PeerDiscoveryState.STOPPED;
 
-  public readonly bootstrapNodes: string[] = [];
   private bootstrapNodePeerIds: PeerId[] = [];
-  private bootstrapNodeEnrs: ENR[] = [];
+  public bootstrapNodeEnrs: ENR[] = [];
+  private trustedPeerEnrs: ENR[] = [];
 
   private startTime = 0;
 
+  private handlers = {
+    onMultiaddrUpdated: this.onMultiaddrUpdated.bind(this),
+    onDiscovered: this.onDiscovered.bind(this),
+    onEnrAdded: this.onEnrAdded.bind(this),
+  };
+
   constructor(
     private peerId: PeerId,
     private config: P2PConfig,
+    private readonly packageVersion: string,
     telemetry: TelemetryClient = getTelemetryClient(),
     private logger = createLogger('p2p:discv5_service'),
+    configOverrides: Partial<IDiscv5CreateOptions> = {},
   ) {
     super();
-    const { tcpAnnounceAddress, udpAnnounceAddress, udpListenAddress, bootstrapNodes } = config;
-    this.bootstrapNodes = bootstrapNodes ?? [];
-    this.bootstrapNodeEnrs = this.bootstrapNodes.map(x => ENR.decodeTxt(x));
-    // create ENR from PeerId
-    this.enr = SignableENR.createFromPeerId(peerId);
-    // Add aztec identification to ENR
-    this.versions = setAztecEnrKey(this.enr, config);
-
-    if (!tcpAnnounceAddress) {
-      throw new Error('You need to provide at least a TCP announce address.');
+    const { p2pIp, p2pPort, p2pBroadcastPort, bootstrapNodes, trustedPeers, privatePeers } = config;
+
+    this.bootstrapNodeEnrs = bootstrapNodes.map(x => ENR.decodeTxt(x));
+    const privatePeerEnrs = new Set(privatePeers);
+    this.trustedPeerEnrs = trustedPeers.filter(x => !privatePeerEnrs.has(x)).map(x => ENR.decodeTxt(x));
+
+    // If no overridden broadcast port is provided, use the p2p port as the broadcast port
+    if (!p2pBroadcastPort) {
+      this.logger.warn('No p2pBroadcastPort provided, using p2pPort as broadcast port');
+      config.p2pBroadcastPort = p2pPort;
     }
 
-    const multiAddrTcp = multiaddr(`${convertToMultiaddr(tcpAnnounceAddress, 'tcp')}/p2p/${peerId.toString()}`);
-    // if no udp announce address is provided, use the tcp announce address
-    const multiAddrUdp = multiaddr(
-      `${convertToMultiaddr(udpAnnounceAddress || tcpAnnounceAddress, 'udp')}/p2p/${peerId.toString()}`,
-    );
+    const bindAddrs: any = {
+      ip4: multiaddr(convertToMultiaddr(config.listenAddress, p2pPort, 'udp')),
+    };
 
-    this.listenMultiAddrUdp = multiaddr(convertToMultiaddr(udpListenAddress, 'udp'));
+    let multiAddrUdp, multiAddrTcp;
+    if (p2pIp) {
+      multiAddrTcp = multiaddr(
+        `${convertToMultiaddr(p2pIp!, config.p2pBroadcastPort!, 'tcp')}/p2p/${peerId.toString()}`,
+      );
+      multiAddrUdp = multiaddr(
+        `${convertToMultiaddr(p2pIp!, config.p2pBroadcastPort!, 'udp')}/p2p/${peerId.toString()}`,
+      );
+    }
 
-    // set location multiaddr in ENR record
-    this.enr.setLocationMultiaddr(multiAddrUdp);
-    this.enr.setLocationMultiaddr(multiAddrTcp);
+    ({ enr: this.enr, versions: this.versions } = createNodeENR(
+      peerId,
+      multiAddrUdp,
+      multiAddrTcp,
+      config,
+      this.packageVersion,
+    ));
 
     const metricsRegistry = new OtelMetricsAdapter(telemetry);
     this.discv5 = Discv5.create({
       enr: this.enr,
       peerId,
-      bindAddrs: { ip4: this.listenMultiAddrUdp },
+      bindAddrs,
       config: {
         lookupTimeout: 2000,
         requestTimeout: 2000,
         allowUnverifiedSessions: true,
+        enrUpdate: !p2pIp ? true : false, // If no p2p IP is set, enrUpdate can automatically resolve it
+        ...configOverrides.config,
       },
       metricsRegistry,
     });
@@ -104,8 +121,17 @@ export class DiscV5Service extends EventEmitter implements PeerDiscoveryService
       }
     };
 
-    this.discv5.on(Discv5Event.DISCOVERED, this.onDiscovered.bind(this));
-    this.discv5.on(Discv5Event.ENR_ADDED, this.onEnrAdded.bind(this));
+    this.discv5.on(Discv5Event.DISCOVERED, this.handlers.onDiscovered);
+    this.discv5.on(Discv5Event.ENR_ADDED, this.handlers.onEnrAdded);
+    this.discv5.on(Discv5Event.MULTIADDR_UPDATED, this.handlers.onMultiaddrUpdated);
+  }
+
+  private onMultiaddrUpdated(m: Multiaddr) {
+    // We want to update our tcp port to match the udp port
+    // p2pBroadcastPort is optional on config, however it is set to default within the p2p client factory
+    const multiAddrTcp = multiaddr(convertToMultiaddr(m.nodeAddress().address, this.config.p2pBroadcastPort!, 'tcp'));
+    this.enr.setLocationMultiaddr(multiAddrTcp);
+    this.logger.info('Multiaddr updated', { multiaddr: multiAddrTcp.toString() });
   }
 
   public async start(): Promise<void> {
@@ -126,10 +152,14 @@ export class DiscV5Service extends EventEmitter implements PeerDiscoveryService
     this.currentState = PeerDiscoveryState.RUNNING;
 
     // Add bootnode ENR if provided
-    if (this.bootstrapNodes?.length) {
+    if (this.bootstrapNodeEnrs?.length) {
       // Do this conversion once since it involves an async function call
       this.bootstrapNodePeerIds = await Promise.all(this.bootstrapNodeEnrs.map(enr => enr.peerId()));
-      this.logger.info(`Adding ${this.bootstrapNodes} bootstrap nodes ENRs: ${this.bootstrapNodes.join(', ')}`);
+      this.logger.info(
+        `Adding ${this.bootstrapNodeEnrs.length} bootstrap nodes ENRs: ${this.bootstrapNodeEnrs
+          .map(enr => enr.encodeTxt())
+          .join(', ')}`,
+      );
       for (const enr of this.bootstrapNodeEnrs) {
         try {
           if (this.config.bootstrapNodeEnrVersionCheck) {
@@ -145,11 +175,23 @@ export class DiscV5Service extends EventEmitter implements PeerDiscoveryService
         }
       }
     }
+
+    // Add trusted peer ENRs if provided
+    if (this.trustedPeerEnrs?.length) {
+      this.logger.info(
+        `Adding ${this.trustedPeerEnrs.length} trusted peer ENRs: ${this.trustedPeerEnrs
+          .map(enr => enr.encodeTxt())
+          .join(', ')}`,
+      );
+      for (const enr of this.trustedPeerEnrs) {
+        this.discv5.addEnr(enr);
+      }
+    }
   }
 
   public async runRandomNodesQuery(): Promise<void> {
     if (this.currentState !== PeerDiscoveryState.RUNNING) {
-      throw new Error('DiscV5Service not running');
+      return;
     }
 
     // First, wait some time before starting the peer discovery
@@ -166,7 +208,7 @@ export class DiscV5Service extends EventEmitter implements PeerDiscoveryService
     }
   }
 
-  public getAllPeers(): ENR[] {
+  public getKadValues(): ENR[] {
     return this.discv5.kadValues();
   }
 
@@ -187,8 +229,12 @@ export class DiscV5Service extends EventEmitter implements PeerDiscoveryService
   }
 
   public async stop(): Promise<void> {
-    await this.discv5.off(Discv5Event.DISCOVERED, this.onDiscovered);
-    await this.discv5.off(Discv5Event.ENR_ADDED, this.onEnrAdded);
+    if (this.currentState !== PeerDiscoveryState.RUNNING) {
+      return;
+    }
+    await this.discv5.off(Discv5Event.DISCOVERED, this.handlers.onDiscovered);
+    await this.discv5.off(Discv5Event.ENR_ADDED, this.handlers.onEnrAdded);
+    await this.discv5.off(Discv5Event.MULTIADDR_UPDATED, this.handlers.onMultiaddrUpdated);
 
     await this.discv5.stop();
 
@@ -231,7 +277,7 @@ export class DiscV5Service extends EventEmitter implements PeerDiscoveryService
     // Check the peer is an aztec peer
     const value = enr.kvs.get(AZTEC_ENR_KEY);
     if (!value) {
-      this.logger.warn(`Peer node ${enr.nodeId} does not have aztec key in ENR`);
+      this.logger.debug(`Peer node ${enr.nodeId} does not have aztec key in ENR`);
       return false;
     }
 

package/src/services/dummy_service.ts

@@ -1,17 +1,40 @@
+import type { EthAddress } from '@aztec/foundation/eth-address';
 import type { PeerInfo } from '@aztec/stdlib/interfaces/server';
-import type { BlockAttestation, BlockProposal, Gossipable } from '@aztec/stdlib/p2p';
-import { TxHash } from '@aztec/stdlib/tx';
+import type { Gossipable, PeerErrorSeverity } from '@aztec/stdlib/p2p';
+import { Tx, TxHash } from '@aztec/stdlib/tx';
 
 import type { PeerId } from '@libp2p/interface';
+import type { ENR } from '@nethermindeth/enr';
 import EventEmitter from 'events';
 
-import type { ReqRespSubProtocol, SubProtocolMap } from './reqresp/interface.js';
-import { type P2PService, type PeerDiscoveryService, PeerDiscoveryState } from './service.js';
+import type { PeerManagerInterface } from './peer-manager/interface.js';
+import type { P2PReqRespConfig } from './reqresp/config.js';
+import { type AuthRequest, StatusMessage } from './reqresp/index.js';
+import type {
+  ReqRespInterface,
+  ReqRespResponse,
+  ReqRespSubProtocol,
+  ReqRespSubProtocolHandler,
+  ReqRespSubProtocolHandlers,
+  ReqRespSubProtocolValidators,
+  SubProtocolMap,
+} from './reqresp/interface.js';
+import type { GoodByeReason } from './reqresp/protocols/goodbye.js';
+import { ReqRespStatus } from './reqresp/status.js';
+import {
+  type P2PBlockReceivedCallback,
+  type P2PCheckpointReceivedCallback,
+  type P2PService,
+  type PeerDiscoveryService,
+  PeerDiscoveryState,
+} from './service.js';
 
 /**
  * A dummy implementation of the P2P Service.
  */
 export class DummyP2PService implements P2PService {
+  updateConfig(_config: Partial<P2PReqRespConfig>): void {}
+
   /** Returns an empty array for peers. */
   getPeers(): PeerInfo[] {
     return [];
@@ -37,7 +60,9 @@ export class DummyP2PService implements P2PService {
    * Called to have the given message propagated through the P2P network.
    * @param _ - The message to be propagated.
    */
-  public propagate<T extends Gossipable>(_: T) {}
+  public propagate<T extends Gossipable>(_: T) {
+    return Promise.resolve();
+  }
 
   /**
    * Called upon receipt of settled transactions.
@@ -48,7 +73,12 @@ export class DummyP2PService implements P2PService {
   /**
    * Register a callback into the validator client for when a block proposal is received
    */
-  public registerBlockReceivedCallback(_: (block: BlockProposal) => Promise<BlockAttestation>) {}
+  public registerBlockReceivedCallback(_callback: P2PBlockReceivedCallback) {}
+
+  /**
+   * Register a callback into the validator client for when a checkpoint proposal is received
+   */
+  public registerCheckpointReceivedCallback(_callback: P2PCheckpointReceivedCallback) {}
 
   /**
    * Sends a request to a peer.
@@ -83,6 +113,25 @@ export class DummyP2PService implements P2PService {
   public getEnr(): undefined {
     return undefined;
   }
+
+  validate(_txs: Tx[]): Promise<void> {
+    return Promise.resolve();
+  }
+
+  addReqRespSubProtocol(
+    _subProtocol: ReqRespSubProtocol,
+    _handler: ReqRespSubProtocolHandler,
+    _validator?: ReqRespSubProtocolValidators[ReqRespSubProtocol],
+  ): Promise<void> {
+    return Promise.resolve();
+  }
+
+  handleAuthRequestFromPeer(_authRequest: AuthRequest, _peerId: PeerId): Promise<StatusMessage> {
+    return Promise.resolve(StatusMessage.random());
+  }
+
+  //this is no-op
+  registerThisValidatorAddresses(_address: EthAddress[]): void {}
 }
 
 /**
@@ -90,7 +139,7 @@ export class DummyP2PService implements P2PService {
  */
 export class DummyPeerDiscoveryService extends EventEmitter implements PeerDiscoveryService {
   private currentState = PeerDiscoveryState.STOPPED;
-  public bootstrapNodes: string[] = [];
+  public bootstrapNodeEnrs: ENR[] = [];
 
   /**
    * Starts the dummy implementation.
@@ -110,9 +159,9 @@ export class DummyPeerDiscoveryService extends EventEmitter implements PeerDisco
   }
   /**
    * Called to discover peers in the network.
-   * @returns An array of discovered peer addresses.
+   * @returns An array of Enrs.
    */
-  public getAllPeers() {
+  public getKadValues() {
     return [];
   }
 
@@ -132,3 +181,98 @@ export class DummyPeerDiscoveryService extends EventEmitter implements PeerDisco
     return undefined;
   }
 }
+
+export class DummyPeerManager implements PeerManagerInterface {
+  constructor(
+    public peerId: PeerId,
+    private peersProvider?: { getPeers: () => PeerId[] },
+  ) {}
+
+  public getPeers(_includePending?: boolean): PeerInfo[] {
+    if (!this.peersProvider) {
+      return [];
+    }
+    return this.peersProvider
+      .getPeers()
+      .filter(peer => !peer.equals(this.peerId))
+      .map(id => ({
+        id: id.toString(),
+        status: 'connected',
+        score: 0,
+      }));
+  }
+
+  public initializePeers(): Promise<void> {
+    return Promise.resolve();
+  }
+  public getPeerScore(_peerId: string): number {
+    return 0;
+  }
+
+  public shouldDisableP2PGossip(_peerId: string): boolean {
+    return false;
+  }
+
+  public stop(): Promise<void> {
+    return Promise.resolve();
+  }
+  public heartbeat(): Promise<void> {
+    return Promise.resolve();
+  }
+  public addTrustedPeer(_peerId: PeerId): void {}
+  public addPrivatePeer(_peerId: PeerId): void {}
+  public goodbyeReceived(_peerId: PeerId, _reason: GoodByeReason): void {}
+  public penalizePeer(_peerId: PeerId, _penalty: PeerErrorSeverity): void {}
+  public addPreferredPeer(_peerId: PeerId): void {}
+  public handleAuthRequestFromPeer(_authRequest: AuthRequest, _peerId: PeerId): Promise<StatusMessage> {
+    return Promise.resolve(StatusMessage.random());
+  }
+
+  //this is no-op
+  registerThisValidatorAddresses(_address: EthAddress[]): void {}
+}
+
+export class DummyReqResp implements ReqRespInterface {
+  updateConfig(_config: Partial<P2PReqRespConfig>): void {}
+  start(
+    _subProtocolHandlers: ReqRespSubProtocolHandlers,
+    _subProtocolValidators: ReqRespSubProtocolValidators,
+  ): Promise<void> {
+    return Promise.resolve();
+  }
+  stop(): Promise<void> {
+    return Promise.resolve();
+  }
+  sendRequest<SubProtocol extends ReqRespSubProtocol>(
+    _subProtocol: SubProtocol,
+    _request: InstanceType<SubProtocolMap[SubProtocol]['request']>,
+  ): Promise<InstanceType<SubProtocolMap[SubProtocol]['response']> | undefined> {
+    return Promise.resolve(undefined);
+  }
+  sendBatchRequest<SubProtocol extends ReqRespSubProtocol>(
+    _subProtocol: SubProtocol,
+    _requests: InstanceType<SubProtocolMap[SubProtocol]['request']>[],
+    _pinnedPeer: PeerId | undefined,
+    _timeoutMs?: number,
+    _maxPeers?: number,
+    _maxRetryAttempts?: number,
+  ): Promise<InstanceType<SubProtocolMap[SubProtocol]['response']>[]> {
+    return Promise.resolve([]);
+  }
+  public sendRequestToPeer(
+    _peerId: PeerId,
+    _subProtocol: ReqRespSubProtocol,
+    _payload: Buffer,
+    _dialTimeout?: number,
+  ): Promise<ReqRespResponse> {
+    return Promise.resolve({ status: ReqRespStatus.SUCCESS, data: Buffer.from([]) });
+  }
+
+  addSubProtocol(
+    _subProtocol: ReqRespSubProtocol,
+    _handler: ReqRespSubProtocolHandler,
+    _validator?: ReqRespSubProtocolValidators[ReqRespSubProtocol],
+  ): Promise<void> {
+    return Promise.resolve();
+  }
+}

package/src/services/encoding.ts

@@ -1,5 +1,7 @@
 // Taken from lodestar: https://github.com/ChainSafe/lodestar
-import { sha256 } from '@aztec/foundation/crypto';
+import { sha256 } from '@aztec/foundation/crypto/sha256';
+import { createLogger } from '@aztec/foundation/log';
+import { TopicType, getTopicFromString } from '@aztec/stdlib/p2p';
 
 import type { RPC } from '@chainsafe/libp2p-gossipsub/message';
 import type { DataTransform } from '@chainsafe/libp2p-gossipsub/types';
@@ -49,31 +51,106 @@ export function getMsgIdFn(message: Message) {
   return sha256(Buffer.concat(vec)).subarray(0, 20);
 }
 
+const DefaultMaxSizesKb: Record<TopicType, number> = {
+  // Tx effects should not exceed 128kb, so 512kb for the full tx obj should be sufficient
+  [TopicType.tx]: 512,
+  // An attestation has roughly 30 fields, which is 1kb, so 5x is plenty
+  [TopicType.checkpoint_attestation]: 5,
+  // Proposals may carry some tx objects, so we allow a larger size capped at 10mb
+  // Note this may not be enough for carrying all tx objects in a block
+  [TopicType.block_proposal]: 1024 * 10,
+  // TODO(palla/mbps): Check size for checkpoint proposal
+  [TopicType.checkpoint_proposal]: 1024 * 10,
+};
+
 /**
  * Snappy transform for libp2p gossipsub
  */
 export class SnappyTransform implements DataTransform {
+  constructor(
+    private maxSizesKb: Record<TopicType, number> = DefaultMaxSizesKb,
+    private defaultMaxSizeKb: number = 10 * 1024,
+    private logger = createLogger('p2p:snappy-transform'),
+  ) {}
+
   // Topic string included to satisfy DataTransform interface
-  inboundTransform(_topicStr: string, data: Uint8Array): Uint8Array {
-    return this.inboundTransformNoTopic(Buffer.from(data));
+  inboundTransform(topicStr: string, data: Uint8Array): Uint8Array {
+    const topic = getTopicFromString(topicStr);
+    return this.inboundTransformData(Buffer.from(data), topic);
   }
 
-  public inboundTransformNoTopic(data: Buffer): Buffer {
+  public inboundTransformData(data: Buffer, topic?: TopicType): Buffer {
     if (data.length === 0) {
       return data;
     }
+    const maxSizeKb = this.maxSizesKb[topic!] ?? this.defaultMaxSizeKb;
+    const { decompressedSize } = readSnappyPreamble(data);
+    if (decompressedSize > maxSizeKb * 1024) {
+      this.logger.warn(`Decompressed size ${decompressedSize} exceeds maximum allowed size of ${maxSizeKb}kb`);
+      throw new Error(`Decompressed size ${decompressedSize} exceeds maximum allowed size of ${maxSizeKb}kb`);
+    }
+
     return Buffer.from(uncompressSync(data, { asBuffer: true }));
   }
 
   // Topic string included to satisfy DataTransform interface
   outboundTransform(_topicStr: string, data: Uint8Array): Uint8Array {
-    return this.outboundTransformNoTopic(Buffer.from(data));
+    return this.outboundTransformData(Buffer.from(data));
   }
 
-  public outboundTransformNoTopic(data: Buffer): Buffer {
+  public outboundTransformData(data: Buffer): Buffer {
     if (data.length === 0) {
       return data;
     }
     return Buffer.from(compressSync(data));
   }
 }
+
+/**
+ * Reads the Snappy preamble from compressed data and returns the expected decompressed size.
+ *
+ * The Snappy format starts with a little-endian varint encoding the uncompressed length.
+ * Varints consist of a series of bytes where:
+ * - Lower 7 bits contain data
+ * - Upper bit (0x80) is set if more bytes follow
+ *
+ * @param data - The compressed data starting with the Snappy preamble
+ * @returns Object containing the decompressed size and the number of bytes read from the preamble
+ * @throws Error if the data is too short or the varint is invalid
+ */
+export function readSnappyPreamble(data: Uint8Array): { decompressedSize: number; bytesRead: number } {
+  if (data.length === 0) {
+    throw new Error('Cannot read preamble from empty data');
+  }
+
+  let result = 0;
+  let shift = 0;
+  let bytesRead = 0;
+
+  // Maximum varint length for 32-bit value is 5 bytes
+  // (7 bits per byte, so 5 bytes = 35 bits, enough for 2^32 - 1)
+  const maxBytes = 5;
+
+  for (let i = 0; i < Math.min(data.length, maxBytes); i++) {
+    const byte = data[i];
+    bytesRead++;
+
+    // Extract lower 7 bits and add to result with appropriate shift
+    // Use >>> 0 to convert to unsigned 32-bit integer to avoid sign issues
+    result = (result | ((byte & 0x7f) << shift)) >>> 0;
+
+    // If upper bit is not set, we're done
+    if ((byte & 0x80) === 0) {
+      return { decompressedSize: result, bytesRead };
+    }
+
+    shift += 7;
+  }
+
+  // If we get here, either we ran out of data or the varint is too long
+  if (bytesRead >= maxBytes) {
+    throw new Error('Varint is too long (max 5 bytes for 32-bit value)');
+  }
+
+  throw new Error('Incomplete varint: data ended before varint termination');
+}

package/src/services/index.ts

@@ -1,2 +1,6 @@
 export * from './service.js';
 export * from './libp2p/libp2p_service.js';
+export * from './tx_provider.js';
+export * from './dummy_service.js';
+export * from './reqresp/index.js';
+export * from './tx_collection/index.js';

package/src/services/libp2p/instrumentation.ts

@@ -0,0 +1,113 @@
+import type { Timer } from '@aztec/foundation/timer';
+import { TopicType } from '@aztec/stdlib/p2p';
+import {
+  Attributes,
+  type BatchObservableResult,
+  type Histogram,
+  Metrics,
+  type ObservableGauge,
+  type TelemetryClient,
+  type UpDownCounter,
+} from '@aztec/telemetry-client';
+
+import { type RecordableHistogram, createHistogram } from 'node:perf_hooks';
+
+export class P2PInstrumentation {
+  private messageValidationDuration: Histogram;
+  private messagePrevalidationCount: UpDownCounter;
+  private messageLatency: Histogram;
+  private txReceivedCount: UpDownCounter;
+
+  private aggLatencyHisto = new Map<TopicType, RecordableHistogram>();
+  private aggValidationHisto = new Map<TopicType, RecordableHistogram>();
+
+  private aggLatencyMetrics: Record<'min' | 'max' | 'p50' | 'p90' | 'avg', ObservableGauge>;
+  private aggValidationMetrics: Record<'min' | 'max' | 'p50' | 'p90' | 'avg', ObservableGauge>;
+
+  constructor(client: TelemetryClient, name: string) {
+    const meter = client.getMeter(name);
+
+    this.messageValidationDuration = meter.createHistogram(Metrics.P2P_GOSSIP_MESSAGE_VALIDATION_DURATION);
+
+    this.messagePrevalidationCount = meter.createUpDownCounter(Metrics.P2P_GOSSIP_MESSAGE_PREVALIDATION_COUNT);
+
+    this.messageLatency = meter.createHistogram(Metrics.P2P_GOSSIP_MESSAGE_LATENCY);
+
+    this.txReceivedCount = meter.createUpDownCounter(Metrics.P2P_GOSSIP_TX_RECEIVED_COUNT);
+
+    this.aggLatencyMetrics = {
+      avg: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_LATENCY_AVG),
+      max: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_LATENCY_MAX),
+      min: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_LATENCY_MIN),
+      p50: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_LATENCY_P50),
+      p90: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_LATENCY_P90),
+    };
+
+    this.aggValidationMetrics = {
+      avg: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_VALIDATION_DURATION_AVG),
+      max: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_VALIDATION_DURATION_MAX),
+      min: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_VALIDATION_DURATION_MIN),
+      p50: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_VALIDATION_DURATION_P50),
+      p90: meter.createObservableGauge(Metrics.P2P_GOSSIP_AGG_MESSAGE_VALIDATION_DURATION_P90),
+    };
+
+    meter.addBatchObservableCallback(this.aggregate, [
+      ...Object.values(this.aggValidationMetrics),
+      ...Object.values(this.aggLatencyMetrics),
+    ]);
+  }
+
+  public recordMessageValidation(topicName: TopicType, timerOrMs: Timer | number) {
+    const ms = Math.ceil(typeof timerOrMs === 'number' ? timerOrMs : timerOrMs.ms());
+    this.messageValidationDuration.record(ms, { [Attributes.TOPIC_NAME]: topicName });
+
+    let validationHistogram = this.aggValidationHisto.get(topicName);
+    if (!validationHistogram) {
+      validationHistogram = createHistogram({ min: 1, max: 5 * 60 * 1000 }); // 5 mins
+      this.aggValidationHisto.set(topicName, validationHistogram);
+    }
+
+    validationHistogram.record(Math.max(ms, 1));
+  }
+
+  public incrementTxReceived(count: number) {
+    this.txReceivedCount.add(count);
+  }
+
+  public incMessagePrevalidationStatus(passed: boolean, topicName: TopicType | undefined) {
+    this.messagePrevalidationCount.add(1, { [Attributes.TOPIC_NAME]: topicName, [Attributes.OK]: passed });
+  }
+
+  public recordMessageLatency(topicName: TopicType, timerOrMs: Timer | number) {
+    const ms = Math.ceil(typeof timerOrMs === 'number' ? timerOrMs : timerOrMs.ms());
+    this.messageLatency.record(ms, { [Attributes.TOPIC_NAME]: topicName });
+
+    let latencyHistogram = this.aggLatencyHisto.get(topicName);
+    if (!latencyHistogram) {
+      latencyHistogram = createHistogram({ min: 1, max: 60 * 1000 }); // Max: 1 minute
+      this.aggLatencyHisto.set(topicName, latencyHistogram);
+    }
+
+    latencyHistogram.record(Math.max(ms, 1));
+  }
+
+  private aggregate = (res: BatchObservableResult) => {
+    for (const [metrics, histograms] of [
+      [this.aggLatencyMetrics, this.aggLatencyHisto],
+      [this.aggValidationMetrics, this.aggValidationHisto],
+    ] as const) {
+      for (const topicName of Object.values(TopicType)) {
+        const histogram = histograms.get(topicName);
+        if (!histogram || histogram.count === 0) {
+          continue;
+        }
+
+        res.observe(metrics.avg, histogram.mean, { [Attributes.TOPIC_NAME]: topicName });
+        res.observe(metrics.max, histogram.max, { [Attributes.TOPIC_NAME]: topicName });
+        res.observe(metrics.min, histogram.min, { [Attributes.TOPIC_NAME]: topicName });
+        res.observe(metrics.p50, histogram.percentile(50), { [Attributes.TOPIC_NAME]: topicName });
+        res.observe(metrics.p90, histogram.percentile(90), { [Attributes.TOPIC_NAME]: topicName });
+      }
+    }
+  };
+}